import type { AxisScore, BenchTask, RawJobResult } from './types.js';
import { logger } from '../logger.js';

export interface JudgeConfig {
  endpoint: string;       // OpenAI-compat /chat/completions base, e.g. http://localhost:11434/v1
  model: string;
  apiKey?: string;
  timeoutMs?: number;
}

interface JudgeRubricResult {
  name: string;
  score: number;
  max: number;
  rationale: string;
}

const SYSTEM_PROMPT = [
  'You are an evaluation assistant for an autonomous agent benchmark.',
  'Score each rubric item with an integer in [0, max_score].',
  'Be strict but fair. Score 0 means the rubric goal is not met at all; max_score means fully met.',
  'Return ONLY a JSON object of the form {"results":[{"name":"...","score":N,"rationale":"..."}, ...]}.',
  'No prose outside the JSON. No code fences.',
].join('\n');

function buildJudgePrompt(task: BenchTask, raw: RawJobResult): string {
  const rubrics = task.grading?.llm_judge?.rubrics ?? [];
  const outputView = Object.entries(raw.outputFiles)
    .slice(0, 5)  // bound size; we don't expect many files
    .map(([name, body]) => `--- ${name} ---\n${body.slice(0, 4000)}`)
    .join('\n\n');

  return [
    `## Original task prompt`,
    task.prompt,
    '',
    `## Job status`,
    `status=${raw.status} duration_ms=${raw.durationMs}`,
    '',
    `## Output files`,
    outputView || '(no output files)',
    '',
    '## Rubrics',
    ...rubrics.map((r, i) => `${i + 1}. name="${r.name}" max_score=${r.max_score ?? 10}\n   ${r.prompt}`),
    '',
    `## Output`,
    `Return JSON with one entry per rubric, in the same order.`,
  ].join('\n');
}

async function callJudge(config: JudgeConfig, userPrompt: string): Promise<string> {
  const body = {
    model: config.model,
    stream: false,
    messages: [
      { role: 'system', content: SYSTEM_PROMPT },
      { role: 'user', content: userPrompt },
    ],
  };
  const controller = new AbortController();
  const timeout = setTimeout(() => controller.abort(), config.timeoutMs ?? 120_000);
  try {
    const headers: Record<string, string> = { 'Content-Type': 'application/json' };
    if (config.apiKey) headers['Authorization'] = `Bearer ${config.apiKey}`;
    const res = await fetch(`${config.endpoint.replace(/\/+$/, '')}/chat/completions`, {
      method: 'POST',
      headers,
      body: JSON.stringify(body),
      signal: controller.signal,
    });
    if (!res.ok) {
      throw new Error(`judge endpoint ${res.status}: ${(await res.text()).slice(0, 400)}`);
    }
    const data = (await res.json()) as { choices?: Array<{ message?: { content?: string } }> };
    return data.choices?.[0]?.message?.content ?? '';
  } finally {
    clearTimeout(timeout);
  }
}

function parseJudgeJson(content: string): JudgeRubricResult[] {
  const trimmed = content.trim().replace(/^```(?:json)?\s*/i, '').replace(/\s*```$/i, '');
  // Some models still embed prose; extract the first {...} block.
  const start = trimmed.indexOf('{');
  const end = trimmed.lastIndexOf('}');
  if (start < 0 || end <= start) throw new Error(`judge response not JSON: ${content.slice(0, 200)}`);
  const parsed = JSON.parse(trimmed.slice(start, end + 1)) as { results?: unknown };
  if (!Array.isArray(parsed.results)) throw new Error('judge response missing results[]');
  return parsed.results.map((r: unknown) => {
    const o = r as Record<string, unknown>;
    return {
      name: String(o['name'] ?? ''),
      score: Number(o['score'] ?? 0),
      max: Number(o['max'] ?? o['max_score'] ?? 10),
      rationale: String(o['rationale'] ?? ''),
    };
  });
}

/**
 * Run the LLM judge for axis D (reasoning). When the task has no llm_judge
 * config or `config` is null, returns a 1.0 score so that absence of judge
 * does not penalize.
 */
export async function gradeReasoning(
  task: BenchTask,
  raw: RawJobResult,
  config: JudgeConfig | null,
): Promise<AxisScore> {
  const rubrics = task.grading?.llm_judge?.rubrics ?? [];
  if (!config || rubrics.length === 0) {
    return { score: 1, details: ['(LLM judge skipped — no config or no rubrics)'] };
  }
  const prompt = buildJudgePrompt(task, raw);
  let raw_response: string;
  try {
    raw_response = await callJudge(config, prompt);
  } catch (err) {
    logger.warn(`[bench/judge] call failed: ${(err as Error).message}`);
    return { score: 0, details: [`judge error: ${(err as Error).message}`] };
  }
  let results: JudgeRubricResult[];
  try {
    results = parseJudgeJson(raw_response);
  } catch (err) {
    logger.warn(`[bench/judge] parse failed: ${(err as Error).message}`);
    return { score: 0, details: [`judge parse error: ${(err as Error).message}`, `raw=${raw_response.slice(0, 200)}`] };
  }
  if (results.length === 0) {
    return { score: 0, details: ['judge returned no results'] };
  }
  const normalized = results.map((r) => Math.max(0, Math.min(1, r.score / Math.max(1, r.max))));
  const score = normalized.reduce((a, b) => a + b, 0) / normalized.length;
  const details = results.map((r) => `${r.name}: ${r.score}/${r.max} — ${r.rationale}`);
  return { score, details };
}

export function loadJudgeConfigFromEnv(fallback: { endpoint: string; model: string; apiKey?: string }): JudgeConfig | null {
  const enabled = process.env['BENCH_JUDGE'] !== 'off';
  if (!enabled) return null;
  return {
    endpoint: process.env['BENCH_JUDGE_ENDPOINT'] ?? fallback.endpoint,
    model: process.env['BENCH_JUDGE_MODEL'] ?? fallback.model,
    apiKey: process.env['BENCH_JUDGE_API_KEY'] ?? fallback.apiKey,
    timeoutMs: 120_000,
  };
}
