import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "@sinclair/typebox";
import { existsSync, statSync } from "fs";
import { mkdir, writeFile } from "fs/promises";
import { resolve, dirname, join } from "path";
import { getServiceKey } from "../shared/auth";
import { errorMessage, isAbortError } from "../shared/errors";
import { parsePdfWithMinerU, parseLocalFileWithMinerU } from "../clients/mineru";
import { renderContentListPages } from "../clients/content-pages";
import { shortenUrl } from "../shared/render";
import { type DocumentParseDetails } from "../types";
import { buildExpandedParams, kylinFrameForCall, kylinFrameForResult } from "../../../vera-theme/src/public";

function truncate(text: string, max = 70): string {
  return text.length > max ? text.slice(0, max - 1) + "…" : text;
}

function buildParamSummary(args: any, theme: any): string {
  const bits: string[] = [];
  if (args.paginate) bits.push("paginate");
  if (args.page_ranges) bits.push(`pages ${args.page_ranges}`);
  if (args.model_version) bits.push(String(args.model_version));
  return (
    theme.fg("text", truncate(shortenUrl(String(args.url ?? args.file_path ?? "")), 58)) +
    (bits.length > 0 ? theme.fg("borderMuted", " · ") + theme.fg("dim", bits.join(" · ")) : "")
  );
}

function contentLines(content: any): string[] {
  if (content?.type !== "text" || typeof content.text !== "string") return [];
  return content.text.split("\n");
}

export function registerDocumentParseTool(pi: ExtensionAPI): void {
  pi.registerTool({
    name: "document_parse",
    label: "Document Parse",
    description:
      "Parse a document (remote URL or local file path) into Markdown using MinerU and WRITE THE RESULT TO output_path — parsed content is never returned inline, so large documents do not flood the context. Provide exactly one of url or file_path, plus a required output_path. Set paginate=true to emit one Markdown file per PDF page (page-NNN.md) plus index.md into the output directory. Returns only a short summary.",
    promptSnippet:
      "Use document_parse to convert a document (url or file_path) into Markdown saved to output_path (content is written to disk, not returned). For large docs set paginate=true to get per-page files + index.md. The tool returns only a summary.",
    parameters: Type.Object({
      url: Type.Optional(Type.String({ description: "Remote URL to the document. Provide either url or file_path, not both." })),
      file_path: Type.Optional(Type.String({ description: "Local file path to the document (PDF, DOCX, PPT, Excel, or image). Provide either url or file_path, not both." })),
      output_path: Type.String({ description: "Required. Where to write the result. Without paginate: the Markdown file path. With paginate=true: a directory that receives per-page files (page-NNN.md) plus index.md." }),
      paginate: Type.Optional(Type.Boolean({ description: "When true, write one Markdown file per PDF page (page-NNN.md) plus index.md into output_path (a directory), reconstructed from MinerU's page structure. Default false (single Markdown file)." })),
      page_ranges: Type.Optional(Type.String({
        description: "Page range to parse (e.g. '1-10', '3,5-8'). Default: all pages.",
      })),
      model_version: Type.Optional(Type.String({
        description: "Parse model: 'vlm' (best quality, default), 'pipeline' (fast)",
        default: "vlm",
      })),
      enable_formula: Type.Optional(Type.Boolean({
        description: "Enable formula recognition (default: true)",
      })),
      enable_table: Type.Optional(Type.Boolean({
        description: "Enable table recognition (default: true)",
      })),
      language: Type.Optional(Type.String({
        description: "Document language hint, e.g. 'en', 'ch' (default: auto-detect)",
      })),
    }),
    renderShell: "self" as const,
    renderCall(args, theme, ctx) {
      const state = ctx.state as { startedAt?: number };
      if (state.startedAt === undefined) state.startedAt = Date.now();
      return kylinFrameForCall(ctx, {
        name: "document_parse",
        theme,
        status: "pending",
        paramSummary: buildParamSummary(args, theme),
        startedAt: state.startedAt,
      });
    },
    renderResult(result, { expanded, isPartial }, theme, context) {
      const details = result.details as DocumentParseDetails | undefined;
      const state = context.state as { startedAt?: number };
      const startedAt = state.startedAt;
      const paramSummary = buildParamSummary(context.args, theme);
      const expandedParams = expanded ? buildExpandedParams([
        { label: "url", value: context.args?.url },
        { label: "file_path", value: context.args?.file_path },
        { label: "output_path", value: context.args?.output_path },
        { label: "paginate", value: context.args?.paginate },
        { label: "page_ranges", value: context.args?.page_ranges },
        { label: "model_version", value: context.args?.model_version },
        { label: "enable_formula", value: context.args?.enable_formula },
        { label: "enable_table", value: context.args?.enable_table },
        { label: "language", value: context.args?.language },
      ], theme) : undefined;

      if (isPartial) {
        const phase = details?.phase ?? "submitting";
        const phaseLabel: Record<string, string> = {
          submitting: "Submitting",
          uploading: "Uploading",
          parsing: "Parsing",
          downloading: "Downloading",
        };
        const label = phaseLabel[phase] ?? phase;
        return kylinFrameForResult(context, {
          name: "document_parse",
          theme,
          status: "pending",
          paramSummary,
          resultSummary: theme.fg("warning", label) + (details?.totalPages ? theme.fg("dim", ` · ${details.totalPages} pages`) : ""),
          expanded,
          expandedParams,
          startedAt,
        });
      }

      if (details?.error) {
        return kylinFrameForResult(context, {
          name: "document_parse",
          theme,
          status: "error",
          paramSummary,
          errorMessage: details.error + (details.taskId ? ` · task ${details.taskId}` : ""),
          expanded,
          expandedParams,
          startedAt,
        });
      }
      if (!details) {
        return kylinFrameForResult(context, {
          name: "document_parse",
          theme,
          status: "error",
          paramSummary,
          errorMessage: "No result",
          expanded,
          expandedParams,
          startedAt,
        });
      }

      const lines = contentLines(result.content[0]);
      const written = details.paginated
        ? `${details.pages ?? 0} pages → ${details.outputPath ?? ""}`
        : `${details.lines ?? 0} lines → ${details.outputPath ?? ""}`;
      return kylinFrameForResult(context, {
        name: "document_parse",
        theme,
        status: "success",
        paramSummary,
        resultSummary: theme.fg("success", "written") + theme.fg("dim", ` · ${written}`),
        expanded,
        expandedParams,
        expandedBody: expanded ? lines.map((line) => theme.fg("dim", line)) : undefined,
        expandedTotalLines: lines.length,
        startedAt,
      });
    },
    async execute(_toolCallId, params, signal, onUpdate) {
      const apiKey = getServiceKey("mineru", "MINERU_API_KEY");
      const url = typeof params.url === "string" ? params.url.trim() : "";
      const filePath = typeof params.file_path === "string" ? params.file_path.trim() : "";
      const outputPath = typeof params.output_path === "string" ? params.output_path.trim() : "";
      const paginate = params.paginate === true;
      const source = url || filePath;

      if ((!url && !filePath) || (url && filePath)) {
        return {
          content: [{ type: "text" as const, text: "Error: provide exactly one of 'url' or 'file_path'." }],
          details: { source, phase: "done", error: "Provide exactly one of url or file_path" } as DocumentParseDetails,
        };
      }

      if (!outputPath) {
        return {
          content: [{ type: "text" as const, text: "Error: 'output_path' is required — document_parse writes results to disk and does not return content inline." }],
          details: { source, phase: "done", error: "output_path is required" } as DocumentParseDetails,
        };
      }

      let absPath = "";
      if (filePath) {
        absPath = resolve(filePath);
        if (!existsSync(absPath) || !statSync(absPath).isFile()) {
          return {
            content: [{ type: "text" as const, text: `Error: file not found: ${absPath}` }],
            details: { source, phase: "done", error: "File not found" } as DocumentParseDetails,
          };
        }
        const sizeMb = statSync(absPath).size / (1024 * 1024);
        if (sizeMb > 200) {
          return {
            content: [{ type: "text" as const, text: `Error: file exceeds MinerU's 200MB limit (${sizeMb.toFixed(1)}MB).` }],
            details: { source, phase: "done", error: "File too large" } as DocumentParseDetails,
          };
        }
      }

      if (!apiKey) {
        return {
          content: [{ type: "text" as const, text: "Error: No MinerU API key found. Set MINERU_API_KEY or add it to agent/auth.json under 'mineru'." }],
          details: { source, phase: "done", error: "No API key configured" } as DocumentParseDetails,
        };
      }

      onUpdate?.({
        content: [{ type: "text" as const, text: "Submitting to MinerU..." }],
        details: { source, phase: "submitting" } as DocumentParseDetails,
      });

      const options = {
        modelVersion: (params.model_version as "vlm" | "pipeline" | "MinerU-HTML") ?? "vlm",
        pageRanges: params.page_ranges,
        enableFormula: params.enable_formula,
        enableTable: params.enable_table,
        language: params.language,
      };

      const onProgress = (phase: string, detail?: string) => {
        const details: DocumentParseDetails = { source, phase: phase as DocumentParseDetails["phase"] };
        if (detail) details.taskId = detail;
        onUpdate?.({
          content: [{ type: "text" as const, text: `${phase}...` }],
          details,
        });
      };

      try {
        const result = url
          ? await parsePdfWithMinerU(url, apiKey, signal, onProgress, options, paginate)
          : await parseLocalFileWithMinerU(absPath, apiKey, signal, onProgress, options, paginate);

        if (result.error) {
          return {
            content: [{ type: "text" as const, text: `Parse error: ${result.error}` }],
            details: { source, phase: "done", taskId: result.taskId, error: result.error } as DocumentParseDetails,
          };
        }

        if (paginate) {
          const v2 = result.contentListV2;
          if (!Array.isArray(v2) || v2.length === 0) {
            return {
              content: [{ type: "text" as const, text: "Parse error: no page content returned for pagination." }],
              details: { source, phase: "done", taskId: result.taskId, error: "No page content" } as DocumentParseDetails,
            };
          }
          const rendered = renderContentListPages(v2);
          const dir = resolve(outputPath);
          await mkdir(dir, { recursive: true });
          let totalBytes = 0;
          for (const page of rendered.pages) {
            totalBytes += Buffer.byteLength(page.markdown, "utf8");
            await writeFile(join(dir, page.name), page.markdown, "utf8");
          }
          await writeFile(join(dir, "index.md"), rendered.index, "utf8");

          const summary = `Parsed ${rendered.pageCount} pages → wrote ${rendered.pageCount} page files + index.md (${totalBytes} bytes) to ${dir}`;
          return {
            content: [{ type: "text" as const, text: summary }],
            details: { source, phase: "done", taskId: result.taskId, totalPages: rendered.pageCount, pages: rendered.pageCount, bytes: totalBytes, outputPath: dir, paginated: true } as DocumentParseDetails,
          };
        }

        // Single-file mode: write full.md to output_path, return a summary + short preview.
        const absOut = resolve(outputPath);
        const bytes = Buffer.byteLength(result.markdown, "utf8");
        const lines = result.markdown.split("\n").length;
        await mkdir(dirname(absOut), { recursive: true });
        await writeFile(absOut, result.markdown, "utf8");

        const head = result.markdown.split("\n").slice(0, 8).join("\n");
        const pagesNote = result.totalPages ? ` (${result.totalPages} pages)` : "";
        const summary = `Parsed${pagesNote} → wrote ${lines} lines, ${bytes} bytes to ${absOut}\n\n--- preview (first 8 lines) ---\n${head}`;
        return {
          content: [{ type: "text" as const, text: summary }],
          details: { source, phase: "done", taskId: result.taskId, totalPages: result.totalPages, lines, bytes, outputPath: absOut, paginated: false } as DocumentParseDetails,
        };
      } catch (error) {
        if (isAbortError(error)) {
          return {
            content: [{ type: "text" as const, text: "Parse cancelled." }],
            details: { source, phase: "done", error: "Cancelled" } as DocumentParseDetails,
          };
        }
        return {
          content: [{ type: "text" as const, text: `Parse error: ${errorMessage(error)}` }],
          details: { source, phase: "done", error: errorMessage(error) } as DocumentParseDetails,
        };
      }
    },
  });
}