/**
 * pi-input-bar
 *
 * Embeds live status into the input editor's own border lines (no extra rows):
 *
 *   ─ ⌂ folder :: ⎇ branch ● ─────── provider/model :: ✦ effort ─   ← top border
 *   > (editor content)
 *   ────────────────── ⛁ 10%/262k :: ↑29k ↓1.6k R24k :: ⚡34 t/s ─   ← bottom border
 *
 * Implementation: a CustomEditor subclass whose render() takes the default
 * editor output and rewrites the two full-rule border lines. The built-in
 * footer is replaced by a minimal one that only shows extension statuses, so
 * no duplicate stats lines appear below the editor.
 *
 * While streaming, a random working word (Claude Code style) replaces pi's
 * default label, re-rolled on every agent run. Tokens/sec: live ~4 chars/token
 * estimate (marked "~") while streaming, exact usage.output / elapsed after.
 */

import type {
  ExtensionAPI,
  ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { CustomEditor } from "@earendil-works/pi-coding-agent";
import { truncateToWidth, visibleWidth as tuiVisibleWidth } from "@earendil-works/pi-tui";
import { existsSync, readFileSync, statSync, writeFileSync } from "node:fs";
import { basename, dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import {
  bottomBarSegments,
  composeBottomLine,
  composeLine,
  DEFAULT_WORDS,
  loopSummary,
  pickWord,
  replaceEdgeChars,
  stripAnsi,
  tokensPerSecond,
  topBarSegments,
  type ColorName,
  type Colorize,
} from "./render.ts";

// Config lives next to the extension file: ./extensions/input-bar.json
const EXT_DIR = dirname(fileURLToPath(import.meta.url));
const CONFIG_PATH = join(EXT_DIR, "input-bar.json");

const DEFAULTS = {
  TOP: true, // embed folder/branch/model/effort into the editor's top border
  BOTTOM: true, // embed context/tokens/tps into the editor's bottom border
  TPS: true, // show tokens/sec
  ANIM: true, // random working word while streaming
  ICONS: true, // glyphs (⌂ ⎇ ⛁ ⚡) in the bars
  WORDS: "", // comma-separated override for the working-word pool
  LOOPS: true, // loop-police detection summary in the top border (if installed)
};

const cfg: typeof DEFAULTS = (() => {
  if (!existsSync(CONFIG_PATH)) {
    try {
      writeFileSync(CONFIG_PATH, JSON.stringify(DEFAULTS, null, 2) + "\n", "utf-8");
    } catch {
      // Read-only install location — keep defaults in memory
    }
  }
  try {
    return { ...DEFAULTS, ...JSON.parse(readFileSync(CONFIG_PATH, "utf-8")) };
  } catch {
    return { ...DEFAULTS };
  }
})();

function wordPool(): string[] {
  const custom = cfg.WORDS.split(",").map((w) => w.trim()).filter(Boolean);
  return custom.length ? custom : DEFAULT_WORDS;
}

// Fallback branch resolution for when the FooterDataProvider isn't captured yet
// (it only reaches us through the setFooter factory). Cached with a short TTL.
let branchCache: { value: string | null; at: number } = { value: null, at: 0 };

function gitBranchFallback(cwd: string): string | null {
  if (Date.now() - branchCache.at < 5000) return branchCache.value;
  let value: string | null = null;
  try {
    let dir = cwd;
    for (;;) {
      const gitPath = join(dir, ".git");
      if (existsSync(gitPath)) {
        let headPath = join(gitPath, "HEAD");
        if (statSync(gitPath).isFile()) {
          // Worktree/submodule: .git is a file pointing at the real git dir
          const m = readFileSync(gitPath, "utf-8").match(/^gitdir:\s*(.+)$/m);
          if (m) headPath = join(m[1].trim(), "HEAD");
        }
        const head = readFileSync(headPath, "utf-8").trim();
        const ref = head.match(/^ref:\s*refs\/heads\/(.+)$/);
        value = ref ? ref[1] : "detached";
        break;
      }
      const parent = dirname(dir);
      if (parent === dir) break;
      dir = parent;
    }
  } catch {
    value = null;
  }
  branchCache = { value, at: Date.now() };
  return value;
}

/** Total characters of text + thinking in an assistant message. */
function messageChars(message: any): number {
  let chars = 0;
  for (const block of message?.content ?? []) {
    if (typeof block?.text === "string") chars += block.text.length;
    else if (typeof block?.thinking === "string") chars += block.thinking.length;
  }
  return chars;
}

export default function (pi: ExtensionAPI) {
  // Latest context seen by any handler — its methods are stable runtime bindings.
  let ctxRef: ExtensionContext | undefined;
  // FooterDataProvider captured from the setFooter factory; the editor reads it too.
  let footerData: { getGitBranch(): string | null } | undefined;
  // TUI handle captured from the editor factory, to repaint on bus events.
  let tuiRef: { requestRender(): void } | undefined;

  // Per-session loop-police detections, keyed by detector name (payload.event).
  // If pi-loop-police isn't installed the event never fires and the map stays
  // empty — the bar renders exactly as before.
  const loopCounts = new Map<string, number>();
  (pi as any).events?.on?.("loop-police:detection", (data: any) => {
    const kind = typeof data?.event === "string" ? data.event : "unknown";
    loopCounts.set(kind, (loopCounts.get(kind) ?? 0) + 1);
    tuiRef?.requestRender();
  });

  // t/s tracking. The clock starts at the FIRST streamed token, not at
  // message_start — otherwise prompt-processing time (long on local servers)
  // dilutes the average and the value creeps up without ever reaching the
  // real generation speed.
  let streaming = false;
  let messageStartMs = 0;
  let firstTokenMs = 0;
  let liveChars = 0;
  let lastTps: number | undefined;
  // Prefill (prompt processing) speed: tokens processed / time until first token
  let lastPps: number | undefined;
  // Live-estimate calibration: learned chars/token ratio from finished messages
  let charsPerToken = 4;

  function snapshotTps(): { tps: number | undefined; live: boolean } {
    if (streaming && firstTokenMs) {
      const est = tokensPerSecond(liveChars / charsPerToken, Date.now() - firstTokenMs);
      if (est !== undefined) return { tps: est, live: true };
    }
    return { tps: lastTps, live: false };
  }

  function makeColorize(theme: { fg(color: any, text: string): string }, effortLevel: string): Colorize {
    const thinkingKey = "thinking" + effortLevel.charAt(0).toUpperCase() + effortLevel.slice(1);
    return (color: ColorName, text: string) =>
      theme.fg(color === "thinking" ? thinkingKey : color, text);
  }

  function currentEffort(): string {
    try {
      return pi.getThinkingLevel();
    } catch {
      return "off";
    }
  }

  /** Top border line with folder/branch/model/effort embedded. */
  function topLine(width: number, theme: any): string {
    const ctx = ctxRef;
    const effort = currentEffort();
    const { left, right } = topBarSegments({
      folder: basename(ctx?.cwd || process.cwd()) || "/",
      branch: footerData?.getGitBranch() ?? gitBranchFallback(ctx?.cwd || process.cwd()),
      streaming,
      provider: ctx?.model?.provider,
      modelId: ctx?.model?.id,
      effort: ctx?.model?.reasoning ? effort : undefined,
      loops: cfg.LOOPS ? loopSummary(loopCounts) : undefined,
      icons: cfg.ICONS,
    });
    const line = composeLine(left, right, width, makeColorize(theme, effort), tuiVisibleWidth);
    return truncateToWidth(line, width);
  }

  /** Bottom border line with context/token/tps stats embedded, right-aligned. */
  function bottomLine(width: number, theme: any): string {
    const ctx = ctxRef;
    let input = 0, output = 0, cacheRead = 0, cacheWrite = 0, cost = 0;
    try {
      for (const entry of ctx?.sessionManager.getEntries() ?? []) {
        if (entry.type === "message" && (entry as any).message.role === "assistant") {
          const u = (entry as any).message.usage;
          input += u.input;
          output += u.output;
          cacheRead += u.cacheRead;
          cacheWrite += u.cacheWrite;
          cost += u.cost.total;
        }
      }
    } catch {
      // Session shape changed — leave totals at 0
    }
    const usage = ctx?.getContextUsage();
    const { tps, live } = snapshotTps();
    const segs = bottomBarSegments({
      percent: usage?.percent ?? null,
      contextWindow: usage?.contextWindow ?? ctx?.model?.contextWindow ?? 0,
      input, output, cacheRead, cacheWrite, cost,
      tps: cfg.TPS ? tps : undefined,
      tpsLive: live,
      pps: cfg.TPS ? lastPps : undefined,
      icons: cfg.ICONS,
    });
    const line = composeBottomLine(segs, width, makeColorize(theme, currentEffort()), tuiVisibleWidth);
    return truncateToWidth(line, width);
  }

  /**
   * Editor that rewrites its full-rule border lines with the info bars.
   * With both bars active the box is closed: rounded corners on the border
   * lines and │ sides painted over the paddingX columns of content lines.
   */
  class BarEditor extends CustomEditor {
    /**
     * pi copies the default editor's paddingX onto custom editors right after
     * the factory runs (setCustomEditorComponent), which would collapse our
     * side columns to the user's setting (often 0). Enforce a minimum of 2 so
     * the │ sides always have a spare space column next to the content/cursor.
     */
    setPaddingX(padding: number): void {
      super.setPaddingX(Math.max(2, Number.isFinite(padding) ? padding : 2));
    }

    render(width: number): string[] {
      const lines = super.render(width);
      const theme = ctxRef?.ui?.theme;
      if (!theme || width < 10) return lines;
      const rule = "─".repeat(width);
      const isRule = (l: string) => stripAnsi(l) === rule;
      // Top border is the first full-rule line (absent while scrolled — the
      // "↑ N more" indicator takes its place and is left untouched).
      const topIdx = cfg.TOP ? lines.findIndex(isRule) : -1;
      // Bottom border is the last full-rule line (autocomplete rows follow it).
      let bottomIdx = -1;
      if (cfg.BOTTOM) {
        for (let i = lines.length - 1; i > topIdx; i--) {
          if (isRule(lines[i])) {
            bottomIdx = i;
            break;
          }
        }
      }
      const closed = topIdx !== -1 && bottomIdx !== -1;
      if (topIdx !== -1) {
        const line = topLine(width, theme);
        lines[topIdx] = closed ? replaceEdgeChars(line, "╭", "╮") : line;
      }
      if (bottomIdx !== -1) {
        const line = bottomLine(width, theme);
        lines[bottomIdx] = closed ? replaceEdgeChars(line, "╰", "╯") : line;
      }
      if (closed) {
        // Paint │ over the outermost padding column of the rows in between.
        const side = theme.fg("dim", "│");
        for (let i = topIdx + 1; i < bottomIdx; i++) {
          const l = lines[i];
          if (l.startsWith(" ") && l.endsWith(" ")) {
            lines[i] = side + l.slice(1, -1) + side;
          }
        }
      }
      return lines;
    }
  }

  function installBars(ctx: ExtensionContext) {
    if (cfg.TOP || cfg.BOTTOM) {
      // paddingX: 2 keeps a spare space column on each side of the content so
      // the │ sides can be painted without touching text or the cursor cell.
      ctx.ui.setEditorComponent(
        (tui, theme, keybindings) => {
          tuiRef = tui;
          return new BarEditor(tui, theme, keybindings, { paddingX: 2 });
        },
      );
    } else {
      ctx.ui.setEditorComponent(undefined);
    }
    if (cfg.BOTTOM) {
      // Minimal footer: kills the built-in stats lines (now embedded in the
      // editor border) but keeps extension statuses and the branch watcher.
      ctx.ui.setFooter((tui, theme, fd) => {
        footerData = fd;
        const unsub = fd.onBranchChange(() => tui.requestRender());
        return {
          render(width: number): string[] {
            const statuses: ReadonlyMap<string, string> = fd.getExtensionStatuses();
            if (statuses.size === 0) return [];
            const line = Array.from(statuses.entries())
              .sort(([a], [b]) => a.localeCompare(b))
              .map(([, text]) => text.replace(/[\r\n\t]/g, " ").trim())
              .join(" ");
            return [truncateToWidth(theme.fg("dim", line), width)];
          },
          invalidate() {},
          dispose() { unsub(); },
        };
      });
    } else {
      ctx.ui.setFooter(undefined);
    }
  }

  pi.on("session_start", (_event, ctx) => {
    ctxRef = ctx;
    loopCounts.clear();
    if (ctx.mode !== "tui") return;
    installBars(ctx);
  });

  pi.on("agent_start", (_event, ctx) => {
    ctxRef = ctx;
    if (ctx.hasUI && cfg.ANIM) {
      ctx.ui.setWorkingMessage(pickWord(wordPool()));
    }
  });

  pi.on("message_start", (event, ctx) => {
    ctxRef = ctx;
    if (event.message.role !== "assistant") return;
    streaming = true;
    messageStartMs = Date.now();
    firstTokenMs = 0;
    liveChars = 0;
  });

  pi.on("message_update", (event) => {
    if (event.message.role !== "assistant") return;
    const chars = messageChars(event.message);
    if (!firstTokenMs && chars > 0) firstTokenMs = Date.now();
    liveChars = chars;
  });

  pi.on("message_end", (event, ctx) => {
    ctxRef = ctx;
    if (event.message.role !== "assistant") return;
    streaming = false;
    if (!firstTokenMs) return;
    const usage = (event.message as any).usage;
    const outTokens = usage?.output ?? 0;
    if (outTokens > 0 && liveChars > 0) {
      charsPerToken = Math.min(8, Math.max(1, liveChars / outTokens));
    }
    const tps = tokensPerSecond(outTokens, Date.now() - firstTokenMs);
    if (tps !== undefined) lastTps = tps;
    // Prefill: prompt tokens actually processed (cached reads cost ~nothing)
    // over the message_start → first-token window. Skip tiny windows where
    // network/server overhead dominates the measurement.
    const prefillTokens = (usage?.input ?? 0) + (usage?.cacheWrite ?? 0);
    const prefillMs = firstTokenMs - messageStartMs;
    if (prefillTokens > 0 && prefillMs > 150) {
      lastPps = prefillTokens / (prefillMs / 1000);
    }
  });

  pi.on("agent_end", (_event, ctx) => {
    ctxRef = ctx;
    streaming = false;
  });

  pi.registerCommand("input-bar", {
    description: "Configure the input bar (show | set KEY=VAL | save | reset)",
    handler: async (args, ctx) => {
      ctxRef = ctx;
      const trimmed = args.trim();
      const notify = (msg: string) => ctx.ui.notify(msg, "info");

      if (!trimmed || trimmed === "show") {
        notify(
          Object.entries(cfg)
            .map(([k, v]) => `${k}=${v}`)
            .join(" "),
        );
        return;
      }
      if (trimmed === "save") {
        try {
          writeFileSync(CONFIG_PATH, JSON.stringify(cfg, null, 2) + "\n", "utf-8");
          notify(`saved ${CONFIG_PATH}`);
        } catch (e) {
          ctx.ui.notify(`could not save: ${e}`, "error");
        }
        return;
      }
      if (trimmed === "reset") {
        Object.assign(cfg, DEFAULTS);
        installBars(ctx);
        notify("input-bar config reset (session only; use save to persist)");
        return;
      }
      const match = trimmed.match(/^set\s+([A-Z]+)\s*=\s*(.*)$/);
      if (!match) {
        ctx.ui.notify("usage: /input-bar [show | set KEY=VAL | save | reset]", "warning");
        return;
      }
      const [, key, raw] = match;
      if (!(key in cfg)) {
        ctx.ui.notify(`unknown key ${key} (${Object.keys(cfg).join(", ")})`, "error");
        return;
      }
      (cfg as any)[key] = typeof (cfg as any)[key] === "boolean" ? raw === "true" || raw === "1" : raw;
      installBars(ctx);
      if (key === "ANIM" && !cfg.ANIM) ctx.ui.setWorkingMessage();
      notify(`${key}=${(cfg as any)[key]} (session only; use /input-bar save to persist)`);
    },
  });
}