/**
 * Golden set annotation via Claude Agent SDK.
 *
 * Uses the user's Claude Code subscription (OAuth) instead of API keys,
 * calling Opus 4.6 through the Agent SDK's `query()` with structured output.
 * Designed for the ~1,200 human-labeled paragraphs.
 *
 * Key differences from Stage 1/2 (OpenRouter):
 * - Full codebook (docs/LABELING-CODEBOOK.md) + v2.5 prompt as system prompt
 * - Saves reasoning traces (Opus adaptive thinking) alongside annotations
 * - Saves raw confidence values before coercion
 * - No API cost — runs on Max subscription
 */

import { readFile } from "node:fs/promises";
import { query } from "@anthropic-ai/claude-agent-sdk";
import { z } from "zod";
import { v4 as uuidv4 } from "uuid";
import {
  LabelOutputRaw,
  ContentCategory,
  SpecificityLabel,
  FactType,
  toLabelOutput,
} from "@sec-cybert/schemas/label.ts";
import type { Annotation } from "@sec-cybert/schemas/annotation.ts";
import type { Paragraph } from "@sec-cybert/schemas/paragraph.ts";
import { SYSTEM_PROMPT, buildUserPrompt, PROMPT_VERSION } from "./prompts.ts";
import { loadCompletedIds } from "../lib/checkpoint.ts";
import { appendJsonl } from "../lib/jsonl.ts";

/**
 * JSON-Schema-safe version of LabelOutputRaw.
 *
 * The real LabelOutputRaw uses .transform() and .pipe() (for Confidence
 * and SpecificFact) which cannot convert to JSON Schema. We define a
 * plain schema here for the Agent SDK's outputFormat, then parse the
 * result through the real LabelOutputRaw for validation + transforms.
 *
 * Confidence is a number (0-1) here since Opus naturally outputs numeric
 * confidence. We save the raw value separately and coerce via LabelOutputRaw.
 */
const GoldenOutputSchema = z.object({
  content_category: z.enum(ContentCategory.options),
  specific_facts: z.array(
    z.object({
      fact: z.string(),
      type: z.enum(FactType.options),
    }),
  ),
  specificity: z.enum(SpecificityLabel.options),
  category_confidence: z.number(),
  specificity_confidence: z.number(),
  reasoning: z.string(),
});

/** Extra fields saved alongside annotation but NOT in the label block. */
interface GoldenExtras {
  /** Opus's reasoning trace from adaptive thinking. */
  thinking: string;
  /** Raw confidence values before coercion to high/medium/low. */
  rawCategoryConfidence: number;
  rawSpecificityConfidence: number;
  /** Specific facts as returned by the model (before any transform). */
  rawSpecificFacts: Array<{ fact: string; type: string }>;
}

/** What we write to the JSONL — Annotation + extras at top level. */
type GoldenAnnotation = Annotation & { golden: GoldenExtras };

export interface GoldenBatchOpts {
  outputPath: string;
  errorsPath: string;
  limit?: number;
  /** Delay between requests in ms. Default 1000 (1 req/s). */
  delayMs?: number;
  /** Number of concurrent workers. Default 1 (serial). */
  concurrency?: number;
}

/** Build the enhanced system prompt: full codebook + v2.5 operational prompt + JSON schema. */
async function buildGoldenSystemPrompt(): Promise<string> {
  const codebookPath = new URL("../../../docs/LABELING-CODEBOOK.md", import.meta.url).pathname;
  const codebook = await readFile(codebookPath, "utf-8");

  // Strip the old "LLM Response Schema" section from the codebook to avoid
  // conflicting with the actual JSON schema we enforce via outputFormat.
  // The old section uses specificity_level (integer) instead of specificity (string label).
  const schemaHeading = "## LLM Response Schema";
  const codebookTrimmed = codebook.includes(schemaHeading)
    ? codebook.slice(0, codebook.indexOf(schemaHeading)).trimEnd()
    : codebook;

  const jsonSchema = JSON.stringify(z.toJSONSchema(GoldenOutputSchema), null, 2);

  return `${codebookTrimmed}

═══════════════════════════════════════════════════════════════════════
OPERATIONAL PROMPT (v${PROMPT_VERSION})
The codebook above is the authoritative reference. The prompt below
is the condensed operational version used for all annotation.
═══════════════════════════════════════════════════════════════════════

${SYSTEM_PROMPT}

═══════════════════════════════════════════════════════════════════════
OUTPUT JSON SCHEMA
You MUST return JSON matching this exact schema. Use text labels for
specificity (not integers). Confidence is a number 0-1.
═══════════════════════════════════════════════════════════════════════

${jsonSchema}`;
}

/** Extract from Agent SDK result messages. */
interface QueryResult {
  structuredOutput: unknown;
  thinking: string;
  inputTokens: number;
  outputTokens: number;
}

/**
 * Annotate a single paragraph via Claude Agent SDK.
 * Each call spawns a fresh, isolated query() — no tools, no hooks,
 * no settings, no session persistence.
 */
async function annotateGolden(
  paragraph: Paragraph,
  runId: string,
  systemPrompt: string,
): Promise<GoldenAnnotation> {
  const requestedAt = new Date().toISOString();
  const start = Date.now();

  let result: QueryResult = {
    structuredOutput: null,
    thinking: "",
    inputTokens: 0,
    outputTokens: 0,
  };

  // Prevent git pull and other non-essential traffic when running concurrently
  process.env.CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC = "1";

  for await (const message of query({
    prompt: buildUserPrompt(paragraph),
    options: {
      model: "claude-opus-4-6",
      systemPrompt,
      outputFormat: {
        type: "json_schema",
        schema: z.toJSONSchema(GoldenOutputSchema),
      },
      // No tools — pure classification
      allowedTools: [],
      disallowedTools: ["Bash", "Read", "Write", "Edit", "Glob", "Grep", "WebSearch", "WebFetch", "Agent", "AskUserQuestion"],
      // Isolation: no hooks, no settings, no session persistence, no plugins
      hooks: {},
      settingSources: [],
      persistSession: false,
      plugins: [],
      // Single-turn: one prompt → one structured response
      maxTurns: 1,
      permissionMode: "dontAsk",
    },
  })) {
    const msg = message as Record<string, unknown>;

    // Extract thinking from assistant messages
    if (msg.type === "assistant") {
      const inner = msg.message as { content?: Array<{ type: string; thinking?: string }> };
      if (inner.content) {
        for (const block of inner.content) {
          if (block.type === "thinking" && block.thinking) {
            result.thinking += block.thinking;
          }
        }
      }
    }

    // Extract structured output and usage from result message
    if (msg.type === "result" && msg.subtype === "success") {
      if (msg.structured_output) {
        result.structuredOutput = msg.structured_output;
      } else if (typeof msg.result === "string") {
        const raw = msg.result;
        const fenceMatch = raw.match(/```(?:json)?\s*\n([\s\S]*?)\n```/);
        const jsonStr = fenceMatch ? fenceMatch[1]! : raw;
        try {
          result.structuredOutput = JSON.parse(jsonStr);
        } catch {
          // not valid JSON
        }
      }

      // Token usage from modelUsage (more detailed than top-level usage)
      const modelUsage = msg.modelUsage as Record<string, { inputTokens?: number; outputTokens?: number; cacheReadInputTokens?: number }> | undefined;
      const opusUsage = modelUsage?.["claude-opus-4-6"];
      if (opusUsage) {
        result.inputTokens = (opusUsage.inputTokens ?? 0) + (opusUsage.cacheReadInputTokens ?? 0);
        result.outputTokens = opusUsage.outputTokens ?? 0;
      }
    }
  }

  const latencyMs = Date.now() - start;

  if (!result.structuredOutput) {
    throw new Error(`No structured output from Opus for ${paragraph.id}`);
  }

  // Save raw values before coercion
  const raw = result.structuredOutput as Record<string, unknown>;
  const rawCategoryConfidence = typeof raw.category_confidence === "number" ? raw.category_confidence : 0;
  const rawSpecificityConfidence = typeof raw.specificity_confidence === "number" ? raw.specificity_confidence : 0;
  const rawSpecificFacts = Array.isArray(raw.specific_facts)
    ? (raw.specific_facts as Array<{ fact: string; type: string }>)
    : [];

  // Parse through the real schema with transforms (Confidence coercion, etc.)
  const parsed = LabelOutputRaw.parse(result.structuredOutput);

  return {
    paragraphId: paragraph.id,
    label: toLabelOutput(parsed),
    provenance: {
      modelId: "anthropic/claude-opus-4-6",
      provider: "anthropic",
      generationId: "agent-sdk",
      stage: "benchmark",
      runId,
      promptVersion: `${PROMPT_VERSION}+codebook`,
      inputTokens: result.inputTokens,
      outputTokens: result.outputTokens,
      reasoningTokens: 0, // included in outputTokens, not broken out by SDK
      costUsd: 0, // subscription — no per-request cost
      latencyMs,
      requestedAt,
    },
    golden: {
      thinking: result.thinking,
      rawCategoryConfidence,
      rawSpecificityConfidence,
      rawSpecificFacts,
    },
  };
}

/**
 * Run golden set annotation through the Agent SDK.
 * Supports concurrent workers for parallelism.
 * Crash-safe with JSONL checkpoint resume.
 */
export async function runGoldenBatch(
  paragraphs: Paragraph[],
  opts: GoldenBatchOpts,
): Promise<void> {
  const { outputPath, errorsPath, limit, delayMs = 1000, concurrency = 1 } = opts;
  const runId = uuidv4();

  // Build system prompt once (codebook + operational prompt)
  const systemPrompt = await buildGoldenSystemPrompt();
  process.stderr.write(`  System prompt: ${(systemPrompt.length / 1024).toFixed(1)}KB\n`);

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

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

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

  process.stderr.write(
    `  Starting golden annotation │ Opus 4.6 (Agent SDK) │ ${total} remaining of ${paragraphs.length} │ concurrency=${concurrency}\n`,
  );

  let processed = 0;
  let errored = 0;
  const startTime = Date.now();

  // Serialized file writes to prevent corruption
  let writeQueue = Promise.resolve();
  function safeAppend(path: string, data: object) {
    writeQueue = writeQueue.then(() => appendJsonl(path, data));
    return writeQueue;
  }

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

  // Dashboard refresh
  function renderStatus() {
    const elapsed = (Date.now() - startTime) / 1000;
    const rate = elapsed > 0 ? (processed / elapsed) * 60 : 0;
    const etaMin = rate > 0 ? Math.round((total - processed) / rate) : 0;
    process.stderr.write(
      `\x1b[2K\r  ${processed}/${total} (${((processed / total) * 100).toFixed(1)}%) │ ${rate.toFixed(1)} para/min │ ETA ${etaMin}m │ ${errored} errors`,
    );
  }
  const dashboardInterval = setInterval(renderStatus, 2000);

  // Worker pool: N concurrent workers pulling from shared queue
  let nextIdx = 0;
  async function worker() {
    while (nextIdx < remaining.length && !stopping) {
      const idx = nextIdx++;
      const paragraph = remaining[idx]!;

      try {
        const annotation = await annotateGolden(paragraph, runId, systemPrompt);
        await safeAppend(outputPath, annotation);
        processed++;
      } catch (error) {
        errored++;
        await safeAppend(errorsPath, {
          paragraphId: paragraph.id,
          error: error instanceof Error ? error.message : String(error),
          modelId: "anthropic/claude-opus-4-6",
          timestamp: new Date().toISOString(),
        });

        if (errored >= 10 && processed === 0) {
          stopping = true;
          process.stderr.write("\n  ✖ 10 errors with no successes. Stopping.\n");
        }
      }

      // Per-worker delay between requests
      if (!stopping) {
        await new Promise((r) => setTimeout(r, delayMs));
      }
    }
  }

  const workers = Array.from(
    { length: Math.min(concurrency, remaining.length) },
    () => worker(),
  );
  await Promise.all(workers);

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

  renderStatus();
  process.stderr.write(
    `\n  ✓ Golden annotation done: ${processed} processed, ${errored} errors\n`,
  );
}
