/**
 * Auto-context files extension for pi.
 *
 * When the agent reads any file, this extension walks up the directory tree
 * from that file's location looking for AGENTS.md or CLAUDE.md files, and
 * prepends them to that read tool result, so the model sees the local
 * instructions before the file contents it just asked to read.
 *
 * Walking stops at the git repository root.
 *
 * Built without module-level mutable state; all discovery and injection
 * bookkeeping is stored in durable session entries so it survives across
 * turns, reloads, and extension runtime resets.
 */

import * as fs from "node:fs";
import * as path from "node:path";
import { execSync } from "node:child_process";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

/** Context file names to auto-discover, in search order. */
const CONTEXT_FILE_NAMES = ["CLAUDE.md", "AGENTS.md"];

/**
 * Memoized git-root lookup per directory.
 */
function getGitRoot(dir: string): string | null {
  try {
    const root = execSync("git rev-parse --show-toplevel", {
      cwd: dir,
      encoding: "utf-8",
      timeout: 5000,
    }).trim();
    return root || null;
  } catch {
    return null;
  }
}

/**
 * Walk up from `startDir` looking for context files, stopping at `stopDir`.
 */
function discoverContextFiles(
  startDir: string,
  stopDir: string,
): Array<{ absPath: string; content: string }> {
  const found: Array<{ absPath: string; content: string }> = [];
  let current = startDir;

  while (
    current &&
    current !== "/" &&
    current.length >= stopDir.length &&
    current !== path.dirname(current)
  ) {
    for (const name of CONTEXT_FILE_NAMES) {
      const absPath = path.join(current, name);
      if (fs.existsSync(absPath) && fs.statSync(absPath).isFile()) {
        try {
          const content = fs.readFileSync(absPath, "utf-8");
          found.unshift({ absPath, content }); // closer file wins → later in prompt
        } catch {
          // ignore read errors
        }
      }
    }

    if (current === stopDir) break;
    const parent = path.dirname(current);
    if (parent === current) break;
    current = parent;
  }

  return found;
}

/** Get the directory that contains a file path. */
function dirOf(filePath: string): string {
  try {
    const stat = fs.statSync(filePath);
    if (stat.isDirectory()) return filePath;
  } catch {
    // fall through
  }
  return path.dirname(filePath);
}

/** Check if the file being read is itself a context file. */
function isContextFile(filePath: string): boolean {
  return CONTEXT_FILE_NAMES.includes(path.basename(filePath));
}

/** Session entry custom types used by this extension. */
const ENTRY_INIT = "auto-context:init";
const ENTRY_DISCOVERED = "auto-context:discovered";
const ENTRY_INJECTED = "auto-context:injected";

interface DiscoveredEntry {
  absPath: string;
  content: string;
  directory: string;
  toolCallId: string;
}

interface InjectedEntry {
  absPath: string;
  toolCallId: string;
  timestamp: number;
}

function isDiscoveredEntry(data: unknown): data is DiscoveredEntry {
  return (
    typeof data === "object" &&
    data !== null &&
    "absPath" in data &&
    typeof (data as any).absPath === "string" &&
    "content" in data &&
    typeof (data as any).content === "string"
  );
}

function isInjectedEntry(data: unknown): data is InjectedEntry {
  return (
    typeof data === "object" &&
    data !== null &&
    "absPath" in data &&
    typeof (data as any).absPath === "string"
  );
}

export default function autoContextFilesExtension(pi: ExtensionAPI) {
  /* ------------------------------------------------------------------ */
  /* 1. Session start — announce presence                               */
  /* ------------------------------------------------------------------ */
  pi.on("session_start", async (_event, ctx) => {
    if (ctx.hasUI) {
      ctx.ui.notify("🔄 Auto-context extension ready", "info");
    }

    pi.appendEntry(ENTRY_INIT, {
      cwd: ctx.cwd,
      gitRoot: getGitRoot(ctx.cwd),
      timestamp: Date.now(),
    });
  });

  /* ------------------------------------------------------------------ */
  /* 2. Tool call — intercept reads, walk up, discover context files   */
  /* ------------------------------------------------------------------ */
  pi.on("tool_call", async (event, ctx) => {
    if (event.toolName !== "read") return;

    const filePath = event.input?.path;
    if (!filePath || typeof filePath !== "string") return;

    // Don't discover when reading a context file itself.
    if (isContextFile(filePath)) return;

    const startDir = dirOf(path.resolve(ctx.cwd, filePath));
    const gitRoot = getGitRoot(startDir);
    const stopDir = gitRoot ?? ctx.cwd;

    // Safety: never walk above git root or cwd. Use path.relative instead of a
    // string prefix check so /repo/foo2 is not treated as inside /repo/foo.
    const relativeToStop = path.relative(stopDir, startDir);
    if (relativeToStop.startsWith("..") || path.isAbsolute(relativeToStop)) return;

    const discovered = discoverContextFiles(startDir, stopDir);
    if (discovered.length === 0) return;

    // Re-read already-injected set from session so we don't duplicate.
    const injectedSet = new Set<string>();
    for (const entry of ctx.sessionManager.getEntries()) {
      if (entry.type === "custom" && entry.customType === ENTRY_INJECTED) {
        if (isInjectedEntry(entry.data)) {
          injectedSet.add(entry.data.absPath);
        }
      }
    }

    for (const { absPath, content } of discovered) {
      if (injectedSet.has(absPath)) continue;

      pi.appendEntry(ENTRY_DISCOVERED, {
        absPath,
        content,
        directory: path.dirname(absPath),
        toolCallId: event.toolCallId,
      } as DiscoveredEntry);

      if (ctx.hasUI) {
        ctx.ui.notify(
          `📄 Discovered ${path.relative(ctx.cwd, absPath)}`,
          "info",
        );
      }
    }
  });

  /* ------------------------------------------------------------------ */
  /* 3. Tool result — prepend context before the file contents          */
  /* ------------------------------------------------------------------ */
  pi.on("tool_result", async (event, ctx) => {
    if (event.toolName !== "read" || event.isError) return;

    const toInject: DiscoveredEntry[] = [];
    for (const entry of ctx.sessionManager.getEntries()) {
      if (entry.type !== "custom" || entry.customType !== ENTRY_DISCOVERED) continue;
      if (isDiscoveredEntry(entry.data) && entry.data.toolCallId === event.toolCallId) {
        toInject.push(entry.data);
      }
    }

    if (toInject.length === 0) return;

    const now = Date.now();
    for (const { absPath } of toInject) {
      pi.appendEntry(ENTRY_INJECTED, {
        absPath,
        toolCallId: event.toolCallId,
        timestamp: now,
      } as InjectedEntry);
    }

    const contextText = toInject
      .map(({ absPath, content }) => {
        const relativePath = path.relative(ctx.cwd, absPath);
        return `<!-- auto-context: ${relativePath} -->\n${content}`;
      })
      .join("\n\n");

    if (ctx.hasUI) {
      ctx.ui.notify(
        `📌 Prepended ${toInject.length} context file(s) to read result`,
        "info",
      );
    }

    return {
      content: [
        {
          type: "text",
          text: `Auto-loaded context for this file:\n\n${contextText}`,
        },
        ...event.content,
      ],
    };
  });

  /* ------------------------------------------------------------------ */
  /* 4. Diagnostic command                                              */
  /* ------------------------------------------------------------------ */
  pi.registerCommand("auto-context", {
    description: "Show auto-context extension status",
    handler: async (_args, ctx) => {
      const entries = ctx.sessionManager.getEntries();
      const discovered: DiscoveredEntry[] = [];
      const injected: InjectedEntry[] = [];

      for (const entry of entries) {
        if (entry.type !== "custom") continue;
        if (entry.customType === ENTRY_DISCOVERED && isDiscoveredEntry(entry.data)) {
          discovered.push(entry.data);
        } else if (entry.customType === ENTRY_INJECTED && isInjectedEntry(entry.data)) {
          injected.push(entry.data);
        }
      }

      const injectedSet = new Set(injected.map((i) => i.absPath));
      const pending = discovered.filter((d) => !injectedSet.has(d.absPath));

      let msg = "📋 Auto-context status:\n";
      msg += `  Discovered: ${discovered.length}\n`;
      msg += `  Injected:   ${injected.length}\n`;
      msg += `  Pending:    ${pending.length}\n`;
      for (const d of pending) {
        msg += `    ⏳ ${path.relative(ctx.cwd, d.absPath)}\n`;
      }
      for (const i of injected) {
        msg += `    ✅ ${path.relative(ctx.cwd, i.absPath)}\n`;
      }

      ctx.ui.notify(msg, "info");
    },
  });
}