/**
 * Generic LLM caller for the analyzer pipeline.
 *
 * One round-trip is:
 *   1. Send (system, user) to the chosen model.
 *   2. Extract the assistant's text content.
 *   3. JSON.parse + Zod-validate against the caller's schema.
 *   4. On failure: retry ONCE with the prior raw response + parse error in a
 *      follow-up user turn, then validate again.
 *   5. After two failures: throw.
 *
 * The system prompt is marked with `cache_control: ephemeral`. Anthropic
 * silently no-ops caching when the prefix is below the model's minimum
 * (~2-4K tokens) — for our short stage prompts this often won't fire, which
 * is fine; cost is unaffected when caching is skipped.
 */

import type Anthropic from '@anthropic-ai/sdk';
import type { ZodType } from 'zod';
import { getAnthropicClient } from './client';
import { estimateCostUsd, type TokenUsage } from './pricing';
import { type ModelId, OPUS } from './models';

export interface LLMCallOptions<T> {
  model: ModelId;
  system: string;
  user: string;
  schema: ZodType<T>;
  maxTokens: number;
  /** Override the singleton (test injection). */
  client?: Anthropic;
}

export interface LLMCallResult<T> {
  data: T;
  usage: TokenUsage;
  estimated_cost_usd: number;
  /** Number of attempts made (1 = first try succeeded; 2 = retry succeeded). */
  attempts: 1 | 2;
  /** Raw response text for debugging / model_traces. */
  raw_response: string;
}

/**
 * Concatenate token usage from two calls (used to track total cost across
 * the original attempt + retry).
 */
function addUsage(a: TokenUsage, b: TokenUsage): TokenUsage {
  return {
    input_tokens: a.input_tokens + b.input_tokens,
    output_tokens: a.output_tokens + b.output_tokens,
    cache_creation_input_tokens:
      (a.cache_creation_input_tokens ?? 0) + (b.cache_creation_input_tokens ?? 0),
    cache_read_input_tokens:
      (a.cache_read_input_tokens ?? 0) + (b.cache_read_input_tokens ?? 0),
  };
}

interface MessagesCreateBody {
  model: string;
  max_tokens: number;
  system: Array<{
    type: 'text';
    text: string;
    cache_control?: { type: 'ephemeral' };
  }>;
  messages: Array<{ role: 'user' | 'assistant'; content: string }>;
}

function buildBody(opts: {
  model: ModelId;
  system: string;
  history: Array<{ role: 'user' | 'assistant'; content: string }>;
  maxTokens: number;
}): MessagesCreateBody {
  return {
    model: opts.model,
    max_tokens: opts.maxTokens,
    system: [
      {
        type: 'text',
        text: opts.system,
        cache_control: { type: 'ephemeral' },
      },
    ],
    messages: opts.history,
  };
}

function extractText(response: Anthropic.Message): string {
  const parts: string[] = [];
  for (const block of response.content) {
    if (block.type === 'text') parts.push(block.text);
  }
  return parts.join('').trim();
}

/**
 * Some models occasionally wrap JSON in code fences despite "respond ONLY with
 * JSON" instructions. Strip a single ```json ... ``` fence if present so the
 * happy path doesn't bounce into the retry just for that.
 */
function unwrapFences(text: string): string {
  const trimmed = text.trim();
  if (trimmed.startsWith('```')) {
    const stripped = trimmed
      .replace(/^```(?:json)?\s*/i, '')
      .replace(/```\s*$/, '')
      .trim();
    return stripped;
  }
  return trimmed;
}

function tryParseValidate<T>(
  text: string,
  schema: ZodType<T>
): { ok: true; value: T } | { ok: false; error: string } {
  let parsed: unknown;
  try {
    parsed = JSON.parse(unwrapFences(text));
  } catch (err) {
    return {
      ok: false,
      error: `JSON.parse failed: ${err instanceof Error ? err.message : String(err)}`,
    };
  }
  const result = schema.safeParse(parsed);
  if (!result.success) {
    return {
      ok: false,
      error: `Zod validation failed: ${JSON.stringify(result.error.issues, null, 2)}`,
    };
  }
  return { ok: true, value: result.data };
}

export async function callLLMStage<T>(
  opts: LLMCallOptions<T>
): Promise<LLMCallResult<T>> {
  const client = opts.client ?? getAnthropicClient();
  const history: Array<{ role: 'user' | 'assistant'; content: string }> = [
    { role: 'user', content: opts.user },
  ];

  // Opus 4.7 rejects `temperature`, `top_p`, `top_k`. We don't pass any of
  // them, so the same body shape works on all three models.
  const firstBody = buildBody({
    model: opts.model,
    system: opts.system,
    history,
    maxTokens: opts.maxTokens,
  });
  void (firstBody satisfies Anthropic.MessageCreateParamsNonStreaming);

  const first = await client.messages.create(firstBody);
  const firstText = extractText(first);
  const firstParse = tryParseValidate(firstText, opts.schema);

  if (firstParse.ok) {
    const usage: TokenUsage = first.usage as TokenUsage;
    return {
      data: firstParse.value,
      usage,
      estimated_cost_usd: estimateCostUsd(opts.model, usage),
      attempts: 1,
      raw_response: firstText,
    };
  }

  // Retry once. Append the model's previous (invalid) response and a follow-up
  // user turn explaining the parse error.
  history.push({ role: 'assistant', content: firstText });
  history.push({
    role: 'user',
    content: [
      'Your previous response could not be parsed. Error:',
      firstParse.error,
      '',
      'Re-emit the SAME response, corrected to be valid JSON that conforms to the requested schema.',
      'Respond ONLY with JSON. No prose, no code fences.',
    ].join('\n'),
  });

  const secondBody = buildBody({
    model: opts.model,
    system: opts.system,
    history,
    maxTokens: opts.maxTokens,
  });
  void (secondBody satisfies Anthropic.MessageCreateParamsNonStreaming);

  const second = await client.messages.create(secondBody);
  const secondText = extractText(second);
  const secondParse = tryParseValidate(secondText, opts.schema);

  const totalUsage = addUsage(
    first.usage as TokenUsage,
    second.usage as TokenUsage
  );

  if (!secondParse.ok) {
    throw new Error(
      `LLM stage on ${opts.model} failed twice. First: ${firstParse.error}. Second: ${secondParse.error}. Last raw response: ${secondText.slice(0, 500)}`
    );
  }

  return {
    data: secondParse.value,
    usage: totalUsage,
    estimated_cost_usd: estimateCostUsd(opts.model, totalUsage),
    attempts: 2,
    raw_response: secondText,
  };
}

// Re-export OPUS so callers don't need a second import for the common case.
export { OPUS };