import { readFile } from 'node:fs/promises';
import { extname } from 'node:path';
import { fetchGeminiWithPool } from './gemini-key-pool.js';

const DEFAULT_GEMINI_JUDGE_ENDPOINT =
  'https://generativelanguage.googleapis.com/v1beta/models/gemini-3.5-flash:generateContent';

export interface JudgeCandidate {
  id: string;
  prompt?: string;
  imagePath?: string;
}

export interface JudgeSceneInput {
  sceneIndex: number;
  intendedPrompt?: string;
  candidates: JudgeCandidate[];
}

export interface JudgeResult {
  sceneIndex: number;
  bestCandidateId: string;
  reasoning: string;
}

type InlineDataPart = { inlineData: { mimeType: string; data: string } };
type TextPart = { text: string };
type GeminiPart = InlineDataPart | TextPart;

/**
 * Maximum total inline-image payload (sum of base64 byte lengths) we will send
 * to Gemini in a single judge request. Gemini's hard per-request ceiling is
 * ~20 MB; we stay comfortably under it so the surrounding JSON/prompt text and
 * transport overhead never push the request over the edge. Configurable via
 * VCLAW_GEMINI_JUDGE_MAX_INLINE_BYTES (bytes). When the budget is exhausted,
 * further images are dropped (with a per-image stderr warning) rather than
 * silently producing an oversized request that the API rejects.
 */
const DEFAULT_MAX_INLINE_BYTES = 15 * 1024 * 1024;

function maxInlineBytes(): number {
  const raw = process.env.VCLAW_GEMINI_JUDGE_MAX_INLINE_BYTES;
  if (raw === undefined) return DEFAULT_MAX_INLINE_BYTES;
  const parsed = Number(raw);
  return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_MAX_INLINE_BYTES;
}

function inlinePartByteLength(part: InlineDataPart): number {
  return part.inlineData.data.length;
}

function imageMimeType(imagePath: string): string {
  switch (extname(imagePath).toLowerCase()) {
    case '.jpg':
    case '.jpeg':
      return 'image/jpeg';
    case '.webp':
      return 'image/webp';
    case '.png':
    default:
      return 'image/png';
  }
}

// Mirrors parseGeminiTextResponse in gemini-analyze.ts: unwrap
// candidates[0].content.parts[].text. Returns '' rather than throwing so the
// caller can omit an unparseable scene instead of failing the whole batch.
function extractGeminiText(payload: unknown): string {
  const candidates = (
    payload as { candidates?: Array<{ content?: { parts?: Array<{ text?: string }> } }> }
  ).candidates;
  return (
    candidates?.[0]?.content?.parts
      ?.map((part) => part.text ?? '')
      .join('\n')
      .trim() ?? ''
  );
}

function stripJsonFences(text: string): string {
  return text
    .replace(/^```json\s*/i, '')
    .replace(/^```\s*/i, '')
    .replace(/\s*```$/i, '')
    .trim();
}

async function readInlineImagePart(imagePath: string): Promise<InlineDataPart | null> {
  try {
    const bytes = await readFile(imagePath);
    return {
      inlineData: {
        mimeType: imageMimeType(imagePath),
        data: bytes.toString('base64'),
      },
    };
  } catch {
    // Missing/unreadable file: skip the attachment, never crash.
    return null;
  }
}

function buildJudgePromptText(scene: JudgeSceneInput): string {
  const lines: string[] = [
    'You are a film director selecting the single best candidate render for one scene of a video.',
    '',
  ];
  if (scene.intendedPrompt) {
    lines.push(`Intended scene prompt:\n${scene.intendedPrompt}`, '');
  }
  lines.push('Candidates (each may have an attached image below, in order):');
  for (const candidate of scene.candidates) {
    lines.push(
      `- id "${candidate.id}"${candidate.prompt ? `: ${candidate.prompt}` : ''}`,
    );
  }
  lines.push(
    '',
    'Pick the candidate that best matches the intended scene prompt with the highest',
    'visual quality, character/style consistency, and absence of artifacts.',
    'Return ONLY valid JSON of the exact shape:',
    '{ "bestCandidateId": "<one of the candidate ids above>", "reasoning": "<short explanation>" }',
    'The bestCandidateId MUST be exactly one of the listed candidate ids.',
    'Do not wrap the JSON in markdown fences.',
  );
  return lines.join('\n');
}

async function judgeOneScene(
  scene: JudgeSceneInput,
  referenceParts: InlineDataPart[],
  endpoint: string,
  fetcher?: typeof fetch,
): Promise<JudgeResult | null> {
  const knownIds = new Set(scene.candidates.map((candidate) => candidate.id));
  const parts: GeminiPart[] = [];

  // Pre-flight inline-image budget. Each base64-encoded image is large; with
  // many candidates plus reference images a single request can blow past
  // Gemini's ~20 MB ceiling, which surfaces as a non-OK response that silently
  // drops the scene. We instead drop images once the budget is exhausted and
  // log each skip, so the failure mode is observable and the scene can still be
  // judged on whatever images (and the prompt text) fit.
  const budget = maxInlineBytes();
  let usedBytes = 0;

  if (referenceParts.length > 0) {
    const fittedRefs: InlineDataPart[] = [];
    let droppedRefs = 0;
    for (const part of referenceParts) {
      const size = inlinePartByteLength(part);
      if (usedBytes + size > budget) {
        droppedRefs += 1;
        continue;
      }
      usedBytes += size;
      fittedRefs.push(part);
    }
    if (droppedRefs > 0) {
      process.stderr.write(
        `[judge/gemini] scene=${scene.sceneIndex} dropped ${droppedRefs} reference image(s) to stay under the ${budget}-byte inline payload budget\n`,
      );
    }
    if (fittedRefs.length > 0) {
      parts.push({ text: 'Reference images (the look/identity to match):' });
      parts.push(...fittedRefs);
    }
  }

  // Attach each candidate's image (if present and readable), labeled so the
  // model can align image -> id. Candidates with no image are judged by prompt
  // text alone. Images beyond the remaining byte budget are skipped (logged),
  // so an oversized scene degrades gracefully instead of failing the request.
  for (const candidate of scene.candidates) {
    if (!candidate.imagePath) continue;
    const part = await readInlineImagePart(candidate.imagePath);
    if (!part) continue;
    const size = inlinePartByteLength(part);
    if (usedBytes + size > budget) {
      process.stderr.write(
        `[judge/gemini] scene=${scene.sceneIndex} candidate "${candidate.id}" image skipped: exceeds the ${budget}-byte inline payload budget; judging it by prompt text only\n`,
      );
      continue;
    }
    usedBytes += size;
    parts.push({ text: `Candidate "${candidate.id}" image:` });
    parts.push(part);
  }

  parts.push({ text: buildJudgePromptText(scene) });

  let response: Response;
  try {
    response = await fetchGeminiWithPool(
      (key) => `${endpoint}${endpoint.includes('?') ? '&' : '?'}key=${encodeURIComponent(key)}`,
      {
        method: 'POST',
        headers: { 'Content-Type': 'application/json', Connection: 'close' },
        body: JSON.stringify({
          contents: [{ parts }],
          generationConfig: {
            temperature: 0.2,
            maxOutputTokens: 600,
            responseMimeType: 'application/json',
          },
        }),
      },
      {
        fetcher,
        onRetry: (label, status) => {
          process.stderr.write(
            `[judge/gemini] scene=${scene.sceneIndex} ${label} returned HTTP ${status}; rotating key\n`,
          );
        },
      },
    );
  } catch (error) {
    process.stderr.write(
      `[judge/gemini] scene=${scene.sceneIndex} request error: ${
        error instanceof Error ? error.message : String(error)
      }; leaving to human\n`,
    );
    return null;
  }

  if (!response.ok) {
    process.stderr.write(
      `[judge/gemini] scene=${scene.sceneIndex} HTTP ${response.status}; leaving to human\n`,
    );
    return null;
  }

  let payload: unknown;
  try {
    payload = await response.json();
  } catch {
    return null;
  }

  const text = extractGeminiText(payload);
  if (!text) return null;

  let parsed: { bestCandidateId?: unknown; reasoning?: unknown };
  try {
    parsed = JSON.parse(stripJsonFences(text)) as typeof parsed;
  } catch {
    process.stderr.write(
      `[judge/gemini] scene=${scene.sceneIndex} unparseable JSON; leaving to human\n`,
    );
    return null;
  }

  const bestCandidateId =
    typeof parsed.bestCandidateId === 'string' ? parsed.bestCandidateId : '';
  if (!bestCandidateId || !knownIds.has(bestCandidateId)) {
    process.stderr.write(
      `[judge/gemini] scene=${scene.sceneIndex} returned unknown candidate id "${bestCandidateId}"; leaving to human\n`,
    );
    return null;
  }

  const reasoning =
    typeof parsed.reasoning === 'string' ? parsed.reasoning : '';
  return { sceneIndex: scene.sceneIndex, bestCandidateId, reasoning };
}

/**
 * LLM-as-judge auto-select. For each scene, asks Gemini to pick the best
 * candidate, conditioned on the intended prompt, each candidate's image (when
 * available on disk), and any shared reference images.
 *
 * Defensive by design: a scene whose response can't be parsed, returns an
 * unknown candidate id, or errors at the network layer is OMITTED from the
 * result (never throws the whole batch) — the caller treats omitted scenes as
 * "leave to human". The network is only ever touched through the injected
 * `fetcher` (via fetchGeminiWithPool), so this is unit-testable without keys.
 */
export async function judgeBestCandidates(input: {
  scenes: JudgeSceneInput[];
  referenceImagePaths?: string[];
  endpoint?: string;
  fetcher?: typeof fetch;
}): Promise<JudgeResult[]> {
  const endpoint =
    input.endpoint ?? process.env.VCLAW_GEMINI_API_ENDPOINT ?? DEFAULT_GEMINI_JUDGE_ENDPOINT;

  const referenceParts: InlineDataPart[] = [];
  for (const refPath of input.referenceImagePaths ?? []) {
    const part = await readInlineImagePart(refPath);
    if (part) referenceParts.push(part);
  }

  const results: JudgeResult[] = [];
  for (const scene of input.scenes) {
    if (scene.candidates.length === 0) continue;
    const result = await judgeOneScene(scene, referenceParts, endpoint, input.fetcher);
    if (result) results.push(result);
  }
  return results;
}