import { isAbsolute, relative, resolve, sep } from "node:path";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { KEY_SWITCH_ENTRY_TYPE, getActiveParleyKeyProfile, shouldShowParleyKey } from "./keys.js";

// A plain-language footer for Pi sessions that use the MIT-parley provider.
// The built-in footer packs everything into one dim line with terse markers
// (↑ ↓ R W CH). This one spells things out on two lines and colors the values
// students care about: thinking level, context fill, cache hit rate, and cost.
//
// It is installed only while an MIT-parley model is active and restores Pi's
// own footer when the user switches to another provider. `/parley-footer`
// toggles it off and on.

const PROVIDER_ID = "MIT-parley";

type Theme = Parameters<Parameters<ExtensionContext["ui"]["setFooter"]>[0] & Function>[1];
type FooterData = Parameters<Parameters<ExtensionContext["ui"]["setFooter"]>[0] & Function>[2];

interface Totals {
  input: number;
  output: number;
  cacheRead: number;
  cacheWrite: number;
  cost: number;
}

// Strip ANSI SGR sequences so padding math sees the visible text only.
// eslint-disable-next-line no-control-regex
const ANSI_RE = /\x1b\[[0-9;]*m/g;
function visibleWidth(text: string): number {
  return text.replace(ANSI_RE, "").length;
}

function truncate(text: string, width: number): string {
  if (visibleWidth(text) <= width) return text;
  const plain = text.replace(ANSI_RE, "");
  return width > 3 ? `${plain.slice(0, width - 3)}...` : plain.slice(0, width);
}

export function formatTokens(count: number): string {
  if (count < 1000) return count.toString();
  if (count < 10000) return `${(count / 1000).toFixed(1)}k`;
  if (count < 1000000) return `${Math.round(count / 1000)}k`;
  if (count < 10000000) return `${(count / 1000000).toFixed(1)}M`;
  return `${Math.round(count / 1000000)}M`;
}

function formatCwd(cwd: string, home: string | undefined): string {
  if (!home) return cwd;
  const rel = relative(resolve(home), resolve(cwd));
  const inside = rel === "" || (rel !== ".." && !rel.startsWith(`..${sep}`) && !isAbsolute(rel));
  if (!inside) return cwd;
  return rel === "" ? "~" : `~${sep}${rel}`;
}

// Sum usage across the current branch, like the built-in footer, and remember
// the cache hit rate of the most recent assistant turn. A /parley-key switch
// leaves a marker entry; counting restarts after the latest one so the cost
// shown is what the active key has been charged.
function collectTotals(ctx: ExtensionContext): { totals: Totals; lastCacheHit: number | undefined; turns: number } {
  let totals: Totals = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 };
  let lastCacheHit: number | undefined;
  let turns = 0;
  const add = (usage: any) => {
    if (!usage) return;
    totals.input += usage.input || 0;
    totals.output += usage.output || 0;
    totals.cacheRead += usage.cacheRead || 0;
    totals.cacheWrite += usage.cacheWrite || 0;
    totals.cost += usage.cost?.total || 0;
  };
  for (const entry of ctx.sessionManager.getBranch() as any[]) {
    if (entry.type === "custom" && entry.customType === KEY_SWITCH_ENTRY_TYPE) {
      totals = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 };
      lastCacheHit = undefined;
      turns = 0;
    } else if (entry.type === "message" && entry.message?.role === "assistant") {
      add(entry.message.usage);
      turns++;
      const u = entry.message.usage;
      const prompt = (u?.input || 0) + (u?.cacheRead || 0) + (u?.cacheWrite || 0);
      lastCacheHit = prompt > 0 ? ((u.cacheRead || 0) / prompt) * 100 : undefined;
    } else if (entry.type === "message" && entry.message?.role === "toolResult") {
      add(entry.message.usage);
    } else if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) {
      add(entry.usage);
    }
  }
  return { totals, lastCacheHit, turns };
}

function thinkingColor(level: string): any {
  switch (level) {
    case "minimal":
      return "thinkingMinimal";
    case "low":
      return "thinkingLow";
    case "medium":
      return "thinkingMedium";
    case "high":
      return "thinkingHigh";
    case "xhigh":
      return "thinkingXhigh";
    case "max":
      return "thinkingMax";
    default:
      return "thinkingOff";
  }
}

export function renderParleyFooter(ctx: ExtensionContext, theme: Theme, footerData: FooterData, width: number): string[] {
  const dim = (s: string) => theme.fg("dim", s);
  const label = (s: string) => dim(`${s} `);
  const sepStr = dim("  ·  ");

  // Line 1: where we are.
  let pwd = formatCwd(ctx.sessionManager.getCwd(), process.env.HOME || process.env.USERPROFILE);
  const branch = footerData.getGitBranch();
  if (branch) pwd = `${pwd} (${branch})`;
  const sessionName = ctx.sessionManager.getSessionName();
  if (sessionName) pwd = `${pwd} • ${sessionName}`;

  // Line 2: model and thinking.
  const model = ctx.model;
  const modelName = model ? `${model.name || model.id}` : "no model";
  const parts2 = [label("model") + theme.fg("accent", modelName) + dim(" via MIT Parley")];
  const keyProfile = getActiveParleyKeyProfile();
  if (keyProfile && shouldShowParleyKey()) {
    parts2.push(label("key") + theme.fg("accent", keyProfile.label) + dim("  (/parley-key)"));
  }
  if (model?.reasoning) {
    const level = ctx.thinkingLevel || "off";
    parts2.push(label("thinking") + theme.fg(thinkingColor(level), level) + dim("  (Shift+Tab)"));
  }

  // Line 3: context, tokens, cache, cost.
  const { totals, lastCacheHit, turns } = collectTotals(ctx);
  const usage = ctx.getContextUsage();
  const contextWindow = usage?.contextWindow ?? model?.contextWindow ?? 0;
  const pct = usage?.percent ?? null;
  const pctText = pct === null ? "?" : `${pct.toFixed(1)}%`;
  const pctColor = pct !== null && pct > 90 ? "error" : pct !== null && pct > 70 ? "warning" : "text";
  const parts3 = [label("context") + theme.fg(pctColor, pctText) + dim(` of ${formatTokens(contextWindow)}`)];

  parts3.push(
    label("sent") + formatTokens(totals.input + totals.cacheRead + totals.cacheWrite) + sepStr + label("received") + formatTokens(totals.output),
  );

  if (turns > 0) {
    let cacheText: string;
    if (lastCacheHit === undefined || (totals.cacheRead === 0 && totals.cacheWrite === 0)) {
      cacheText = dim("none");
    } else {
      const color = lastCacheHit >= 50 ? "success" : lastCacheHit > 0 ? "warning" : "muted";
      cacheText = theme.fg(color, `${lastCacheHit.toFixed(0)}% hit`);
    }
    parts3.push(label("cache") + cacheText);
  }

  const costColor = totals.cost >= 2 ? "error" : totals.cost >= 0.5 ? "warning" : "success";
  parts3.push(label("est. cost") + theme.fg(costColor, `$${totals.cost.toFixed(3)}`));

  const lines = [
    truncate(dim(pwd), width),
    truncate(parts2.join(sepStr), width),
    truncate(parts3.join(sepStr), width),
  ];

  // Statuses from other extensions, as the built-in footer shows them. The key
  // status is ours and already on line 2.
  const statuses = footerData.getExtensionStatuses();
  if (statuses.size > 0) {
    const text = Array.from(statuses.entries())
      .filter(([key]) => key !== "parley-key")
      .sort(([a], [b]) => a.localeCompare(b))
      .map(([, t]) => t.replace(/[\r\n\t]/g, " ").replace(/ +/g, " ").trim())
      .join(" ");
    lines.push(truncate(text, width));
  }
  return lines;
}

export function registerParleyFooter(pi: ExtensionAPI): void {
  let latest: ExtensionContext | undefined;
  let enabled = true;
  let installed = false;

  const sync = (ctx: ExtensionContext) => {
    latest = ctx;
    if (!ctx.hasUI || ctx.mode !== "tui") return;
    const wanted = enabled && ctx.model?.provider === PROVIDER_ID;
    if (wanted && !installed) {
      ctx.ui.setFooter((tui, theme, footerData) => ({
        render(width: number) {
          return latest ? renderParleyFooter(latest, theme, footerData, width) : [];
        },
        invalidate() {},
        dispose: footerData.onBranchChange(() => tui.requestRender()),
      }));
      installed = true;
    } else if (!wanted && installed) {
      ctx.ui.setFooter(undefined);
      installed = false;
    }
  };

  pi.on("session_start", (_event, ctx) => sync(ctx));
  pi.on("model_select", (_event, ctx) => sync(ctx));
  pi.on("thinking_level_select", (_event, ctx) => sync(ctx));
  pi.on("turn_end", (_event, ctx) => sync(ctx));
  pi.on("agent_end", (_event, ctx) => sync(ctx));

  pi.registerCommand("parley-footer", {
    description: "Toggle the MIT Parley footer (model, thinking, context, cache, cost)",
    handler: async (_args, ctx) => {
      enabled = !enabled;
      sync(ctx);
      ctx.ui.notify(`MIT Parley footer ${enabled ? "on" : "off"}`, "info");
    },
  });
}