import { mkdtemp, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join, resolve } from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import {
  DEFAULT_MAX_BYTES,
  DEFAULT_MAX_LINES,
  formatSize,
  truncateHead,
  withFileMutationQueue,
} from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";

const DEFAULT_TOP_K = 5;
const MAX_RENDERED_QUERY_LENGTH = 72;

const SearchParameters = Type.Object(
  {
    query: Type.String({
      minLength: 1,
      description: "Natural-language or text query for project documentation",
    }),
    topK: Type.Optional(
      Type.Integer({
        minimum: 1,
        maximum: 20,
        description: `Number of results (default: ${DEFAULT_TOP_K})`,
      }),
    ),
  },
  { additionalProperties: false },
);

interface SearchDetails {
  docsRoot: string;
  topK: number;
  truncated: boolean;
  fullOutputPath?: string;
  locations?: string[];
  matchCount?: number;
}

function compactQuery(query: string) {
  const singleLine = query.replace(/\s+/g, " ").trim();
  if (singleLine.length <= MAX_RENDERED_QUERY_LENGTH) return singleLine;
  return `${singleLine.slice(0, MAX_RENDERED_QUERY_LENGTH - 3)}...`;
}

function summarizeOutput(
  output: string,
): Pick<SearchDetails, "locations" | "matchCount"> {
  try {
    const parsed = JSON.parse(output) as unknown;
    if (!parsed || typeof parsed !== "object" || !("results" in parsed))
      return {};

    const results = (parsed as { results?: unknown }).results;
    if (!Array.isArray(results)) return {};

    const locations = results.flatMap((result) => {
      if (!result || typeof result !== "object") return [];
      const { file_path, start_line, end_line } = result as Record<
        string,
        unknown
      >;
      if (
        typeof file_path !== "string" ||
        typeof start_line !== "number" ||
        typeof end_line !== "number"
      )
        return [];
      return [`${file_path}:${start_line}-${end_line}`];
    });
    return { locations, matchCount: results.length };
  } catch {
    return {};
  }
}

async function prepareOutput(output: string, details: SearchDetails) {
  const truncation = truncateHead(output, {
    maxLines: DEFAULT_MAX_LINES,
    maxBytes: DEFAULT_MAX_BYTES,
  });

  if (!truncation.truncated) {
    return { text: truncation.content.trimEnd(), details };
  }

  const tempDirectory = await mkdtemp(join(tmpdir(), "pi-semble-docs-"));
  const fullOutputPath = join(tempDirectory, "results.json");
  await withFileMutationQueue(fullOutputPath, () =>
    writeFile(fullOutputPath, output, "utf8"),
  );

  details.truncated = true;
  details.fullOutputPath = fullOutputPath;

  const text = `${truncation.content}\n\n[Output truncated: showing ${truncation.outputLines} of ${truncation.totalLines} lines (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}). Full output saved to: ${fullOutputPath}]`;
  return { text, details };
}

export default function sembleDocsExtension(pi: ExtensionAPI) {
  pi.registerTool({
    name: "semble_docs_search",
    label: "Semble Docs Search",
    description: `Search documentation under docs/**/* with Semble. The search root and content type are fixed; output is truncated to ${DEFAULT_MAX_LINES} lines or ${formatSize(DEFAULT_MAX_BYTES)}.`,
    promptSnippet: "Semantically search documentation under docs/**/* only",
    promptGuidelines: [
      "Use semble_docs_search for conceptual or natural-language searches in project documentation; it cannot search outside docs/**/*.",
    ],
    parameters: SearchParameters,

    renderCall(args, theme) {
      const title = theme.fg("toolTitle", theme.bold("semble_docs_search "));
      return new Text(
        `${title}${theme.fg("accent", `"${compactQuery(args.query)}"`)}`,
        0,
        0,
      );
    },

    renderResult(result, { expanded, isPartial }, theme) {
      if (isPartial)
        return new Text(theme.fg("warning", "Searching docs..."), 0, 0);

      const details = result.details as SearchDetails | undefined;
      const count = details?.matchCount;
      const label =
        count === undefined
          ? "Search complete"
          : `${count} ${count === 1 ? "result" : "results"} in docs/`;
      let text = theme.fg("success", label);

      if (details?.truncated) text += theme.fg("warning", " (truncated)");
      if (expanded && details?.locations) {
        for (const location of details.locations)
          text += `\n  ${theme.fg("dim", location)}`;
      }
      if (expanded && details?.fullOutputPath) {
        text += `\n${theme.fg("dim", `Full output: ${details.fullOutputPath}`)}`;
      }

      return new Text(text, 0, 0);
    },

    async execute(_toolCallId, params, signal, _onUpdate, ctx) {
      const docsRoot = resolve(ctx.cwd, "docs");
      const topK = params.topK ?? DEFAULT_TOP_K;
      const result = await pi.exec(
        "semble",
        [
          "search",
          params.query,
          docsRoot,
          "--top-k",
          String(topK),
          "--content",
          "docs",
        ],
        { signal },
      );

      if (result.killed || signal?.aborted) {
        throw new Error("Semble search cancelled.");
      }
      if (result.code !== 0) {
        const reason =
          result.stderr.trim() ||
          result.stdout.trim() ||
          `exit code ${result.code}`;
        throw new Error(`Semble search failed: ${reason}`);
      }

      const details: SearchDetails = {
        docsRoot,
        topK,
        truncated: false,
        ...summarizeOutput(result.stdout),
      };
      const output = await prepareOutput(result.stdout, details);
      return {
        content: [{ type: "text", text: output.text }],
        details: output.details,
      };
    },
  });
}