/**
 * session-files extension — Track and review files edited/created by the AI in
 * the current session.
 *
 * Provides:
 *  - Automatic tracking of every file touched by the edit and write tools.
 *  - Per-file diff statistics (lines added/removed, edit count).
 *  - /session-files TUI: directory-grouped tree view with keyboard shortcuts
 *    to open files or folders in VS Code (`code`).
 */

import type { ExtensionAPI, Theme } from "@earendil-works/pi-coding-agent";
import { isEditToolResult, isWriteToolResult } from "@earendil-works/pi-coding-agent";
import { Key, matchesKey, truncateToWidth } from "@earendil-works/pi-tui";
import { exec } from "node:child_process";
import { basename, dirname, isAbsolute, relative, resolve } from "node:path";

// ---------------------------------------------------------------------------
// Data types
// ---------------------------------------------------------------------------

interface EditRecord {
  additions: number;
  deletions: number;
}

interface FileInfo {
  absPath: string;
  relPath: string;
  dir: string; // parent directory (relative), displayed as group header
  records: EditRecord[];
  totalAdditions: number;
  totalDeletions: number;
  /** True when the write tool was used (indicates file creation / overwrite). */
  wasCreated: boolean;
}

type FlatItem =
  | { kind: "dir"; label: string; fileCount: number }
  | { kind: "file"; info: FileInfo };

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

function parseDiffStats(diff: string): { additions: number; deletions: number } {
  let additions = 0;
  let deletions = 0;
  for (const line of diff.split("\n")) {
    if (line.startsWith("+") && !line.startsWith("+++")) additions++;
    else if (line.startsWith("-") && !line.startsWith("---")) deletions++;
  }
  return { additions, deletions };
}

function toAbs(cwd: string, p: string): string {
  return isAbsolute(p) ? p : resolve(cwd, p);
}

function toRel(cwd: string, p: string): string {
  const abs = toAbs(cwd, p);
  const rel = relative(cwd, abs);
  return rel || basename(abs);
}

function padRight(s: string, n: number): string {
  return s.length >= n ? s : s + " ".repeat(n - s.length);
}

function openInCode(target: string): void {
  exec(`code "${target}"`, (err) => {
    if (err) {
      // Silently ignore — VS Code may not be installed.
    }
  });
}

function openTerminal(cwd: string): void {
  const plat = process.platform;
  let cmd: string;

  if (plat === "win32") {
    // Windows: prefer Windows Terminal, fallback to cmd
    cmd = `wt -d "${cwd}" 2>nul || start "" cmd /K "cd /d ${cwd}"`;
  } else if (plat === "darwin") {
    cmd = `open -a Terminal "${cwd}"`;
  } else {
    // Linux: try common terminal emulators, fallback to xterm
    cmd = `cd "${cwd}" && (gnome-terminal --working-directory="${cwd}" 2>/dev/null || konsole --workdir "${cwd}" 2>/dev/null || xfce4-terminal --working-directory="${cwd}" 2>/dev/null || x-terminal-emulator 2>/dev/null || xterm) &`;
  }

  exec(cmd, (err) => {
    if (err) {
      // Silently ignore
    }
  });
}

// ---------------------------------------------------------------------------
// TUI component
// ---------------------------------------------------------------------------

class SessionFilesComponent {
  private allDirs: Array<{ label: string; files: FileInfo[] }>;
  private expandedDirs: Set<string>;
  private filterQuery = "";
  /** 'file' = only matching files shown; 'folder' = dir match → show all files */
  private filterMode: "file" | "folder" = "file";

  private items: FlatItem[];
  private selectedIndex: number;
  private theme: Theme;
  private onClose: () => void;
  private cwd: string;

  // Virtual scrolling
  private maxVisible = 15;
  private scrollOffset = 0;

  // Render cache
  private cachedWidth?: number;
  private cachedLines?: string[];

  constructor(
    fileMap: Map<string, FileInfo>,
    theme: Theme,
    cwd: string,
    onClose: () => void,
  ) {
    this.theme = theme;
    this.cwd = cwd;
    this.onClose = onClose;

    // Build sorted dir list
    const byDir = new Map<string, FileInfo[]>();
    for (const info of fileMap.values()) {
      const d = info.dir;
      if (!byDir.has(d)) byDir.set(d, []);
      byDir.get(d)!.push(info);
    }
    this.allDirs = [...byDir.entries()]
      .sort(([a], [b]) => a.localeCompare(b))
      .map(([label, files]) => ({
        label,
        files: files.sort((a, b) =>
          basename(a.relPath).localeCompare(basename(b.relPath)),
        ),
      }));

    this.expandedDirs = new Set();
    this.rebuildItems();
    this.selectedIndex = this.items.findIndex((i) => i.kind === "file");
    if (this.selectedIndex < 0) this.selectedIndex = 0;
    this.clampScroll();
  }

  // ---- Component interface ----

  render(width: number): string[] {
    if (this.cachedLines && this.cachedWidth === width) {
      return this.cachedLines;
    }
    this.cachedWidth = width;
    this.cachedLines = this.doRender(width);
    return this.cachedLines;
  }

  invalidate(): void {
    this.cachedWidth = undefined;
    this.cachedLines = undefined;
  }

  handleInput(data: string): void {
    // ---- Navigation ----
    if (matchesKey(data, Key.up)) {
      this.selectedIndex = Math.max(0, this.selectedIndex - 1);
      this.clampScroll();
      this.invalidate();
      return;
    }
    if (matchesKey(data, Key.down)) {
      this.selectedIndex = Math.min(
        this.items.length - 1,
        this.selectedIndex + 1,
      );
      this.clampScroll();
      this.invalidate();
      return;
    }
    if (matchesKey(data, Key.pageUp)) {
      this.selectedIndex = Math.max(0, this.selectedIndex - this.maxVisible);
      this.clampScroll();
      this.invalidate();
      return;
    }
    if (matchesKey(data, Key.pageDown)) {
      this.selectedIndex = Math.min(
        this.items.length - 1,
        this.selectedIndex + this.maxVisible,
      );
      this.clampScroll();
      this.invalidate();
      return;
    }

    // ---- Toggle / expand / collapse ----
    if (matchesKey(data, Key.enter) || matchesKey(data, Key.space)) {
      this.activateOrToggle();
      return;
    }
    if (matchesKey(data, Key.right)) {
      this.expandCurrentDir();
      return;
    }
    if (matchesKey(data, Key.left)) {
      this.collapseCurrentDir();
      return;
    }

    // ---- Actions ----
    if (matchesKey(data, Key.ctrl("o"))) {
      this.activateSelected("folder");
      return;
    }
    if (matchesKey(data, Key.ctrl("t"))) {
      this.activateSelected("terminal");
      return;
    }

    // ---- Tab: toggle filter mode ----
    if (matchesKey(data, Key.tab)) {
      this.filterMode = this.filterMode === "file" ? "folder" : "file";
      this.rebuildItems();
      this.selectedIndex = 0;
      this.scrollOffset = 0;
      this.invalidate();
      return;
    }

    // ---- Escape: clear filter first, then close ----
    if (matchesKey(data, Key.escape)) {
      if (this.filterQuery) {
        this.filterQuery = "";
        this.rebuildItems();
        this.selectedIndex = 0;
        this.scrollOffset = 0;
        this.invalidate();
      } else {
        this.onClose();
      }
      return;
    }

    // ---- Filter editing ----
    if (matchesKey(data, Key.backspace) || matchesKey(data, Key.delete)) {
      if (this.filterQuery.length > 0) {
        this.filterQuery = this.filterQuery.slice(0, -1);
        this.rebuildItems();
        this.selectedIndex = 0;
        this.scrollOffset = 0;
        this.invalidate();
      }
      return;
    }
    if (matchesKey(data, Key.ctrl("u"))) {
      this.filterQuery = "";
      this.rebuildItems();
      this.selectedIndex = 0;
      this.scrollOffset = 0;
      this.invalidate();
      return;
    }

    // ---- Type to filter (printable chars & multi-byte for CJK) ----
    if (data.length === 1) {
      const c = data.charCodeAt(0);
      if (c >= 32 && c !== 127) {
        this.filterQuery += data;
        this.rebuildItems();
        this.selectedIndex = 0;
        this.scrollOffset = 0;
        this.invalidate();
        return;
      }
    }
    if (data.length > 1 && !data.startsWith("\x1b")) {
      this.filterQuery += data;
      this.rebuildItems();
      this.selectedIndex = 0;
      this.scrollOffset = 0;
      this.invalidate();
    }
  }

  // ---- Directory expand / collapse ----

  private expandCurrentDir(): void {
    const item = this.items[this.selectedIndex];
    if (!item || item.kind !== "dir") return;
    if (!this.expandedDirs.has(item.label)) {
      this.expandedDirs.add(item.label);
      this.rebuildItems();
      this.clampScroll();
      this.invalidate();
    }
  }

  private collapseCurrentDir(): void {
    const item = this.items[this.selectedIndex];
    if (!item) return;
    const label = item.kind === "dir"
      ? item.label
      : item.info.dir;
    if (this.expandedDirs.has(label)) {
      this.expandedDirs.delete(label);
      this.rebuildItems();
      this.clampScroll();
      this.invalidate();
    }
  }

  private activateOrToggle(): void {
    const item = this.items[this.selectedIndex];
    if (!item) return;

    if (item.kind === "dir") {
      // Toggle expand/collapse
      if (this.expandedDirs.has(item.label)) {
        this.expandedDirs.delete(item.label);
      } else {
        this.expandedDirs.add(item.label);
      }
      this.rebuildItems();
      this.clampScroll();
      this.invalidate();
    } else {
      // File: open in code
      this.activateSelected("file");
    }
  }

  // ---- Rebuild flat item list ----

  private rebuildItems(): void {
    const q = this.filterQuery.toLowerCase().trim();
    const isFolderMode = this.filterMode === "folder";
    const items: FlatItem[] = [];

    for (const dir of this.allDirs) {
      const dirName = dir.label.toLowerCase();
      const dirMatches = !q || dirName.includes(q);

      // Files that match the query (by filename)
      const matchingFiles = q
        ? dir.files.filter((f) => basename(f.relPath).toLowerCase().includes(q))
        : [];

      // Should the directory header be visible?
      // In both modes: show if dir name matches, or any file matches
      const hasVisibleFiles =
        (isFolderMode && dirMatches) ||  // folder mode + dir match → all files
        (!q) ||                           // no query → all files (when expanded)
        matchingFiles.length > 0;         // file mode + matching files

      const showDir = !q || dirMatches || hasVisibleFiles;
      if (!showDir) continue;

      items.push({ kind: "dir", label: dir.label, fileCount: dir.files.length });

      // Show files if expanded, OR if something would be visible due to filter
      const showFiles =
        this.expandedDirs.has(dir.label) ||
        (q && hasVisibleFiles);

      if (showFiles) {
        // Decide which files to show
        let filesToShow: FileInfo[];
        if (!q) {
          filesToShow = dir.files;
        } else if (isFolderMode && dirMatches) {
          // Folder mode: dir name matched → show ALL files in this dir
          filesToShow = dir.files;
        } else {
          // File mode: only show files whose name matches
          filesToShow = matchingFiles;
        }
        for (const f of filesToShow) {
          items.push({ kind: "file", info: f });
        }
      }
    }

    this.items = items;
  }

  // ---- Virtual scroll ----

  private clampScroll(): void {
    if (this.selectedIndex < this.scrollOffset) {
      this.scrollOffset = this.selectedIndex;
    } else if (this.selectedIndex >= this.scrollOffset + this.maxVisible) {
      this.scrollOffset = this.selectedIndex - this.maxVisible + 1;
    }
    this.scrollOffset = Math.max(
      0,
      Math.min(this.scrollOffset, Math.max(0, this.items.length - this.maxVisible)),
    );
  }

  // ---- Render ----

  private doRender(width: number): string[] {
    const t = this.theme;
    const totalFiles = this.allDirs.reduce((s, d) => s + d.files.length, 0);
    const totalEdits = this.allDirs.reduce(
      (s, d) => s + d.files.reduce((ss, f) => ss + f.records.length, 0),
      0,
    );

    // Empty state
    if (this.allDirs.length === 0) {
      return [t.fg("dim", " No files have been edited or created in this session.")];
    }

    const lines: string[] = [];

    // Header
    lines.push(
      t.fg("accent", t.bold(` Session Files (${totalFiles} files, ${totalEdits} edits)`)),
    );

    // Filter bar
    const matchFiles = this.items.filter((i) => i.kind === "file").length;
    const modeLabel = this.filterMode === "file" ? "[file]" : "[folder]";
    if (this.filterQuery) {
      lines.push(
        t.fg("accent", ` ${this.filterQuery} `) +
          t.fg("dim", `${matchFiles} file${matchFiles !== 1 ? "s" : ""}  `) +
          t.fg("muted", modeLabel),
      );
    } else {
      lines.push(
        t.fg("dim", ` type to filter  `) +
          t.fg("muted", `Tab ${modeLabel}`),
      );
    }

    lines.push(t.fg("dim", " " + "─".repeat(Math.min(width - 2, 50))));

    const canScrollUp = this.scrollOffset > 0;
    const canScrollDown = this.scrollOffset + this.maxVisible < this.items.length;

    // Up overflow
    if (canScrollUp) {
      lines.push(t.fg("dim", ` ↑ ${this.scrollOffset} more items above`));
    } else {
      lines.push("");
    }

    // No items after filter
    if (this.items.length === 0) {
      lines.push(t.fg("dim", " No matching files."));
    }

    // Visible window
    const end = Math.min(this.scrollOffset + this.maxVisible, this.items.length);
    for (let i = this.scrollOffset; i < end; i++) {
      const item = this.items[i]!;
      const selected = i === this.selectedIndex;

      if (item.kind === "dir") {
        this.renderDirLine(lines, item, selected, width);
      } else {
        this.renderFileLine(lines, item.info, selected, width);
      }
    }

    // Down overflow
    if (canScrollDown) {
      lines.push(
        t.fg("dim", ` ↓ ${this.items.length - this.scrollOffset - this.maxVisible} more items below`),
      );
    } else {
      lines.push("");
    }

    // Footer
    lines.push(t.fg("dim", " " + "─".repeat(Math.min(width - 2, 50))));
    lines.push(
      t.fg(
        "dim",
        " ↑↓ PgUp/Dn  ←→ fold  Enter=toggle  Tab=mode  Ctrl+O=folder  Ctrl+T=term  Esc=close",
      ),
    );

    return lines;
  }

  private renderDirLine(
    lines: string[],
    item: FlatItem & { kind: "dir" },
    selected: boolean,
    width: number,
  ): void {
    const t = this.theme;
    const expanded = this.expandedDirs.has(item.label);
    const arrow = expanded ? "▼" : "▶";
    const prefix = selected
      ? t.fg("accent", arrow)
      : t.fg("dim", arrow);
    const icon = t.fg("accent", "📁");
    const label = selected
      ? t.fg("accent", t.bold(` ${item.label}/`))
      : t.fg("accent", ` ${item.label}/`);
    const count = t.fg("dim", ` (${item.fileCount} file${item.fileCount !== 1 ? "s" : ""})`);
    lines.push(truncateToWidth(` ${prefix} ${icon}${label}${count}`, width));
  }

  private renderFileLine(
    lines: string[],
    info: FileInfo,
    selected: boolean,
    width: number,
  ): void {
    const t = this.theme;
    const prefix = selected ? t.fg("accent", "▶") : " ";
    const icon = info.wasCreated
      ? t.fg("success", "●")
      : t.fg("dim", "○");
    const name = selected
      ? t.fg("accent", t.bold(` ${basename(info.relPath)}`))
      : t.fg("muted", ` ${basename(info.relPath)}`);

    const addStr = info.totalAdditions > 0
      ? t.fg("success", `+${info.totalAdditions}`)
      : " 0";
    const delStr = info.totalDeletions > 0
      ? t.fg("error", `-${info.totalDeletions}`)
      : " 0";
    const statStr = `${padRight(addStr, 6)} ${padRight(delStr, 6)}`;
    const editStr = t.fg("dim", `${info.records.length} edit${info.records.length > 1 ? "s" : ""}`);

    const leftPart = `   ${prefix} ${icon}${name}`;
    const rightPart = ` ${statStr} ${editStr}`;

    const leftVisible = leftPart.replace(/\x1b\[[0-9;]*m/g, "").length;
    const rightVisible = rightPart.replace(/\x1b\[[0-9;]*m/g, "").length;
    const padding = Math.max(1, width - leftVisible - rightVisible);

    lines.push(
      truncateToWidth(leftPart + " ".repeat(padding) + rightPart, width),
    );
  }

  // ---- Actions ----

  private activateSelected(action: "file" | "folder" | "terminal"): void {
    const item = this.items[this.selectedIndex];
    if (!item) return;

    let targetDir: string;
    if (item.kind === "dir") {
      targetDir = toAbs(this.cwd, item.label);
    } else {
      targetDir = toAbs(this.cwd, item.info.dir);
    }

    if (action === "terminal") {
      openTerminal(targetDir);
      return;
    }

    if (item.kind === "dir") {
      openInCode(targetDir);
      return;
    }

    if (action === "file") {
      openInCode(item.info.absPath);
    } else {
      openInCode(targetDir);
    }
  }
}

// ---------------------------------------------------------------------------
// Persistence (via session JSONL custom entries)
// ---------------------------------------------------------------------------

const STATE_ENTRY_TYPE = "session-files-state";

/** Serializable snapshot stored in the session file. */
interface FileMapSnapshot {
  files: Array<{
    absPath: string;
    relPath: string;
    dir: string;
    totalAdditions: number;
    totalDeletions: number;
    wasCreated: boolean;
    recordCount: number;
  }>;
}

function snapshotFromMap(fileMap: Map<string, FileInfo>): FileMapSnapshot {
  const files: FileMapSnapshot["files"] = [];
  for (const info of fileMap.values()) {
    files.push({
      absPath: info.absPath,
      relPath: info.relPath,
      dir: info.dir,
      totalAdditions: info.totalAdditions,
      totalDeletions: info.totalDeletions,
      wasCreated: info.wasCreated,
      recordCount: info.records.length,
    });
  }
  return { files };
}

// ---------------------------------------------------------------------------
// Extension entry point
// ---------------------------------------------------------------------------

export default function (pi: ExtensionAPI) {
  // filePath (absolute) → FileInfo
  let fileMap = new Map<string, FileInfo>();

  function persistState(): void {
    pi.appendEntry(STATE_ENTRY_TYPE, snapshotFromMap(fileMap));
  }

  // ---- Restore state from session entries ----
  // Used for session_start (new / reload / resume / fork) and
  // session_tree (tree navigation within the same session file).
  function restoreFromSession(ctx: { sessionManager: { getEntries(): ReadonlyArray<{ id: string; parentId?: string | null; type: string; customType?: string; data?: unknown }>; getLeafId(): string } }): void {
    const allEntries = ctx.sessionManager.getEntries();
    const leafId = ctx.sessionManager.getLeafId();

    // Build the set of entry IDs on the current active path
    const activePathIds = new Set<string>();
    const entriesById = new Map<string, (typeof allEntries)[number]>();
    for (const e of allEntries) {
      entriesById.set(e.id, e);
    }
    let current: string | null | undefined = leafId;
    while (current) {
      activePathIds.add(current);
      const entry = entriesById.get(current);
      current = entry?.parentId;
    }

    // Only merge snapshots that are on the active path
    const restored = new Map<string, FileInfo>();
    for (const entry of allEntries) {
      if (
        entry.type === "custom" &&
        entry.customType === STATE_ENTRY_TYPE &&
        entry.data &&
        activePathIds.has(entry.id)
      ) {
        const snap = entry.data as FileMapSnapshot;
        if (snap.files) {
          for (const f of snap.files) {
            restored.set(f.absPath, {
              absPath: f.absPath,
              relPath: f.relPath,
              dir: f.dir,
              records: new Array(f.recordCount).fill({ additions: 0, deletions: 0 }),
              totalAdditions: f.totalAdditions,
              totalDeletions: f.totalDeletions,
              wasCreated: f.wasCreated,
            });
          }
        }
      }
    }

    fileMap = restored;
  }

  pi.on("session_start", (_event, ctx) => restoreFromSession(ctx));

  pi.on("session_tree", (_event, ctx) => restoreFromSession(ctx));

  // After compaction the active path changes — re-persist current state so
  // it stays on the new active branch.
  pi.on("session_compact", () => {
    if (fileMap.size > 0) persistState();
  });

  // Save final state on shutdown so nothing is lost
  pi.on("session_shutdown", () => {
    if (fileMap.size > 0) persistState();
  });

  function ensureFile(
    absPath: string,
    relPath: string,
    dir: string,
    isWrite: boolean,
  ): FileInfo {
    let info = fileMap.get(absPath);
    if (!info) {
      info = {
        absPath,
        relPath,
        dir,
        records: [],
        totalAdditions: 0,
        totalDeletions: 0,
        wasCreated: isWrite,
      };
      fileMap.set(absPath, info);
    } else if (isWrite) {
      info.wasCreated = true;
    }
    return info;
  }

  // ---- Track tool results ----
  pi.on("tool_result", async (event, ctx) => {
    if (event.isError) return;

    if (isEditToolResult(event)) {
      const rawPath = (event.input as Record<string, unknown>).path;
      if (typeof rawPath !== "string") return;

      const absPath = toAbs(ctx.cwd, rawPath);
      const relPath = toRel(ctx.cwd, rawPath);
      const dir = dirname(relPath) || ".";
      const diff = event.details?.diff;

      let additions = 0;
      let deletions = 0;
      if (diff) {
        const stats = parseDiffStats(diff);
        additions = stats.additions;
        deletions = stats.deletions;
      }

      const info = ensureFile(absPath, relPath, dir, false);
      info.records.push({ additions, deletions });
      info.totalAdditions += additions;
      info.totalDeletions += deletions;

      // Persist after every edit so reload / resume never lose data
      persistState();
      return;
    }

    if (isWriteToolResult(event)) {
      const rawPath = (event.input as Record<string, unknown>).path;
      if (typeof rawPath !== "string") return;

      const absPath = toAbs(ctx.cwd, rawPath);
      const relPath = toRel(ctx.cwd, rawPath);
      const dir = dirname(relPath) || ".";

      const info = ensureFile(absPath, relPath, dir, true);
      info.records.push({ additions: 0, deletions: 0 });

      persistState();
    }
  });

  // ---- /codedot command ----
  pi.registerCommand("codedot", {
    description: "Open the current pi working directory in VS Code",
    handler: async (_args, ctx) => {
      openInCode(ctx.cwd);
      ctx.ui.notify(`Opened ${ctx.cwd} in VS Code`, "info");
    },
  });

  // ---- /session-files command ----
  pi.registerCommand("session-files", {
    description: "Browse files edited by the AI this session",
    handler: async (_args, ctx) => {
      await ctx.ui.custom<void>((_tui, theme, _kb, done) => {
        const comp = new SessionFilesComponent(
          fileMap,
          theme,
          ctx.cwd,
          () => done(),
        );
        return {
          render: (w: number) => comp.render(w),
          invalidate: () => comp.invalidate(),
          handleInput: (data: string) => comp.handleInput(data),
        };
      });
    },
  });
}