import type { ScenarioName, ScenarioSummary, PerTaskRow, ModelReport, Aggregate, AggregateScenario, AggregateModel } from "@duck/types";
import { SCENARIO_NAMES } from "@duck/types";

function round(x: number, d = 2): number {
  if (!Number.isFinite(x)) return x;
  const f = 10 ** d;
  return Math.round(x * f) / f;
}

function avg(arr: number[]): number {
  if (!arr.length) return 0;
  return arr.reduce((a, b) => a + b, 0) / arr.length;
}

export interface BuildReportOpts {
  model: string;
  mcpUrl: string;
  rows: PerTaskRow[];
  meta?: Record<string, unknown>;
  prompts?: Partial<Record<ScenarioName, string>>;
}

export function buildReport({ model, mcpUrl, rows, meta = {}, prompts = {} }: BuildReportOpts): ModelReport {
  const scenarios = SCENARIO_NAMES;
  const byScenario: Record<string, { total: number; correct: number; reviewed: number; duckUsed: number; toolCalls: number; promptTokens: number[]; genTokens: number[]; duckTokens: number[]; rows: PerTaskRow[] }> = {};
  for (const s of scenarios) byScenario[s] = { total: 0, correct: 0, reviewed: 0, duckUsed: 0, toolCalls: 0, promptTokens: [], genTokens: [], duckTokens: [], rows: [] };

  for (const row of rows) {
    const s = row.scenario;
    if (!byScenario[s]) continue;
    byScenario[s].total += 1;
    if (row.correct !== null && row.correct !== undefined) {
      byScenario[s].reviewed += 1;
      if (row.correct) byScenario[s].correct += 1;
    }
    if (row.duckUsed) byScenario[s].duckUsed += 1;
    byScenario[s].toolCalls += row.toolCalls ?? 0;
    byScenario[s].promptTokens.push(row.promptTokens ?? 0);
    byScenario[s].genTokens.push(row.genTokens ?? 0);
    byScenario[s].duckTokens.push(row.duckTokens ?? 0);
    byScenario[s].rows.push(row);
  }

  const summaries: Record<ScenarioName, ScenarioSummary> = {} as Record<ScenarioName, ScenarioSummary>;
  for (const s of scenarios) {
    const g = byScenario[s];
    summaries[s] = {
      tasks: g.total,
      reviewed: g.reviewed,
      correct: g.correct,
      accuracy: g.reviewed ? round((g.correct / g.reviewed) * 100, 1) : 0,
      pending: g.total - g.reviewed,
      duckUsed: g.duckUsed,
      toolCalls: g.toolCalls,
      avgCalls: g.total ? round(g.toolCalls / g.total, 2) : 0,
      avgPromptTokens: round(avg(g.promptTokens)),
      avgGenTokens: round(avg(g.genTokens)),
      avgDuckTokens: round(avg(g.duckTokens)),
    };
  }

  return {
    generatedAt: new Date().toISOString(),
    model,
    mcpUrl,
    scenarios: summaries,
    prompts: prompts as Record<ScenarioName, string>,
    perTask: rows.map((r) => ({
      id: r.id,
      scenario: r.scenario,
      correct: r.correct,
      duckUsed: r.duckUsed,
      toolCalls: r.toolCalls,
      promptTokens: r.promptTokens,
      genTokens: r.genTokens,
      duckTokens: r.duckTokens,
      response: r.response,
      expected: r.expected,
      question: r.question,
    })),
    meta,
  };
}

export function buildAggregate(modelsReport: Record<string, ModelReport>): Aggregate {
  const perScenario: Record<ScenarioName, AggregateScenario> = {} as Record<ScenarioName, AggregateScenario>;
  for (const s of SCENARIO_NAMES) {
    let total = 0, reviewed = 0, correct = 0, duckUsed = 0, toolCalls = 0;
    const byModel: Record<string, AggregateScenario["byModel"][string]> = {};
    for (const [name, rep] of Object.entries(modelsReport)) {
      const g = rep.scenarios[s];
      if (!g) continue;
      total += g.tasks;
      reviewed += g.reviewed;
      correct += g.correct;
      duckUsed += g.duckUsed;
      toolCalls += g.toolCalls;
      byModel[name] = {
        tasks: g.tasks,
        reviewed: g.reviewed,
        correct: g.correct,
        accuracy: g.reviewed ? g.accuracy : 0,
        duckUsed: g.duckUsed,
      };
    }
    perScenario[s] = {
      tasks: total,
      reviewed,
      correct,
      accuracy: reviewed ? round((correct / reviewed) * 100, 1) : 0,
      pending: total - reviewed,
      duckUsed,
      toolCalls,
      avgCalls: total ? round(toolCalls / total, 2) : 0,
      byModel,
    };
  }

  const perModel: Record<string, AggregateModel> = {};
  for (const [name, rep] of Object.entries(modelsReport)) {
    let duckUsed = 0, toolCalls = 0;
    const gen: number[] = [];
    let reviewed = 0, correct = 0;
    for (const s of SCENARIO_NAMES) {
      const g = rep.scenarios[s];
      duckUsed += g.duckUsed;
      toolCalls += g.toolCalls;
      reviewed += g.reviewed;
      correct += g.correct;
      for (const row of rep.perTask.filter((r) => r.scenario === s)) {
        gen.push(row.genTokens ?? 0);
      }
    }
    const taskCount = rep.scenarios.control?.tasks ?? 0;
    const totalRows = SCENARIO_NAMES.reduce((acc, s) => acc + (rep.scenarios[s].tasks || 0), 0);
    perModel[name] = {
      tasks: taskCount,
      totalRows,
      reviewed,
      correct,
      accuracy: reviewed ? round((correct / reviewed) * 100, 1) : 0,
      duckUsed,
      toolCalls,
      avgGenTokens: round(avg(gen)),
    };
  }

  return { perScenario, perModel };
}