import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { execFile } from "node:child_process";
import { promisify } from "node:util";
import { readFileSync, writeFileSync, mkdirSync } from "node:fs";
import { homedir } from "node:os";
import { join, dirname } from "node:path";

const pexec = promisify(execFile);

type EnvDef = { profile?: string; context?: string; prod?: boolean; accountId?: string };
type Config = { envs?: Record<string, EnvDef>; prodAccounts?: string[] };

const CONFIG_PATH = join(homedir(), ".pi", "agent", "aws-accounts.json");

// Active selection. Set by the human (/aws-switch), by startup adoption of the live context,
// or by an in-session `kubectl config use-context`.
let current: ({ env: string } & EnvDef) | null = null;

// Config is OPTIONAL — everything is auto-derived. It only adds prod marking, friendly
// aliases, and manual pairing overrides.
let raw: Config = loadConfig();
// Effective env map: derived from ~/.aws/config + kube contexts, unioned with raw.envs.
let envs: Record<string, EnvDef> = {};

const TOUCHES_CLOUD = /\b(kubectl|aws|eksctl|helm)\b/;
const READONLY =
  /\b(kubectl\s+(get|describe|logs|top|explain|version|cluster-info|api-resources|api-versions|config\s+(get-contexts|current-context|view|get-clusters))|aws\s+\S+\s+(describe|list|get)[\w-]*|aws\s+sts\s+get-caller-identity|helm\s+(list|status|get|history|version))\b/;
const MUTATING =
  /\b(delete|destroy|terminate|apply|create|replace|patch|edit|scale|drain|cordon|uncordon|rollout|annotate|label|taint|exec|install|upgrade|uninstall|rollback|put-|update-|modify-|remove-|stop-|start-|reboot-|run-instances|detach-|attach-|deregister-|revoke-|disable-|enable-|associate-|disassociate-)\b/i;

function loadConfig(): Config {
  try {
    const c = JSON.parse(readFileSync(CONFIG_PATH, "utf8"));
    return c && typeof c === "object" ? c : {};
  } catch {
    return {}; // absent/invalid config => pure zero-config mode
  }
}

// AWS account ids are 12 digits; pull from any arn (eks context or iam role_arn).
function acctOf(arn?: string): string | undefined {
  return arn?.match(/:(\d{12}):/)?.[1];
}

// profile name -> account id, parsed from ~/.aws/config (account lives in role_arn).
function awsProfiles(): Record<string, string | undefined> {
  const out: Record<string, string | undefined> = {};
  try {
    let cur = "";
    for (const ln of readFileSync(join(homedir(), ".aws", "config"), "utf8").split("\n")) {
      const line = ln.trim();
      const sec = line.match(/^\[(?:profile\s+)?([^\]]+)\]$/);
      if (sec) {
        cur = sec[1].trim();
        out[cur] = undefined;
        continue;
      }
      const role = line.match(/^role_arn\s*=\s*(\S+)/);
      if (role && cur) out[cur] = acctOf(role[1]);
    }
  } catch {
    /* no ~/.aws/config */
  }
  return out;
}

async function kubeContexts(): Promise<string[]> {
  try {
    return (await pexec("kubectl", ["config", "get-contexts", "-o", "name"])).stdout
      .split("\n")
      .map((s) => s.trim())
      .filter(Boolean);
  } catch {
    return [];
  }
}

// Build one env per kube context (so multiple clusters under one account each get an entry),
// pairing each to an AWS profile by account id. Union explicit raw.envs; mark prod by account.
async function rebuild(): Promise<void> {
  const profiles = awsProfiles();
  const profByAcct: Record<string, string> = {};
  for (const [n, a] of Object.entries(profiles)) if (a && !(a in profByAcct)) profByAcct[a] = n;
  const merged: Record<string, EnvDef> = {};
  for (const c of await kubeContexts()) {
    const accountId = acctOf(c);
    let key = clusterName(c);
    if (merged[key]) key = `${key}-${accountId || "x"}`; // disambiguate duplicate cluster names
    merged[key] = { profile: accountId ? profByAcct[accountId] : undefined, context: c, accountId };
  }
  for (const [name, e] of Object.entries(raw.envs || {})) {
    merged[name] = { ...merged[name], ...e, accountId: e.accountId || acctOf(e.context) || merged[name]?.accountId };
  }
  const prod = new Set((raw.prodAccounts || []).map(String));
  for (const e of Object.values(merged)) {
    if (e.prod !== true && e.accountId && prod.has(e.accountId)) e.prod = true;
  }
  envs = merged;
}

function shortCtx(c?: string): string {
  if (!c) return "-";
  const p = c.split("/");
  return p[p.length - 1].slice(0, 24);
}

function clusterName(ctx: string): string {
  const m = ctx.match(/cluster\/(.+)$/);
  return m ? m[1] : ctx;
}

// Picker label: "<key> — <profile|account> · <cluster> [PROD]". Includes the key so labels are unique.
function labelFor(k: string): string {
  const e = envs[k];
  const who = e.profile || "no-profile";
  const acct = e.accountId ? ` (${e.accountId})` : "";
  const cl = e.context ? clusterName(e.context) : "(no cluster)";
  return `${k}  —  ${who}${acct} · ${cl}${e.prod ? "  [PROD]" : ""}`;
}

function setStatus(ctx: ExtensionContext): void {
  if (!current) {
    ctx.ui.setStatus("aws-accounts", "aws: no env selected");
    return;
  }
  const tag = current.prod ? "PROD\u26a0" : current.env;
  ctx.ui.setStatus("aws-accounts", `aws:${tag} [${current.profile || "-"}] k8s:${shortCtx(current.context)}`);
}

export default function (pi: ExtensionAPI) {
  pi.on("session_start", async (_e, ctx) => {
    try {
      raw = loadConfig();
      await rebuild();
      // Adopt the live kube context so we never show "no env" while actually pointed at prod.
      try {
        const live = (await pexec("kubectl", ["config", "current-context"])).stdout.trim();
        if (live) {
          const hit = Object.entries(envs).find(([, e]) => e.context === live);
          current = hit
            ? { env: hit[0], ...hit[1], context: live, profile: process.env.AWS_PROFILE }
            : { env: "(unmapped)", context: live, profile: process.env.AWS_PROFILE };
        }
      } catch {
        /* kubectl absent */
      }
      setStatus(ctx);
      const n = Object.keys(envs).length;
      if (n === 0)
        ctx.ui.notify("pi-aws-accounts: no AWS profiles or kube contexts detected — nothing to switch (see README).", "info");
      else if (current?.prod)
        ctx.ui.notify(`\u26a0 on PROD context (${current.env}) — guard armed.`, "info");
      else if (!Object.values(envs).some((e) => e.prod))
        ctx.ui.notify(`pi-aws-accounts: ${n} env(s) ready. Prod guard OFF — add "prodAccounts":["<id>"] to ${CONFIG_PATH} to arm it.`, "info");
    } catch {
      /* never break session startup */
    }
  });

  pi.registerCommand("aws-switch", {
    description: "Switch active AWS profile + kube context, e.g. /aws-switch prod",
    handler: async (args: string, ctx: ExtensionContext) => {
      try {
        if (Object.keys(envs).length === 0) await rebuild();
        const names = Object.keys(envs);
        if (names.length === 0) {
          ctx.ui.notify("no AWS/EKS envs detected — nothing to switch (see README).", "info");
          return;
        }
        let name = (args || "").trim();
        if (!name) {
          const labels = names.map(labelFor);
          const picked = await ctx.ui.select("Switch AWS / EKS context", labels);
          if (!picked) return;
          name = names[labels.indexOf(picked)];
        }
        const e = envs[name];
        if (!e) {
          ctx.ui.notify(`unknown env "${name}". available: ${names.join(", ") || "(none)"}`, "error");
          return;
        }
        // Always clear stale static creds so the profile (or default) wins.
        delete process.env.AWS_ACCESS_KEY_ID;
        delete process.env.AWS_SECRET_ACCESS_KEY;
        delete process.env.AWS_SESSION_TOKEN;
        // Unset AWS_PROFILE for context-only envs, so kube never drifts from a stale AWS account.
        if (e.profile) process.env.AWS_PROFILE = e.profile;
        else delete process.env.AWS_PROFILE;
        let note = "";
        if (e.context) {
          try {
            await pexec("kubectl", ["config", "use-context", e.context]);
          } catch {
            note = " (kube switch failed: kubectl or context unavailable)";
          }
        }
        current = { env: name, ...e };
        setStatus(ctx);
        ctx.ui.notify(`switched to ${name}${e.prod ? " \u26a0 [PROD — guard armed]" : ""}${note}`, "info");
      } catch (err: any) {
        ctx.ui.notify("aws-switch failed: " + (err?.message || String(err)), "error");
      }
    },
  });

  pi.registerCommand("aws-whoami", {
    description: "Show active AWS identity + kube context",
    handler: async (_args: string, ctx: ExtensionContext) => {
      try {
        let id: string;
        try {
          id = (await pexec("aws", ["sts", "get-caller-identity", "--output", "text"])).stdout.trim();
        } catch (e: any) {
          id = "aws unavailable (" + (e?.message || e) + ")";
        }
        let kc: string;
        try {
          kc = (await pexec("kubectl", ["config", "current-context"])).stdout.trim();
        } catch (e: any) {
          kc = "kubectl unavailable (" + (e?.message || e) + ")";
        }
        ctx.ui.notify(
          `env: ${current?.env || "none"}\nprofile: ${process.env.AWS_PROFILE || "(default)"}\nidentity: ${id}\nkube-context: ${kc}`,
          "info",
        );
      } catch (err: any) {
        ctx.ui.notify("aws-whoami failed: " + (err?.message || String(err)), "error");
      }
    },
  });

  pi.registerCommand("aws-set-prod", {
    description: "Mark a cluster as production (arms the prod guard for it)",
    handler: async (_args: string, ctx: ExtensionContext) => {
      try {
        if (Object.keys(envs).length === 0) await rebuild();
        const names = Object.keys(envs);
        if (names.length === 0) {
          ctx.ui.notify("no envs detected — nothing to mark.", "info");
          return;
        }
        const labels = names.map(labelFor);
        const picked = await ctx.ui.select("Mark which cluster is PRODUCTION", labels);
        if (!picked) return;
        const key = names[labels.indexOf(picked)];
        const cfg = loadConfig();
        cfg.envs = cfg.envs || {};
        cfg.envs[key] = { ...cfg.envs[key], prod: true };
        mkdirSync(dirname(CONFIG_PATH), { recursive: true });
        writeFileSync(CONFIG_PATH, JSON.stringify(cfg, null, 2) + "\n");
        raw = cfg;
        await rebuild();
        if (current?.env === key) {
          current.prod = true;
          setStatus(ctx);
        }
        ctx.ui.notify(`marked "${key}" as PROD — guard armed. Saved to ${CONFIG_PATH}`, "info");
      } catch (err: any) {
        ctx.ui.notify("aws-set-prod failed: " + (err?.message || String(err)), "error");
      }
    },
  });

  // Prod guardrail. Confirms/blocks mutating aws|kubectl|helm that will run against a prod context.
  // Resolves the REAL kube context at command time, so a switch from ANY terminal is caught.
  pi.on("tool_call", async (event: any, ctx: ExtensionContext) => {
    try {
      if (event.toolName !== "bash") return;
      const cmd: string = event.input?.command || "";
      if (!TOUCHES_CLOUD.test(cmd)) return;

      // Which kube context will this command actually hit?
      //  inline `use-context` wins; else the live current-context (catches external switches);
      //  else the last-known context (if kubectl is unavailable).
      const inline = cmd.match(/kubectl\s+config\s+use-context\s+["']?([^"'\s;&|]+)/);
      let liveCtx = inline?.[1];
      if (!liveCtx) {
        try {
          liveCtx = (await pexec("kubectl", ["config", "current-context"])).stdout.trim();
        } catch {
          liveCtx = current?.context;
        }
      }
      const hit = liveCtx ? Object.entries(envs).find(([, e]) => e.context === liveCtx) : undefined;

      // Sync the indicator to the resolved live context FIRST, so it's truthful even if we block.
      if (liveCtx && current?.context !== liveCtx) {
        current = hit
          ? { env: hit[0], ...hit[1], context: liveCtx, profile: process.env.AWS_PROFILE }
          : { env: "(unmapped)", context: liveCtx, profile: process.env.AWS_PROFILE };
        setStatus(ctx);
      }

      const readOnly = READONLY.test(cmd) && !MUTATING.test(cmd);
      if (hit?.[1].prod && !readOnly) {
        const ok = await ctx.ui.confirm(
          "\u26a0 PROD guard",
          `This runs against PROD (${hit[0]}).\nAllow this command?\n\n${cmd}`,
        );
        if (!ok) return { block: true, reason: "Blocked by pi-aws-accounts prod guard" };
      }
    } catch (err: any) {
      // Fail closed: if the guard itself errors, block rather than silently allow.
      return { block: true, reason: "prod guard error: " + (err?.message || String(err)) };
    }
  });
}