import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "@sinclair/typebox";
import { getServiceKey } from "../shared/auth";
import { errorMessage, isAbortError } from "../shared/errors";
import { searchOpenAlex } from "../clients/openalex";
import { type AcademicSearchDetails } 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 filters: string[] = [];
  if (args.type) filters.push(String(args.type));
  if (args.year_from || args.year_to) filters.push(`${args.year_from ?? "?"}–${args.year_to ?? "?"}`);
  return (
    theme.fg("text", `"${truncate(String(args.query ?? ""), 58)}"`) +
    (filters.length > 0 ? theme.fg("borderMuted", " · ") + theme.fg("dim", filters.join(" · ")) : "")
  );
}

function extractTextPreview(content: any, maxLines: number): { lines: string[]; totalLines: number } {
  if (content?.type !== "text" || typeof content.text !== "string") return { lines: [], totalLines: 0 };
  const lines = content.text.split("\n").filter((line: string) => line.trim());
  return { lines: lines.slice(0, maxLines), totalLines: lines.length };
}

export function registerAcademicSearchTool(pi: ExtensionAPI): void {
  pi.registerTool({
    name: "academic_search",
    label: "Academic Search",
    description:
      "Search academic papers and books via OpenAlex (470M+ scholarly works). Returns titles, authors, citations, DOI, open access PDF links, and abstracts. Use this for finding research papers, literature reviews, or scholarly sources.",
    promptSnippet:
      "Use academic_search to find scholarly papers, books, and research articles. Supports filtering by type (article, book, etc.) and year range.",
    parameters: Type.Object({
      query: Type.String({ description: "The search query (e.g. 'transformer attention mechanism', 'reinforcement learning robotics')" }),
      type: Type.Optional(Type.String({
        description: "Work type filter: article, book, book-chapter, preprint, dissertation, conference-paper, dataset, review, or editorial",
      })),
      year_from: Type.Optional(Type.Number({ description: "Earliest publication year (e.g. 2020)" })),
      year_to: Type.Optional(Type.Number({ description: "Latest publication year (e.g. 2024)" })),
    }),
    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: "academic_search",
        theme,
        status: "pending",
        paramSummary: buildParamSummary(args, theme),
        startedAt: state.startedAt,
      });
    },
    renderResult(result, { expanded, isPartial }, theme, context) {
      const details = result.details as AcademicSearchDetails | undefined;
      const state = context.state as { startedAt?: number };
      const startedAt = state.startedAt;
      const paramSummary = buildParamSummary(context.args, theme);
      const expandedParams = expanded ? buildExpandedParams([
        { label: "query", value: context.args?.query, multiline: true },
        { label: "type", value: context.args?.type },
        { label: "year_from", value: context.args?.year_from },
        { label: "year_to", value: context.args?.year_to },
      ], theme) : undefined;

      if (isPartial) {
        return kylinFrameForResult(context, {
          name: "academic_search",
          theme,
          status: "pending",
          paramSummary,
          resultSummary: theme.fg("warning", "OpenAlex"),
          expanded,
          expandedParams,
          startedAt,
        });
      }

      if (details?.error) {
        return kylinFrameForResult(context, {
          name: "academic_search",
          theme,
          status: "error",
          paramSummary,
          errorMessage: details.error,
          expanded,
          expandedParams,
          startedAt,
        });
      }
      if (!details) {
        return kylinFrameForResult(context, {
          name: "academic_search",
          theme,
          status: "error",
          paramSummary,
          errorMessage: "No result",
          expanded,
          expandedParams,
          startedAt,
        });
      }

      const preview = extractTextPreview(result.content[0], 40);
      const shown = details.returnedResults !== undefined ? ` · ${details.returnedResults} shown` : "";
      return kylinFrameForResult(context, {
        name: "academic_search",
        theme,
        status: "success",
        paramSummary,
        resultSummary: theme.fg("success", `${details.totalResults?.toLocaleString() ?? "?"} results`) + theme.fg("dim", shown),
        expanded,
        expandedParams,
        expandedBody: expanded ? preview.lines.map((line) => theme.fg("dim", line)) : undefined,
        expandedTotalLines: preview.totalLines,
        startedAt,
      });
    },
    async execute(_toolCallId, params, signal, onUpdate) {
      const apiKey = getServiceKey("openalex", "OPENALEX_API_KEY");

      if (!apiKey) {
        return {
          content: [{ type: "text" as const, text: "Error: No OpenAlex API key found. Set OPENALEX_API_KEY or add it to agent/auth.json under 'openalex'." }],
          details: { query: params.query, phase: "done", error: "No API key configured" } as AcademicSearchDetails,
        };
      }

      onUpdate?.({
        content: [{ type: "text" as const, text: "Searching OpenAlex..." }],
        details: { query: params.query, phase: "searching" } as AcademicSearchDetails,
      });

      try {
        const { results, totalResults, error } = await searchOpenAlex(
          params.query,
          apiKey,
          signal,
          params.type,
          params.year_from,
          params.year_to,
        );

        if (error) {
          return {
            content: [{ type: "text" as const, text: `OpenAlex search error: ${error}` }],
            details: { query: params.query, phase: "done", error } as AcademicSearchDetails,
          };
        }

        const returnedCount = results.split("\n").filter((l) => /^\[\d+\]/.test(l)).length;

        return {
          content: [{ type: "text" as const, text: results || "No results found." }],
          details: {
            query: params.query,
            phase: "done",
            totalResults,
            returnedResults: returnedCount,
          } as AcademicSearchDetails,
        };
      } catch (error) {
        if (isAbortError(error)) {
          return {
            content: [{ type: "text" as const, text: "Search cancelled." }],
            details: { query: params.query, phase: "done", error: "Cancelled" } as AcademicSearchDetails,
          };
        }
        return {
          content: [{ type: "text" as const, text: `Search error: ${errorMessage(error)}` }],
          details: { query: params.query, phase: "done", error: errorMessage(error) } as AcademicSearchDetails,
        };
      }
    },
  });
}