import pLimit from "p-limit";
import { v4 as uuidv4 } from "uuid";
import { loadCompletedIds } from "../lib/checkpoint.ts";
import { appendJsonl } from "../lib/jsonl.ts";
import { classifyError } from "../lib/retry.ts";
import type { Annotation } from "@sec-cybert/schemas/annotation.ts";
import type { SessionLog } from "@sec-cybert/schemas/session.ts";
import type { Paragraph } from "@sec-cybert/schemas/paragraph.ts";
import { annotateParagraph, type AnnotateOpts } from "./annotate.ts";
import { PROMPT_VERSION } from "./prompts.ts";

export interface BatchOpts {
  modelId: string;
  stage: "stage1" | "stage2-judge" | "benchmark";
  outputPath: string;
  errorsPath: string;
  sessionsPath: string;
  concurrency?: number;
  promptVersion?: string;
  reasoningEffort?: "low" | "medium" | "high";
  limit?: number; // max paragraphs to process (for pilots)
}

interface LiveStats {
  processed: number;
  errored: number;
  retries: number;
  totalCostUsd: number;
  totalReasoningTokens: number;
  latencies: number[];
  startTime: number;
}

function percentile(sorted: number[], p: number): number {
  if (sorted.length === 0) return 0;
  const idx = Math.ceil((p / 100) * sorted.length) - 1;
  return sorted[Math.max(0, idx)]!;
}

function renderDashboard(
  stats: LiveStats,
  total: number,
  modelId: string,
  promptVersion: string,
  command: string,
): void {
  const elapsed = (Date.now() - stats.startTime) / 1000;
  const pct = total > 0 ? ((stats.processed / total) * 100).toFixed(1) : "0.0";
  const rate = elapsed > 0 ? (stats.processed / elapsed) * 60 : 0;
  const remaining = rate > 0 ? (total - stats.processed) / rate : 0;
  const etaMin = Math.floor(remaining);
  const etaSec = Math.round((remaining - etaMin) * 60);

  const sortedLat = [...stats.latencies].sort((a, b) => a - b);
  const p50 = percentile(sortedLat, 50);
  const p95 = percentile(sortedLat, 95);
  const p99 = percentile(sortedLat, 99);

  const avgReasoning =
    stats.processed > 0
      ? Math.round(stats.totalReasoningTokens / stats.processed)
      : 0;
  const reasoningPct =
    stats.totalReasoningTokens > 0
      ? ((stats.totalReasoningTokens / (stats.totalReasoningTokens + stats.processed * 50)) * 100).toFixed(1)
      : "0.0";

  const elapsedMin = Math.floor(elapsed / 60);
  const elapsedSec = Math.round(elapsed % 60);

  process.stderr.write(
    `\x1b[2K\r SEC-cyBERT │ ${command} │ ${modelId} │ ${promptVersion}\n` +
      ` ${"─".repeat(72)}\n` +
      ` Progress   ${stats.processed.toLocaleString()} / ${total.toLocaleString()}  (${pct}%)     ETA ${etaMin}m ${etaSec.toString().padStart(2, "0")}s\n` +
      ` Session    $${stats.totalCostUsd.toFixed(4)} │ ${elapsedMin}m ${elapsedSec.toString().padStart(2, "0")}s elapsed │ ${rate.toFixed(1)} para/min\n` +
      ` Errors     ${stats.errored} (${total > 0 ? ((stats.errored / Math.max(1, stats.processed + stats.errored)) * 100).toFixed(2) : "0.00"}%) │ ${stats.retries} retries\n` +
      ` Latency    p50: ${Math.round(p50)}ms │ p95: ${Math.round(p95)}ms │ p99: ${Math.round(p99)}ms\n` +
      ` Reasoning  avg ${avgReasoning} tokens/para │ ${reasoningPct}% of output tokens\n` +
      `\x1b[7A`, // move cursor back up
  );
}

/**
 * Run batch annotation on a list of paragraphs.
 * - Crash-safe: appends each annotation to JSONL immediately
 * - Resumable: skips already-completed paragraphs
 * - Observable: live dashboard to stderr
 * - Graceful shutdown on SIGINT
 */
export async function runBatch(
  paragraphs: Paragraph[],
  opts: BatchOpts,
): Promise<void> {
  const {
    modelId,
    stage,
    outputPath,
    errorsPath,
    sessionsPath,
    concurrency = 12,
    promptVersion = PROMPT_VERSION,
    reasoningEffort = "low",
    limit,
  } = opts;

  const runId = uuidv4();
  const command = `label:annotate`;

  // Resume: load already-completed IDs
  const { completedIds, skippedLines } = await loadCompletedIds(outputPath);
  if (skippedLines > 0) {
    process.stderr.write(`  ⚠ Skipped ${skippedLines} corrupt lines in ${outputPath}\n`);
  }

  // Filter to remaining work
  let remaining = paragraphs.filter((p) => !completedIds.has(p.id));
  if (limit !== undefined) remaining = remaining.slice(0, limit);
  const total = remaining.length;
  const totalWithResume = paragraphs.length;

  if (total === 0) {
    process.stderr.write(`  ✓ All ${totalWithResume} paragraphs already completed\n`);
    return;
  }

  process.stderr.write(
    `  Starting ${modelId} │ ${total} remaining of ${totalWithResume} │ concurrency=${concurrency}\n`,
  );

  const stats: LiveStats = {
    processed: 0,
    errored: 0,
    retries: 0,
    totalCostUsd: 0,
    totalReasoningTokens: 0,
    latencies: [],
    startTime: Date.now(),
  };

  // Graceful shutdown
  let stopping = false;
  const onSignal = () => {
    if (stopping) return;
    stopping = true;
    process.stderr.write("\n  ⏸ Stopping — waiting for in-flight requests...\n");
  };
  process.on("SIGINT", onSignal);
  process.on("SIGTERM", onSignal);

  // Dashboard refresh
  const dashboardInterval = setInterval(() => {
    renderDashboard(stats, total, modelId, promptVersion, command);
  }, 1000);

  // Concurrency limiter
  const limit_ = pLimit(concurrency);
  let consecutiveErrors = 0;

  const annotateOpts: AnnotateOpts = { modelId, stage, runId, promptVersion, reasoningEffort };

  const tasks = remaining.map((paragraph) =>
    limit_(async () => {
      if (stopping) return;
      if (consecutiveErrors >= 10) return;

      try {
        const annotation: Annotation = await annotateParagraph(paragraph, annotateOpts);
        await appendJsonl(outputPath, annotation);

        stats.processed++;
        stats.totalCostUsd += annotation.provenance.costUsd;
        stats.totalReasoningTokens += annotation.provenance.reasoningTokens;
        stats.latencies.push(annotation.provenance.latencyMs);
        consecutiveErrors = 0;
      } catch (error) {
        const kind = classifyError(error);

        if (kind === "budget") {
          stopping = true;
          process.stderr.write("\n  ✖ Budget exhausted (402). Stopping.\n");
          return;
        }

        stats.errored++;
        consecutiveErrors++;

        await appendJsonl(errorsPath, {
          paragraphId: paragraph.id,
          error: error instanceof Error ? error.message : String(error),
          errorKind: kind,
          modelId,
          timestamp: new Date().toISOString(),
        });

        if (consecutiveErrors >= 10) {
          stopping = true;
          process.stderr.write(
            "\n  ✖ 10 consecutive errors. Likely systemic — stopping.\n",
          );
        }
      }
    }),
  );

  await Promise.all(tasks);

  // Cleanup
  clearInterval(dashboardInterval);
  process.off("SIGINT", onSignal);
  process.off("SIGTERM", onSignal);

  // Final dashboard render
  process.stderr.write("\n".repeat(7)); // clear dashboard space
  renderDashboard(stats, total, modelId, promptVersion, command);
  process.stderr.write("\n".repeat(7)); // move past dashboard

  // Session log
  const endedAt = new Date().toISOString();
  const durationSeconds = (Date.now() - stats.startTime) / 1000;

  const sortedLat = [...stats.latencies].sort((a, b) => a - b);
  const exitReason = stopping
    ? consecutiveErrors >= 10
      ? "error_threshold"
      : "interrupted"
    : "complete";

  const session: SessionLog = {
    sessionId: runId,
    command,
    modelId,
    stage,
    promptVersion,
    startedAt: new Date(stats.startTime).toISOString(),
    endedAt,
    durationSeconds: Math.round(durationSeconds),
    paragraphsTotal: totalWithResume,
    paragraphsProcessed: stats.processed,
    paragraphsSkippedResume: completedIds.size,
    paragraphsErrored: stats.errored,
    costUsd: stats.totalCostUsd,
    reasoningTokensTotal: stats.totalReasoningTokens,
    avgLatencyMs: stats.latencies.length > 0
      ? Math.round(stats.latencies.reduce((a, b) => a + b, 0) / stats.latencies.length)
      : 0,
    p95LatencyMs: Math.round(percentile(sortedLat, 95)),
    throughputPerMinute: durationSeconds > 0
      ? Math.round((stats.processed / durationSeconds) * 60 * 10) / 10
      : 0,
    concurrency,
    exitReason,
  };

  await appendJsonl(sessionsPath, session);

  process.stderr.write(
    `\n  ✓ Done: ${stats.processed} processed, ${stats.errored} errors, $${stats.totalCostUsd.toFixed(4)} cost\n`,
  );
}
