import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";

// Several Parley keys, chosen by nickname, all read from the environment so
// nothing secret is written to disk by this extension.
//
//   PARLEY_API_KEY               the unnamed "default" key
//                                (MIT_PARLEY_API_KEY / MY_MIT_PARLEY_API_KEY also work)
//   PARLEY_<NICK>_API_KEY        the key for profile <nick>
//   PARLEY_<NICK>_BASE_URL       optional base URL for that profile
//   PARLEY_<NICK>_LABEL          optional display name (env var names cannot
//                                hold dots, so "6.101" lives here)
//   PARLEY_KEY_PROFILE           nickname to start with
//
// Under nono, a custom credential named `parley_<nick>` yields exactly this
// PARLEY_<NICK>_API_KEY (phantom token) / PARLEY_<NICK>_BASE_URL (proxy route)
// pair inside the sandbox. Switching profiles switches proxy routes, and the
// real keys stay on the host. Variable names are matched without regard to
// case because nono upper-cases the names it derives.

export const DEFAULT_PARLEY_BASE_URL = "https://parley.api.mit.edu/v1";

export const PARLEY_BASE_URL_ENV_VARS = [
  "MIT_PARLEY_BASE_URL",
  "PARLEY_BASE_URL",
  // Derived form some credential proxies use for a service named "MIT-parley".
  "MIT-PARLEY_BASE_URL",
  "MY_MIT_PARLEY_BASE_URL",
];

// Parley's own docs use PARLEY_API_KEY; the other names are kept for existing setups.
export const PARLEY_API_KEY_ENV_VARS = ["MY_MIT_PARLEY_API_KEY", "MIT_PARLEY_API_KEY", "PARLEY_API_KEY"];

export const DEFAULT_PROFILE_NICK = "default";
export const PROFILE_SELECT_ENV_VAR = "PARLEY_KEY_PROFILE";
export const KEY_SWITCH_ENTRY_TYPE = "parley-key";
const STATUS_KEY = "parley-key";

export interface ParleyKeyProfile {
  /** Lower-case nickname used for matching and as the /parley-key argument. */
  nick: string;
  /** Display name: PARLEY_<NICK>_LABEL when set, else the nickname. */
  label: string;
  /** Exact environment variable name holding the key. */
  apiKeyVar: string;
  /** Exact environment variable name holding the base URL, when one is set. */
  baseUrlVar?: string;
  isDefault: boolean;
}

export interface KeySwitchEntryData {
  profile: string;
  label?: string;
}

type Env = Record<string, string | undefined>;

// Upper-case name -> actual name, for non-empty variables only.
function envIndex(env: Env): Map<string, string> {
  const index = new Map<string, string>();
  for (const [name, value] of Object.entries(env)) {
    if (value?.trim()) index.set(name.toUpperCase(), name);
  }
  return index;
}

function firstSet(env: Env, names: string[]): string | undefined {
  return names.find((name) => env[name]?.trim());
}

export function listParleyKeyProfiles(env: Env = process.env): ParleyKeyProfile[] {
  const index = envIndex(env);
  const lookup = (upperName: string) => index.get(upperName);
  const value = (name: string | undefined) => (name ? env[name]?.trim() : undefined);

  const profiles: ParleyKeyProfile[] = [];
  const defaultVar = firstSet(env, PARLEY_API_KEY_ENV_VARS);
  if (defaultVar) {
    profiles.push({
      nick: DEFAULT_PROFILE_NICK,
      label: value(lookup("PARLEY_LABEL")) || DEFAULT_PROFILE_NICK,
      apiKeyVar: defaultVar,
      baseUrlVar: firstSet(env, PARLEY_BASE_URL_ENV_VARS),
      isDefault: true,
    });
  }

  const seen = new Set<string>();
  for (const [name, raw] of Object.entries(env)) {
    const match = /^PARLEY_(.+)_API_KEY$/i.exec(name);
    if (!match || !raw?.trim()) continue;
    const nick = match[1].toLowerCase();
    if (nick === DEFAULT_PROFILE_NICK || seen.has(nick)) continue;
    seen.add(nick);
    const upper = match[1].toUpperCase();
    profiles.push({
      nick,
      label: value(lookup(`PARLEY_${upper}_LABEL`)) || nick,
      apiKeyVar: name,
      baseUrlVar: lookup(`PARLEY_${upper}_BASE_URL`),
      isDefault: false,
    });
  }

  return profiles.sort((a, b) => {
    if (a.isDefault !== b.isDefault) return a.isDefault ? -1 : 1;
    return a.nick.localeCompare(b.nick);
  });
}

export function findParleyKeyProfile(nick: string, profiles = listParleyKeyProfiles()): ParleyKeyProfile | undefined {
  const wanted = nick.trim().toLowerCase();
  return profiles.find((p) => p.nick === wanted || p.label.toLowerCase() === wanted);
}

// An explicit selector must resolve: never silently bill a different key.
// Without one, use the default key, then the first named profile.
export function pickStartupProfile(profiles = listParleyKeyProfiles(), env: Env = process.env): ParleyKeyProfile | undefined {
  const wanted = env[PROFILE_SELECT_ENV_VAR]?.trim();
  if (wanted) {
    const found = findParleyKeyProfile(wanted, profiles);
    if (found) return found;
    throw new Error(`Parley key profile "${wanted}" is unavailable. Configure its key or unset ${PROFILE_SELECT_ENV_VAR}.`);
  }
  return profiles.find((p) => p.isDefault) ?? profiles[0];
}

let activeNick: string | undefined;

export function getActiveParleyKeyProfile(): ParleyKeyProfile | undefined {
  const profiles = listParleyKeyProfiles();
  if (activeNick) {
    const found = profiles.find((p) => p.nick === activeNick);
    if (found) return found;
  }
  return pickStartupProfile(profiles);
}

export function setActiveParleyKeyProfile(nick: string | undefined): ParleyKeyProfile | undefined {
  if (nick === undefined) {
    activeNick = undefined;
    return getActiveParleyKeyProfile();
  }
  const found = findParleyKeyProfile(nick);
  if (found) {
    activeNick = found.nick;
    // Export only the selector. Fresh --no-session children inherit the same
    // profile (and its URL override) without putting a key in argv or artifacts.
    process.env[PROFILE_SELECT_ENV_VAR] = found.nick;
  }
  return found;
}

/** The key for the active profile, if any is configured. */
export function resolveParleyApiKey(): string | undefined {
  const profile = getActiveParleyKeyProfile();
  return profile ? process.env[profile.apiKeyVar]?.trim() || undefined : undefined;
}

/** Base URL for the active profile, falling back to the generic override
 *  variables and then to Parley's public endpoint. */
export function resolveParleyBaseUrl(): string {
  const profile = getActiveParleyKeyProfile();
  if (profile?.baseUrlVar) {
    const value = process.env[profile.baseUrlVar]?.trim();
    if (value) return value;
  }
  for (const key of PARLEY_BASE_URL_ENV_VARS) {
    const value = process.env[key]?.trim();
    if (value) return value;
  }
  return DEFAULT_PARLEY_BASE_URL;
}

/** Whether the footer should bother naming the key: more than one profile,
 *  or a named one selected explicitly. */
export function shouldShowParleyKey(): boolean {
  const profiles = listParleyKeyProfiles();
  if (profiles.length > 1) return true;
  const active = getActiveParleyKeyProfile();
  return active !== undefined && !active.isDefault;
}

function describeProfile(profile: ParleyKeyProfile, active: ParleyKeyProfile | undefined): string {
  const key = process.env[profile.apiKeyVar]?.trim() ?? "";
  const tail = key.length >= 8 ? `…${key.slice(-4)}` : "";
  const name = profile.label === profile.nick ? profile.nick : `${profile.label} (${profile.nick})`;
  const mark = active?.nick === profile.nick ? "  ← active" : "";
  return `${name}  ${profile.apiKeyVar}${tail}${mark}`;
}

// Find the most recent switch recorded on the current branch so a resumed
// session comes back on the key it was using.
export function restoreParleyKeyFromSession(ctx: ExtensionContext): { restored?: ParleyKeyProfile; missing?: KeySwitchEntryData } {
  const entries = ctx.sessionManager.getBranch() as any[];
  for (let i = entries.length - 1; i >= 0; i--) {
    const entry = entries[i];
    if (entry?.type !== "custom" || entry.customType !== KEY_SWITCH_ENTRY_TYPE) continue;
    const data = entry.data as KeySwitchEntryData | undefined;
    if (!data?.profile) return {};
    const restored = setActiveParleyKeyProfile(data.profile);
    return restored ? { restored } : { missing: data };
  }
  return {};
}

export function registerParleyKeySwitcher(
  pi: ExtensionAPI,
  onProfileChange?: (ctx: ExtensionContext) => Promise<void>,
): void {
  const syncStatus = (ctx: ExtensionContext) => {
    if (!ctx.hasUI) return;
    const active = getActiveParleyKeyProfile();
    ctx.ui.setStatus(STATUS_KEY, active && shouldShowParleyKey() ? `key: ${active.label}` : undefined);
  };

  pi.on("session_start", async (_event, ctx) => {
    const { missing } = restoreParleyKeyFromSession(ctx);
    if (missing && ctx.hasUI) {
      const fallback = getActiveParleyKeyProfile();
      ctx.ui.notify(
        `This session used Parley key "${missing.label ?? missing.profile}", which is not in the environment.` +
          (fallback ? ` Using "${fallback.label}" instead.` : " No Parley key is configured."),
        "warning",
      );
    }
    syncStatus(ctx);
    await onProfileChange?.(ctx);
  });

  pi.registerCommand("parley-key", {
    description: "Switch between Parley keys (PARLEY_<NICK>_API_KEY) and reset the session cost counter",
    getArgumentCompletions: (prefix) => {
      const items = listParleyKeyProfiles()
        .filter((p) => p.nick.startsWith(prefix.toLowerCase()))
        .map((p) => ({ value: p.nick, label: p.label === p.nick ? p.nick : `${p.nick} — ${p.label}` }));
      return items.length ? items : null;
    },
    handler: async (args, ctx) => {
      const profiles = listParleyKeyProfiles();
      if (profiles.length === 0) {
        ctx.ui.notify("No Parley keys found. Export PARLEY_API_KEY or PARLEY_<NICK>_API_KEY and restart Pi.", "error");
        return;
      }
      const active = getActiveParleyKeyProfile();
      const known = profiles.map((p) => p.nick).join(", ");

      let target: ParleyKeyProfile | undefined;
      const arg = args.trim();
      if (arg) {
        target = findParleyKeyProfile(arg, profiles);
        if (!target) {
          ctx.ui.notify(`Unknown Parley key "${arg}". Known: ${known}`, "error");
          return;
        }
      } else if (ctx.hasUI) {
        const options = profiles.map((p) => describeProfile(p, active));
        const choice = await ctx.ui.select("Parley key", options);
        if (choice === undefined) return;
        target = profiles[options.indexOf(choice)];
      } else {
        ctx.ui.notify(`Usage: /parley-key <nickname>. Known: ${known}`, "info");
        return;
      }

      if (target.nick === active?.nick) {
        ctx.ui.notify(`Parley key is already "${target.label}".`, "info");
        return;
      }

      setActiveParleyKeyProfile(target.nick);
      pi.appendEntry<KeySwitchEntryData>(KEY_SWITCH_ENTRY_TYPE, { profile: target.nick, label: target.label });
      syncStatus(ctx);
      ctx.ui.notify(`Parley key: ${target.label}. Session cost counter reset.`, "info");
      await onProfileChange?.(ctx);
    },
  });
}