import { mkdir, writeFile } from 'node:fs/promises';
import { basename, dirname, extname, resolve } from 'node:path';
import type { AssistantMessage, UserMessage, Usage } from '@earendil-works/pi-ai';
import {
  canonicalLabel,
  itemsToMarkdown,
  parseLayoutBlocks,
  swapGeminiBbox,
} from './layout.js';
import { readConfig } from './config.js';
import { renderPdfPages, pdfPageCount } from './pdf.js';
import { promptRef } from './prompts.js';
import {
  formatAuthResolutionFailure,
  formatProviderFailure,
} from './provider-errors.js';
import { selectModel } from './models.js';
import type {
  CompleteModel,
  LayoutBlock,
  ModelCandidate,
  ModelRegistryView,
  PageParse,
  PageSelection,
  ParseOptions,
  ParsePlan,
  ParseProgress,
  ParseResult,
  UsageSummary,
  WrittenArtifacts,
} from './types.js';

const EMPTY_USAGE: UsageSummary = {
  input: 0,
  output: 0,
  cacheRead: 0,
  cacheWrite: 0,
  reasoning: 0,
  totalTokens: 0,
  referenceCostUsd: 0,
};

function reportProgress(options: ParseOptions, progress: ParseProgress): void {
  options.onProgress?.(progress);
}

export function parsePageSelection(value: string | undefined, pageCount: number): PageSelection {
  if (pageCount < 1) throw new Error('The PDF contains no pages.');
  if (!value || value.trim().toLowerCase() === 'all') {
    return {
      pages: Array.from({ length: pageCount }, (_, index) => index + 1),
      description: pageCount === 1 ? '1' : `1-${pageCount}`,
    };
  }
  const pages = new Set<number>();
  const parts = value.split(',').map((part) => part.trim());
  if (parts.some((part) => !part)) {
    throw new Error(`Invalid page selection "${value}". Use values such as 2, 2-5, or 1-3,8.`);
  }
  for (const part of parts) {
    const match = /^(\d+)(?:-(\d+))?$/.exec(part);
    if (!match) {
      throw new Error(`Invalid page selection "${value}". Use values such as 2, 2-5, or 1-3,8.`);
    }
    const from = Number.parseInt(match[1], 10);
    const to = Number.parseInt(match[2] ?? match[1], 10);
    if (from < 1 || to < from || to > pageCount) {
      throw new Error(`Page selection "${value}" is outside this ${pageCount}-page PDF.`);
    }
    for (let page = from; page <= to; page++) pages.add(page);
  }
  const selected = [...pages].sort((a, b) => a - b);
  return { pages: selected, description: value.replace(/\s+/g, '') };
}

function costNote(candidate: ModelCandidate, selectedPages: number): string {
  if (candidate.access === 'subscription') return 'uses the provider subscription/quota';
  if (candidate.access === 'local') return 'runs against a local provider ($0 API cost)';
  if (candidate.estimatedUsdPerPage !== undefined) {
    return `estimated $${(candidate.estimatedUsdPerPage * selectedPages).toFixed(4)} for ${selectedPages} page${selectedPages === 1 ? '' : 's'}`;
  }
  return 'provider price is unknown';
}

export async function planPdfParse(
  source: string,
  registry: ModelRegistryView,
  options: ParseOptions = {},
): Promise<{ plan: ParsePlan; candidate: ModelCandidate }> {
  const absoluteSource = resolve(source);
  reportProgress(options, {
    phase: 'planning',
    message: `Inspecting ${basename(absoluteSource)} and selecting an authenticated vision model...`,
    completedPages: 0,
    blockCount: 0,
    warningCount: 0,
  });
  const pageCount = await pdfPageCount(absoluteSource);
  const selectedPages = parsePageSelection(options.pages, pageCount);
  const selectedPageCount = selectedPages.pages.length;
  const config = options.config ?? (await readConfig());
  const candidate = selectModel(registry, { requested: options.model, config });
  const prompt = candidate.prompt;
  if (!prompt) throw new Error(`No parser prompt matched ${candidate.fullId}.`);
  const estimatedCostUsd =
    candidate.access === 'metered' && candidate.estimatedUsdPerPage !== undefined
      ? candidate.estimatedUsdPerPage * selectedPageCount
      : undefined;

  if (options.maxCostUsd !== undefined) {
    if (!Number.isFinite(options.maxCostUsd) || options.maxCostUsd < 0) {
      throw new Error('max_cost_usd must be a non-negative number.');
    }
    if (candidate.access === 'unknown') {
      throw new Error(
        `Cannot enforce a $${options.maxCostUsd.toFixed(4)} cost limit because ${candidate.fullId} has no known Pi pricing. Choose a priced, subscription, or local model.`,
      );
    }
    if (estimatedCostUsd !== undefined && estimatedCostUsd > options.maxCostUsd) {
      throw new Error(
        `Estimated cost $${estimatedCostUsd.toFixed(4)} exceeds max_cost_usd $${options.maxCostUsd.toFixed(4)}. Select fewer pages or another model.`,
      );
    }
  }

  const plan: ParsePlan = {
    source: absoluteSource,
    pageCount,
    selectedPages,
    selectedPageCount,
    model: candidate.fullId,
    provider: candidate.model.provider,
    prompt: prompt.id,
    promptVersion: prompt.version,
    access: candidate.access,
    ...(estimatedCostUsd !== undefined ? { estimatedCostUsd } : {}),
    costNote: costNote(candidate, selectedPageCount),
  };
  reportProgress(options, {
    phase: 'planned',
    message: `Selected ${candidate.fullId} with ${promptRef(prompt)} for ${selectedPageCount} page${selectedPageCount === 1 ? '' : 's'}.`,
    completedPages: 0,
    totalPages: selectedPageCount,
    blockCount: 0,
    warningCount: 0,
    model: candidate.fullId,
    prompt: promptRef(prompt),
  });
  return {
    candidate,
    plan,
  };
}

function assistantText(message: AssistantMessage, model: ModelCandidate['model']): string {
  if (message.stopReason === 'aborted') {
    throw new Error('PDF parse aborted.');
  }
  if (message.stopReason === 'error') {
    throw new Error(
      formatProviderFailure(model, message.errorMessage ?? 'Provider stopped with an error.'),
    );
  }
  const text = message.content
    .filter((part): part is { type: 'text'; text: string } => part.type === 'text')
    .map((part) => part.text)
    .join('\n')
    .trim();
  if (!text) throw new Error('The model returned no text.');
  return text;
}

function summarizeUsage(usage: Usage): UsageSummary {
  return {
    input: usage.input,
    output: usage.output,
    cacheRead: usage.cacheRead,
    cacheWrite: usage.cacheWrite,
    reasoning: usage.reasoning ?? 0,
    totalTokens: usage.totalTokens,
    referenceCostUsd: usage.cost.total,
  };
}

function addUsage(a: UsageSummary, b: UsageSummary): UsageSummary {
  return {
    input: a.input + b.input,
    output: a.output + b.output,
    cacheRead: a.cacheRead + b.cacheRead,
    cacheWrite: a.cacheWrite + b.cacheWrite,
    reasoning: a.reasoning + b.reasoning,
    totalTokens: a.totalTokens + b.totalTokens,
    referenceCostUsd: a.referenceCostUsd + b.referenceCostUsd,
  };
}

export class PdfParseDecodeError extends Error {
  artifacts?: WrittenArtifacts;

  constructor(
    message: string,
    readonly plan: ParsePlan,
    readonly result: ParseResult,
  ) {
    super(message);
    this.name = 'PdfParseDecodeError';
  }
}

export async function parsePdf(
  source: string,
  registry: ModelRegistryView,
  completeModel: CompleteModel,
  options: ParseOptions = {},
): Promise<{ plan: ParsePlan; result: ParseResult }> {
  const startedAt = Date.now();
  const { plan, candidate } = await planPdfParse(source, registry, options);
  const prompt = candidate.prompt;
  if (!prompt) throw new Error(`No parser prompt matched ${candidate.fullId}.`);

  reportProgress(options, {
    phase: 'authenticating',
    message: `Resolving Pi credentials for ${candidate.fullId}...`,
    completedPages: 0,
    totalPages: plan.selectedPageCount,
    blockCount: 0,
    warningCount: 0,
    model: candidate.fullId,
    prompt: promptRef(prompt),
  });
  const auth = await registry.getApiKeyAndHeaders(candidate.model);
  if (!auth.ok) {
    throw new Error(formatAuthResolutionFailure(candidate.model, auth.error));
  }

  reportProgress(options, {
    phase: 'rendering',
    message: `Rendering ${plan.selectedPageCount} PDF page${plan.selectedPageCount === 1 ? '' : 's'} locally...`,
    completedPages: 0,
    totalPages: plan.selectedPageCount,
    blockCount: 0,
    warningCount: 0,
    model: candidate.fullId,
    prompt: promptRef(prompt),
  });
  const renderedPages = await renderPdfPages(
    plan.source,
    plan.selectedPages,
    options.dpi,
  );
  const pages: PageParse[] = [];
  const warnings: string[] = [];
  let usage = { ...EMPTY_USAGE };
  let decodedBlockCount = 0;

  for (const [index, rendered] of renderedPages.entries()) {
    if (options.signal?.aborted) throw new Error('PDF parse aborted.');
    reportProgress(options, {
      phase: 'requesting',
      message: `Parsing page ${index + 1}/${renderedPages.length} (PDF page ${rendered.page}) with ${candidate.fullId}...`,
      completedPages: index,
      totalPages: renderedPages.length,
      currentPage: rendered.page,
      blockCount: decodedBlockCount,
      warningCount: warnings.length,
      model: candidate.fullId,
      prompt: promptRef(prompt),
    });
    const message: UserMessage = {
      role: 'user',
      content: [
        { type: 'text', text: prompt.user },
        {
          type: 'image',
          data: Buffer.from(rendered.png).toString('base64'),
          mimeType: 'image/png',
        },
      ],
      timestamp: Date.now(),
    };
    let response: AssistantMessage;
    try {
      response = await completeModel(
        candidate.model,
        { systemPrompt: prompt.system, messages: [message] },
        {
          apiKey: auth.apiKey,
          headers: auth.headers,
          env: auth.env,
          maxTokens: Math.min(8192, candidate.model.maxTokens),
          signal: options.signal,
        },
      );
    } catch (error) {
      if (options.signal?.aborted) throw new Error('PDF parse aborted.');
      throw new Error(
        formatProviderFailure(
          candidate.model,
          error instanceof Error ? error.message : String(error),
        ),
      );
    }
    const text = assistantText(response, candidate.model);
    const decoded = parseLayoutBlocks(text);
    const normalized = prompt.bboxOrder === 'yxyx' ? swapGeminiBbox(decoded) : decoded;
    const blocks: LayoutBlock[] = normalized.map((block) => ({
      label: canonicalLabel(block.label),
      bbox: block.bbox,
      text: block.text,
      page: rendered.page,
    }));
    decodedBlockCount += blocks.length;
    if (blocks.length === 0) warnings.push(`page ${rendered.page}: 0 layout blocks decoded`);
    const pageUsage = summarizeUsage(response.usage);
    usage = addUsage(usage, pageUsage);
    pages.push({
      page: rendered.page,
      markdown: itemsToMarkdown(normalized),
      blocks,
      rawText: text,
      usage: pageUsage,
    });
    reportProgress(options, {
      phase: 'decoded',
      message: `Page ${index + 1}/${renderedPages.length}: decoded ${blocks.length} layout block${blocks.length === 1 ? '' : 's'} (${decodedBlockCount} total).`,
      completedPages: index + 1,
      totalPages: renderedPages.length,
      currentPage: rendered.page,
      blockCount: decodedBlockCount,
      warningCount: warnings.length,
      model: candidate.fullId,
      prompt: promptRef(prompt),
    });

    if (
      candidate.access === 'metered' &&
      options.maxCostUsd !== undefined &&
      usage.referenceCostUsd > options.maxCostUsd
    ) {
      throw new Error(
        `Provider-reported cost reached $${usage.referenceCostUsd.toFixed(4)} after page ${rendered.page}, above max_cost_usd $${options.maxCostUsd.toFixed(4)}. No further pages were sent.`,
      );
    }
  }

  const blocks = pages.flatMap((page) => page.blocks);
  const markdown = pages.map((page) => page.markdown).join('\n\n');
  const result: ParseResult = {
    markdown,
    blocks,
    pages,
    usage,
    meta: {
      parserId: 'layout-vlm',
      provider: candidate.model.provider,
      model: candidate.fullId,
      prompt: prompt.id,
      promptVersion: prompt.version,
      access: candidate.access,
      pageCount: pages.length,
      durationMs: Date.now() - startedAt,
      ...(plan.estimatedCostUsd !== undefined ? { estimatedCostUsd: plan.estimatedCostUsd } : {}),
      ...(candidate.access === 'metered' ? { costUsd: usage.referenceCostUsd } : {}),
      ...(candidate.access !== 'metered' && usage.referenceCostUsd > 0
        ? { referenceCostUsd: usage.referenceCostUsd }
        : {}),
      warnings,
    },
  };
  if (blocks.length === 0) {
    throw new PdfParseDecodeError(
      `${candidate.fullId} returned text for ${pages.length} page${pages.length === 1 ? '' : 's'}, but ${promptRef(prompt)} decoded 0 layout blocks.`,
      plan,
      result,
    );
  }
  return { plan, result };
}

export function defaultOutputDir(source: string): string {
  const extension = extname(source);
  const stem = basename(source, extension);
  return resolve(dirname(source), `${stem}.okra`);
}

export async function writeParseArtifacts(
  source: string,
  result: ParseResult,
  outputDir = defaultOutputDir(source),
): Promise<WrittenArtifacts> {
  const absoluteOutputDir = resolve(outputDir);
  await mkdir(absoluteOutputDir, { recursive: true });
  const markdownPath = resolve(absoluteOutputDir, 'doc.md');
  const blocksPath = resolve(absoluteOutputDir, 'blocks.json');
  const manifestPath = resolve(absoluteOutputDir, 'manifest.json');
  const rawDir = resolve(absoluteOutputDir, 'raw');
  await mkdir(rawDir, { recursive: true });
  const rawPaths = result.pages.map((page) =>
    resolve(rawDir, `page-${String(page.page).padStart(4, '0')}.txt`),
  );
  const manifest = {
    meta: result.meta,
    usage: result.usage,
    pages: result.pages.map((page) => ({
      page: page.page,
      blockCount: page.blocks.length,
      rawResponse: `raw/page-${String(page.page).padStart(4, '0')}.txt`,
      usage: page.usage,
    })),
  };
  await Promise.all([
    writeFile(markdownPath, `${result.markdown}\n`, 'utf8'),
    writeFile(blocksPath, `${JSON.stringify(result.blocks, null, 2)}\n`, 'utf8'),
    writeFile(manifestPath, `${JSON.stringify(manifest, null, 2)}\n`, 'utf8'),
    ...result.pages.map((page, index) => writeFile(rawPaths[index], `${page.rawText}\n`, 'utf8')),
  ]);
  return {
    outputDir: absoluteOutputDir,
    markdownPath,
    blocksPath,
    manifestPath,
    rawDir,
    rawPaths,
  };
}

function diagnosticErrorMessage(error: PdfParseDecodeError): string {
  const rawPreview = error.result.pages
    .map((page) => page.rawText)
    .join('\n')
    .replace(/\s+/g, ' ')
    .trim()
    .slice(0, 240);
  const diagnostics = error.artifacts
    ? `Diagnostics: ${error.artifacts.outputDir}\nRaw responses: ${error.artifacts.rawDir}`
    : 'The raw response could not be written to disk.';
  return [
    error.message,
    diagnostics,
    ...(rawPreview
      ? [`Response preview: ${rawPreview}${rawPreview.length === 240 ? '…' : ''}`]
      : []),
    'Try another validated model, or inspect the raw response to revalidate this model/prompt pair.',
  ].join('\n');
}

export async function parsePdfToArtifacts(
  source: string,
  registry: ModelRegistryView,
  completeModel: CompleteModel,
  options: ParseOptions = {},
  outputDir?: string,
): Promise<{ plan: ParsePlan; result: ParseResult; artifacts: WrittenArtifacts }> {
  try {
    const { plan, result } = await parsePdf(source, registry, completeModel, options);
    reportProgress(options, {
      phase: 'writing',
      message: `Writing markdown, layout blocks, manifest, and raw responses...`,
      completedPages: result.pages.length,
      totalPages: result.pages.length,
      blockCount: result.blocks.length,
      warningCount: result.meta.warnings.length,
      model: plan.model,
      prompt: `${plan.prompt}@${plan.promptVersion}`,
    });
    const artifacts = await writeParseArtifacts(source, result, outputDir);
    reportProgress(options, {
      phase: 'complete',
      message: `Parsed ${result.pages.length} page${result.pages.length === 1 ? '' : 's'} into ${result.blocks.length} layout block${result.blocks.length === 1 ? '' : 's'}. Artifacts: ${artifacts.outputDir}`,
      completedPages: result.pages.length,
      totalPages: result.pages.length,
      blockCount: result.blocks.length,
      warningCount: result.meta.warnings.length,
      model: plan.model,
      prompt: `${plan.prompt}@${plan.promptVersion}`,
    });
    return { plan, result, artifacts };
  } catch (error) {
    if (error instanceof PdfParseDecodeError) {
      reportProgress(options, {
        phase: 'writing',
        message: 'No layout blocks decoded; preserving the raw model response and diagnostics...',
        completedPages: error.result.pages.length,
        totalPages: error.plan.selectedPageCount,
        blockCount: 0,
        warningCount: error.result.meta.warnings.length,
        model: error.plan.model,
        prompt: `${error.plan.prompt}@${error.plan.promptVersion}`,
      });
      try {
        error.artifacts = await writeParseArtifacts(source, error.result, outputDir);
      } catch (artifactError) {
        error.message = `${error.message}\nCould not write diagnostics: ${artifactError instanceof Error ? artifactError.message : String(artifactError)}`;
      }
      error.message = diagnosticErrorMessage(error);
    }
    reportProgress(options, {
      phase: 'failed',
      message: error instanceof Error ? error.message : String(error),
      completedPages: error instanceof PdfParseDecodeError ? error.result.pages.length : 0,
      totalPages: error instanceof PdfParseDecodeError ? error.plan.selectedPageCount : undefined,
      blockCount: error instanceof PdfParseDecodeError ? error.result.blocks.length : 0,
      warningCount: error instanceof PdfParseDecodeError ? error.result.meta.warnings.length : 0,
      model: error instanceof PdfParseDecodeError ? error.plan.model : undefined,
      prompt:
        error instanceof PdfParseDecodeError
          ? `${error.plan.prompt}@${error.plan.promptVersion}`
          : undefined,
    });
    throw error;
  }
}