import { resolve } from 'node:path';
import { Type } from '@earendil-works/pi-ai';
import { complete } from '@earendil-works/pi-ai/compat';
import type { ExtensionAPI } from '@earendil-works/pi-coding-agent';
import { Text } from '@earendil-works/pi-tui';
import { runPdfParseCommand } from './commands.js';
import { readConfig } from './config.js';
import { parsePdfToArtifacts, planPdfParse } from './parser.js';
import { renderPdfScreenshots } from './screenshots.js';
import { searchPdf } from './search.js';
import type { CompleteModel, ModelRegistryView, ParseProgress } from './types.js';

const completeModel: CompleteModel = (model, context, options) =>
  complete(model, context, options);

function progressStatus(progress: ParseProgress): string {
  const pages =
    progress.totalPages === undefined
      ? ''
      : ` ${progress.completedPages}/${progress.totalPages} pages`;
  return `${progress.phase}${pages} · ${progress.blockCount} blocks`;
}

function progressWidget(progress: ParseProgress): string[] {
  const counts = [
    progress.totalPages === undefined
      ? undefined
      : `Pages ${progress.completedPages}/${progress.totalPages}`,
    `Blocks ${progress.blockCount}`,
    progress.warningCount > 0 ? `Warnings ${progress.warningCount}` : undefined,
  ].filter((value): value is string => value !== undefined);
  return [`PDF Parse · ${progress.phase}`, progress.message, counts.join(' · ')];
}

export default function okraPiExtension(pi: ExtensionAPI): void {
  pi.registerCommand('pdf-parse', {
    description: 'Parse PDFs with Pi-authenticated vision models',
    handler: async (args, ctx) => {
      const onProgress = (progress: ParseProgress): void => {
        ctx.ui.setWorkingMessage(progress.message);
        ctx.ui.setStatus('pdf-parse', progressStatus(progress));
        ctx.ui.setWidget('pdf-parse-progress', progressWidget(progress), {
          placement: 'belowEditor',
        });
      };
      try {
        const message = await runPdfParseCommand(args, {
          cwd: ctx.cwd,
          registry: ctx.modelRegistry as ModelRegistryView,
          completeModel,
          onProgress,
        });
        ctx.ui.notify(message, 'info');
      } catch (error) {
        ctx.ui.notify(error instanceof Error ? error.message : String(error), 'error');
      } finally {
        ctx.ui.setWidget('pdf-parse-progress', undefined);
        ctx.ui.setStatus('pdf-parse', undefined);
        ctx.ui.setWorkingMessage();
      }
    },
  });

  pi.registerTool({
    name: 'pdf_parse',
    label: 'PDF Parse',
    description:
      'Parse a local PDF into layout-aware markdown and 0-1000 bounding boxes using a Pi-authenticated vision model with a validated prompt.',
    promptSnippet: 'Parse local PDFs with authenticated vision models and layout-aware prompts',
    promptGuidelines: [
      'Use dry_run when the user asks to inspect model choice or estimated cost before parsing.',
      'Honor max_cost_usd when the user specifies a budget.',
    ],
    parameters: Type.Object({
      source: Type.String({ description: 'Local PDF path, absolute or relative to the Pi working directory.' }),
      model: Type.Optional(
        Type.String({ description: 'Optional provider/model override from /pdf-parse models.' }),
      ),
      pages: Type.Optional(
        Type.String({ description: 'Optional pages, for example 2, 2-5, or 1-3,8.' }),
      ),
      dpi: Type.Optional(
        Type.Integer({
          minimum: 72,
          maximum: 300,
          description: 'PDF render resolution; defaults to 175 DPI.',
        }),
      ),
      max_cost_usd: Type.Optional(
        Type.Number({ minimum: 0, description: 'Reject a metered parse whose estimate exceeds this USD limit.' }),
      ),
      output_dir: Type.Optional(
        Type.String({ description: 'Artifact directory; defaults to <source-name>.okra beside the PDF.' }),
      ),
      dry_run: Type.Optional(
        Type.Boolean({ description: 'Plan model, prompt, pages, and cost without sending page images.' }),
      ),
    }),
    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      const registry = ctx.modelRegistry as ModelRegistryView;
      const config = await readConfig();
      const source = resolve(ctx.cwd, params.source);
      const reportProgress = (progress: ParseProgress): void => {
        onUpdate?.({
          content: [{ type: 'text', text: progress.message }],
          details: { progress },
        });
      };
      const options = {
        model: params.model,
        pages: params.pages,
        dpi: params.dpi,
        maxCostUsd: params.max_cost_usd,
        config,
        signal,
        onProgress: reportProgress,
      };

      if (params.dry_run) {
        const { plan } = await planPdfParse(source, registry, options);
        return {
          content: [
            {
              type: 'text',
              text: `PDF parse dry run: ${plan.model} with ${plan.prompt}@${plan.promptVersion}; ${plan.selectedPageCount} page(s); ${plan.costNote}. No images sent.`,
            },
          ],
          details: { plan },
        };
      }

      const { plan, result, artifacts } = await parsePdfToArtifacts(
        source,
        registry,
        completeModel,
        options,
        params.output_dir ? resolve(ctx.cwd, params.output_dir) : undefined,
      );
      return {
        content: [
          {
            type: 'text',
            text: `Parsed ${result.meta.pageCount} PDF page(s) into ${result.blocks.length} layout blocks with ${plan.model}. Artifacts: ${artifacts.outputDir}`,
          },
        ],
        details: { plan, meta: result.meta, usage: result.usage, artifacts },
      };
    },
    renderCall(args, theme) {
      let text = theme.fg('toolTitle', theme.bold('pdf_parse '));
      text += theme.fg('accent', args.source);
      if (args.pages) text += theme.fg('dim', ` pages=${args.pages}`);
      if (args.dry_run) text += theme.fg('warning', ' dry-run');
      return new Text(text, 0, 0);
    },
    renderResult(result, { expanded, isPartial }, theme, context) {
      const details = result.details as {
        progress?: ParseProgress;
        plan?: { model?: string; selectedPageCount?: number; costNote?: string };
        meta?: { pageCount?: number; warnings?: string[]; costUsd?: number };
        artifacts?: { outputDir?: string };
      };
      const content = result.content[0];
      const fullText = content?.type === 'text' ? content.text : '';
      if (isPartial) {
        const progress = details?.progress;
        return new Text(
          theme.fg('warning', progress?.message ?? 'Parsing PDF...'),
          0,
          0,
        );
      }
      if (context.isError) {
        return new Text(theme.fg('error', fullText.split('\n')[0] || 'PDF parse failed'), 0, 0);
      }
      const parts = [
        details?.meta?.pageCount !== undefined
          ? theme.fg('success', `${details.meta.pageCount} pages`)
          : theme.fg('success', 'planned'),
        details?.plan?.model ? theme.fg('muted', details.plan.model) : undefined,
        details?.artifacts?.outputDir
          ? theme.fg('dim', `→ ${details.artifacts.outputDir}`)
          : undefined,
      ].filter((part): part is string => part !== undefined);
      let text = parts.join(theme.fg('dim', ' · '));
      if (expanded && fullText) {
        text += `\n${theme.fg('toolOutput', fullText)}`;
      }
      return new Text(text, 0, 0);
    },
  });

  pi.registerTool({
    name: 'pdf_search',
    label: 'PDF Search',
    description:
      'Search a local PDF for text and return page numbers plus normalized 0-1000 bounding boxes. Uses existing pdf_parse artifacts when available, otherwise the native PDF text layer.',
    promptSnippet: 'Search PDFs for text with page and bounding-box locations',
    promptGuidelines: [
      'Use pdf_search before pdf_screenshot when looking for a known phrase.',
      'For scanned PDFs, run pdf_parse first so semantic blocks are available to search.',
    ],
    parameters: Type.Object({
      source: Type.String({ description: 'Local PDF path, absolute or relative to the Pi working directory.' }),
      query: Type.String({ minLength: 1, description: 'Text or phrase to locate.' }),
      pages: Type.Optional(
        Type.String({ description: 'Optional pages, for example 1-3,8.' }),
      ),
      case_sensitive: Type.Optional(
        Type.Boolean({ description: 'Use case-sensitive matching; defaults to false.' }),
      ),
      max_results: Type.Optional(
        Type.Integer({
          minimum: 1,
          maximum: 500,
          description: 'Maximum hits to return; defaults to 50.',
        }),
      ),
      artifact_dir: Type.Optional(
        Type.String({ description: 'Optional directory containing blocks.json from pdf_parse.' }),
      ),
    }),
    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      if (signal?.aborted) throw new Error('PDF search aborted.');
      onUpdate?.({
        content: [{ type: 'text', text: `Searching ${params.source} for "${params.query}"...` }],
        details: {},
      });
      const source = resolve(ctx.cwd, params.source);
      const result = await searchPdf(source, params.query, {
        pages: params.pages,
        caseSensitive: params.case_sensitive,
        maxResults: params.max_results,
        artifactDir: params.artifact_dir
          ? resolve(ctx.cwd, params.artifact_dir)
          : undefined,
      });
      const lines = [
        `Searched: ${source}`,
        `Query: ${params.query}`,
        `Source: ${result.searched}`,
        `Hits: ${result.hits.length}`,
        ...result.hits.slice(0, 20).map(
          (hit) =>
            `p${hit.page} [${hit.bbox.join(',')}]${hit.label ? ` ${hit.label}` : ''}: ${hit.text.replace(/\s+/g, ' ').slice(0, 240)}`,
        ),
      ];
      if (result.hits.length > 20) {
        lines.push(`...${result.hits.length - 20} more hits are available in tool details.`);
      }
      return {
        content: [{ type: 'text', text: lines.join('\n') }],
        details: { source, query: params.query, ...result },
      };
    },
  });

  pi.registerTool({
    name: 'pdf_screenshot',
    label: 'PDF Screenshot',
    description:
      'Render up to eight selected PDF pages as PNG image blocks for direct visual inspection and save them beside pdf_parse artifacts.',
    promptSnippet: 'Render selected PDF pages as images for visual inspection',
    promptGuidelines: [
      'Use pdf_search first for known text, then screenshot only the relevant pages.',
      'Keep screenshot calls to one to four pages unless broader visual comparison is necessary.',
    ],
    parameters: Type.Object({
      source: Type.String({ description: 'Local PDF path, absolute or relative to the Pi working directory.' }),
      pages: Type.Optional(
        Type.String({ description: 'Optional pages, for example 1-3,8; defaults to all with an 8-page cap.' }),
      ),
      dpi: Type.Optional(
        Type.Integer({
          minimum: 72,
          maximum: 300,
          description: 'Screenshot resolution; defaults to 150 DPI.',
        }),
      ),
      output_dir: Type.Optional(
        Type.String({ description: 'Base artifact directory; defaults to <source-name>.okra.' }),
      ),
    }),
    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      if (signal?.aborted) throw new Error('PDF screenshot rendering aborted.');
      const source = resolve(ctx.cwd, params.source);
      onUpdate?.({
        content: [{ type: 'text', text: `Rendering ${params.pages ?? 'selected'} PDF pages...` }],
        details: {},
      });
      const result = await renderPdfScreenshots(source, {
        pages: params.pages,
        dpi: params.dpi,
        outputDir: params.output_dir ? resolve(ctx.cwd, params.output_dir) : undefined,
      });
      if (signal?.aborted) throw new Error('PDF screenshot rendering aborted.');
      const content: Array<
        { type: 'text'; text: string } | { type: 'image'; data: string; mimeType: string }
      > = [
        {
          type: 'text',
          text: [
            `Rendered ${result.screenshots.length} page screenshot${result.screenshots.length === 1 ? '' : 's'}.`,
            `Output: ${result.outputDir}`,
            ...result.screenshots.map((screenshot) => `p${screenshot.page}: ${screenshot.path}`),
          ].join('\n'),
        },
        ...result.rendered.map((page) => ({
          type: 'image' as const,
          mimeType: 'image/png',
          data: Buffer.from(page.png).toString('base64'),
        })),
      ];
      return {
        content,
        details: {
          source,
          outputDir: result.outputDir,
          screenshots: result.screenshots,
        },
      };
    },
  });
}

export * from './commands.js';
export * from './config.js';
export * from './models.js';
export * from './layout.js';
export * from './parser.js';
export * from './pdf.js';
export * from './prompts.js';
export * from './screenshots.js';
export * from './search.js';
export * from './types.js';