import type {
  ExtensionAPI,
  ExtensionContext
} from "@earendil-works/pi-coding-agent";
import { Key, type KeyId } from "@earendil-works/pi-tui";
import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";

import { dirname, join } from "node:path";

const APP_NAME = "PreApeXis";
const FALLBACK_APP_VERSION = "dev";

const WELCOME_TEXT = "Welcome back!";
const LOGO_LINES = [
  "                   ░▒▓███▓▒░                   ",
  "                  ░▒▓█████▓▒░                  ",
  "                   ░▒▓███▓▒░                   "
];

const CLEAR_HEADER_ON_AGENT_START = false;

const TIPS_HEADER = "Tips for getting started";
const TIPS = ["Run /init to create a AGENTS.md file with instructions for Pi"];

const USAGE_STATUS_ID = "preapexis-usage";

// F8 is usually the most reliable.
// Alt+U and Ctrl+Shift+U are included as extra attempts.
const USAGE_SHORTCUTS = [
  "f8",
  Key.alt("u"),
  Key.ctrlShift("u")
] satisfies readonly KeyId[];

interface Theme {
  fg(color: string, text: string): string;
}

interface HeaderRenderer {
  render(width: number): string[];
  invalidate(): void;
}

interface HeaderUI {
  setHeader?: (
    factory?: (tui: unknown, theme: Theme) => HeaderRenderer
  ) => void;
}

interface UsageCost {
  input?: number;
  output?: number;
  cacheRead?: number;
  cacheWrite?: number;
  total?: number;
}

interface Usage {
  input?: number;
  output?: number;
  cacheRead?: number;
  cacheWrite?: number;
  totalTokens?: number;
  cost?: UsageCost;
}

interface AssistantMessageLike {
  role?: string;
  provider?: string;
  model?: string;
  usage?: Usage;
}

interface UsageBucket {
  usd: number;
  input: number;
  output: number;
  cacheRead: number;
  cacheWrite: number;
  totalTokens: number;
}

interface UsageStore {
  version: 1;
  updatedAt: string;
  days: Record<string, UsageBucket>;
  lifetime: UsageBucket;
  byModel: Record<string, UsageBucket>;
}

function getHeaderUI(ctx: ExtensionContext): HeaderUI {
  return ctx.ui as unknown as HeaderUI;
}

function stripAnsi(input: string): string {
  return input.replace(/\x1b\[[0-9;]*m/g, "");
}

function visibleLength(input: string): number {
  return stripAnsi(input).length;
}

function centerVisible(input: string, width: number): string {
  const len = visibleLength(input);
  if (len >= width) return input;

  const left = Math.floor((width - len) / 2);
  const right = width - len - left;

  return " ".repeat(left) + input + " ".repeat(right);
}

function padEndVisible(input: string, width: number): string {
  const len = visibleLength(input);
  if (len >= width) return input;

  return input + " ".repeat(width - len);
}

function truncateVisible(input: string, width: number): string {
  const len = visibleLength(input);
  if (len <= width) return input;

  const stripped = stripAnsi(input);

  if (width <= 1) {
    return "…";
  }

  return stripped.slice(0, Math.max(0, width - 1)) + "…";
}

function getWorkingDirectoryLabel(): string {
  const cwd = process.cwd();
  const home = process.env.HOME ?? process.env.USERPROFILE ?? "";

  if (home && cwd.startsWith(home)) {
    return `~${cwd.slice(home.length)}`;
  }

  return cwd;
}

function findPackageJson(startDir = process.cwd()): string | undefined {
  let dir = startDir;

  while (true) {
    const candidate = join(dir, "package.json");

    if (existsSync(candidate)) {
      return candidate;
    }

    const parent = dirname(dir);

    if (parent === dir) {
      return undefined;
    }

    dir = parent;
  }
}

function getOwnPackageVersion(): string | undefined {
  const packageJsonPath = findPackageJson();

  if (!packageJsonPath) {
    return undefined;
  }

  try {
    const json = JSON.parse(readFileSync(packageJsonPath, "utf8")) as {
      version?: unknown;
    };

    return typeof json.version === "string" ? json.version : undefined;
  } catch {
    return undefined;
  }
}

function makeVersionLabel(): string {
  const ownVersion = getOwnPackageVersion() ?? FALLBACK_APP_VERSION;
  return `${APP_NAME} v${ownVersion}`;
}

function emptyBucket(): UsageBucket {
  return {
    usd: 0,
    input: 0,
    output: 0,
    cacheRead: 0,
    cacheWrite: 0,
    totalTokens: 0
  };
}

function safeNumber(value: unknown): number {
  return typeof value === "number" && Number.isFinite(value) ? value : 0;
}

function normalizeBucket(value: Partial<UsageBucket> | undefined): UsageBucket {
  return {
    usd: safeNumber(value?.usd),
    input: safeNumber(value?.input),
    output: safeNumber(value?.output),
    cacheRead: safeNumber(value?.cacheRead),
    cacheWrite: safeNumber(value?.cacheWrite),
    totalTokens: safeNumber(value?.totalTokens)
  };
}

function getLocalDayKey(date = new Date()): string {
  const year = date.getFullYear();
  const month = String(date.getMonth() + 1).padStart(2, "0");
  const day = String(date.getDate()).padStart(2, "0");

  return `${year}-${month}-${day}`;
}

function getUsageFilePath(): string {
  const home = process.env.HOME ?? process.env.USERPROFILE ?? process.cwd();
  return join(home, ".preapexis", "usage.json");
}

function loadUsageStore(): UsageStore {
  const filePath = getUsageFilePath();

  if (!existsSync(filePath)) {
    return {
      version: 1,
      updatedAt: new Date().toISOString(),
      days: {},
      lifetime: emptyBucket(),
      byModel: {}
    };
  }

  try {
    const parsed = JSON.parse(
      readFileSync(filePath, "utf8")
    ) as Partial<UsageStore>;

    const days: Record<string, UsageBucket> = {};
    for (const [day, bucket] of Object.entries(parsed.days ?? {})) {
      days[day] = normalizeBucket(bucket);
    }

    const byModel: Record<string, UsageBucket> = {};
    for (const [model, bucket] of Object.entries(parsed.byModel ?? {})) {
      byModel[model] = normalizeBucket(bucket);
    }

    return {
      version: 1,
      updatedAt:
        typeof parsed.updatedAt === "string"
          ? parsed.updatedAt
          : new Date().toISOString(),
      days,
      lifetime: normalizeBucket(parsed.lifetime),
      byModel
    };
  } catch {
    return {
      version: 1,
      updatedAt: new Date().toISOString(),
      days: {},
      lifetime: emptyBucket(),
      byModel: {}
    };
  }
}

function saveUsageStore(store: UsageStore): void {
  const filePath = getUsageFilePath();

  mkdirSync(dirname(filePath), { recursive: true });
  writeFileSync(filePath, JSON.stringify(store, null, 2));
}

function addUsage(bucket: UsageBucket, usage: Usage): void {
  const input = safeNumber(usage.input);
  const output = safeNumber(usage.output);
  const cacheRead = safeNumber(usage.cacheRead);
  const cacheWrite = safeNumber(usage.cacheWrite);

  const totalTokens =
    safeNumber(usage.totalTokens) || input + output + cacheRead + cacheWrite;

  bucket.usd += safeNumber(usage.cost?.total);
  bucket.input += input;
  bucket.output += output;
  bucket.cacheRead += cacheRead;
  bucket.cacheWrite += cacheWrite;
  bucket.totalTokens += totalTokens;
}

function resetBucket(bucket: UsageBucket): void {
  bucket.usd = 0;
  bucket.input = 0;
  bucket.output = 0;
  bucket.cacheRead = 0;
  bucket.cacheWrite = 0;
  bucket.totalTokens = 0;
}

function formatUsd(value: number): string {
  if (value <= 0) {
    return "$0.0000";
  }

  if (value < 0.01) {
    return `$${value.toFixed(4)}`;
  }

  return `$${value.toFixed(2)}`;
}

function formatTokens(value: number): string {
  return Math.round(value).toLocaleString();
}

function renderWelcomeLines(
  theme: Theme,
  version: string,
  width: number
): string[] {
  const panelWidth = Math.min(Math.max(width - 4, 88), 132);
  const innerWidth = panelWidth - 4;

  const leftWidth = Math.floor(innerWidth * 0.42);
  const rightWidth = innerWidth - leftWidth - 3;
  const hPadding = Math.max(0, Math.floor((width - panelWidth) / 2));

  const lines: string[] = [];

  lines.push(" ".repeat(hPadding));

  const title = `─── ${version} `;
  const topFill = "─".repeat(Math.max(0, panelWidth - title.length - 2));

  lines.push(" ".repeat(hPadding) + theme.fg("accent", `╭${title}${topFill}╮`));

  const cwd = getWorkingDirectoryLabel();

  const leftRows = [
    "",
    theme.fg("muted", WELCOME_TEXT),
    "",
    ...LOGO_LINES,
    "",
    theme.fg("dim", truncateVisible(cwd, leftWidth)),
    ""
  ];

  const rightRows = [
    "",
    theme.fg("accent", TIPS_HEADER),
    "",
    theme.fg("muted", truncateVisible(TIPS[0] ?? "", rightWidth)),
    "",
    theme.fg("dim", "─".repeat(rightWidth)),
    "",
    theme.fg("dim", "Model, effort, mode, and context stay in the footer."),
    theme.fg("dim", "Press F8 to view usage and estimated spend."),
    ""
  ];

  const maxRows = Math.max(leftRows.length, rightRows.length);

  for (let i = 0; i < maxRows; i++) {
    const left = leftRows[i] ?? "";
    const right = rightRows[i] ?? "";

    const leftCell = centerVisible(left, leftWidth);
    const rightCell = padEndVisible(right, rightWidth);

    lines.push(" ".repeat(hPadding) + `│ ${leftCell} │ ${rightCell} │`);
  }

  lines.push(
    " ".repeat(hPadding) + theme.fg("accent", `╰${"─".repeat(panelWidth - 2)}╯`)
  );

  return lines;
}

export default function welcomeExtension(pi: ExtensionAPI): void {
  let shown = false;

  const usageStore = loadUsageStore();
  const sessionUsage = emptyBucket();

  function getTodayUsage(): UsageBucket {
    const today = getLocalDayKey();
    usageStore.days[today] ??= emptyBucket();
    return usageStore.days[today];
  }

  function makeUsageFooterText(): string {
    const today = getTodayUsage();

    return `Today ${formatUsd(today.usd)} • Session ${formatUsd(
      sessionUsage.usd
    )} • F8`;
  }

  function updateUsageFooter(ctx: ExtensionContext): void {
    if (!ctx.hasUI) {
      return;
    }

    ctx.ui.setStatus(USAGE_STATUS_ID, makeUsageFooterText());
  }

  function recordAssistantUsage(message: AssistantMessageLike): void {
    if (message.role !== "assistant" || !message.usage) {
      return;
    }

    const today = getLocalDayKey();
    usageStore.days[today] ??= emptyBucket();

    addUsage(sessionUsage, message.usage);
    addUsage(usageStore.days[today], message.usage);
    addUsage(usageStore.lifetime, message.usage);

    const provider = message.provider ?? "unknown-provider";
    const model = message.model ?? "unknown-model";
    const modelKey = `${provider}/${model}`;

    usageStore.byModel[modelKey] ??= emptyBucket();
    addUsage(usageStore.byModel[modelKey], message.usage);

    usageStore.updatedAt = new Date().toISOString();
    saveUsageStore(usageStore);
  }

  function makeUsageDetailsText(): string {
    const today = getTodayUsage();

    const modelLines = Object.entries(usageStore.byModel)
      .sort(([, a], [, b]) => b.usd - a.usd)
      .slice(0, 5)
      .map(
        ([model, bucket]) =>
          `- ${model}: ${formatUsd(bucket.usd)} • ${formatTokens(
            bucket.totalTokens
          )} tokens`
      );

    return [
      "PreApeXis Usage",
      "",
      `Session: ${formatUsd(sessionUsage.usd)} • ${formatTokens(
        sessionUsage.totalTokens
      )} tokens`,
      `Today: ${formatUsd(today.usd)} • ${formatTokens(
        today.totalTokens
      )} tokens`,
      `Lifetime: ${formatUsd(usageStore.lifetime.usd)} • ${formatTokens(
        usageStore.lifetime.totalTokens
      )} tokens`,
      "",
      "Top models:",
      ...(modelLines.length > 0
        ? modelLines
        : ["- No model usage recorded yet."]),
      "",
      "Shortcuts:",
      "- F8",
      "- Alt+U",
      "- Ctrl+Shift+U",
      "",
      `Saved to: ${getUsageFilePath()}`
    ].join("\n");
  }

  async function showUsagePopup(ctx: ExtensionContext): Promise<void> {
    if (!ctx.hasUI) {
      return;
    }

    ctx.ui.notify(makeUsageDetailsText(), "info");
    updateUsageFooter(ctx);
  }

  function setPreApexisHeader(ctx: ExtensionContext): void {
    if (!ctx.hasUI || ctx.mode !== "tui") {
      return;
    }

    const ui = getHeaderUI(ctx);
    const version = makeVersionLabel();

    if (!ui.setHeader) {
      return;
    }

    ui.setHeader((_tui, theme) => ({
      render(width: number): string[] {
        return renderWelcomeLines(theme, version, width);
      },
      invalidate(): void {
        // no-op
      }
    }));
  }

  function clearPreApexisHeader(ctx: ExtensionContext): void {
    if (!ctx.hasUI || ctx.mode !== "tui") {
      return;
    }

    const ui = getHeaderUI(ctx);

    if (!ui.setHeader) {
      return;
    }

    ui.setHeader((_tui, _theme) => ({
      render(): string[] {
        return [];
      },
      invalidate(): void {
        // no-op
      }
    }));
  }

  async function showWelcome(ctx: ExtensionContext): Promise<void> {
    if (!ctx.hasUI || ctx.mode !== "tui") {
      return;
    }

    shown = true;
    setPreApexisHeader(ctx);
  }

  function hideWelcome(ctx: ExtensionContext): void {
    if (!shown) {
      return;
    }

    shown = false;
    clearPreApexisHeader(ctx);
  }

  pi.on("session_start", async (_event, ctx) => {
    updateUsageFooter(ctx);
    await showWelcome(ctx);
  });

  pi.on("agent_start", async (_event, ctx) => {
    if (CLEAR_HEADER_ON_AGENT_START) {
      hideWelcome(ctx);
      return;
    }

    updateUsageFooter(ctx);

    if (shown) {
      setPreApexisHeader(ctx);
    }
  });

  pi.on("message_end", async (event, ctx) => {
    recordAssistantUsage(event.message as AssistantMessageLike);
    updateUsageFooter(ctx);
  });

  for (const shortcut of USAGE_SHORTCUTS) {
    pi.registerShortcut(shortcut, {
      description: "Show PreApeXis usage and spend",
      handler: async (ctx) => {
        await showUsagePopup(ctx);
      }
    });
  }

  // Keep this as a fallback/debug command.
  // If /usage works but F8 does not, the issue is terminal key delivery.
  pi.registerCommand("usage", {
    description: "Show PreApeXis token usage and spend",
    handler: async (_args, ctx) => {
      await showUsagePopup(ctx);
    }
  });

  pi.registerCommand("usage-reset-session", {
    description: "Reset only the current PreApeXis session usage counter",
    handler: async (_args, ctx) => {
      resetBucket(sessionUsage);
      updateUsageFooter(ctx);
      ctx.ui.notify("PreApeXis session usage reset.", "info");
    }
  });

  pi.registerCommand("welcome", {
    description: "Show the PreApeXis welcome banner",
    handler: async (_args, ctx) => {
      await showWelcome(ctx);
      updateUsageFooter(ctx);
    }
  });
}