//! Comparison-matrix reporters: a test-case-by-model grid.
//!
//! promptfoo's signature output is a matrix: rows are test cases, columns are
//! models/prompts, each cell shows pass/fail (and the score) with drill-down
//! into the failing rows. We produce the matrix run data already (an agent
//! test fanned across a `models:` matrix), carried on [`Report::matrix`]; these
//! reporters render it as a self-contained HTML file ([`MatrixReporter`]) or a
//! GitHub-flavored Markdown grid ([`MatrixMarkdownReporter`]), one file, no
//! server, consistent with the file/stdout output model.
//!
//! When a report carries no matrix section (an ordinary, non-matrix run) the
//! grid degrades to a single `result` column with one row per test, so the
//! reporter is always well-defined rather than empty.
//!
//! Redaction happens upstream: the dispatch site calls [`Report::redacted`]
//! (which redacts the matrix labels and row names too) before handing the
//! report to any reporter, so these reporters render the report as given.

use std::fmt::Write as _;
use std::io::Write;

use crate::report::{Report, Reporter, TestStatus};

/// Status of one grid cell after joining its rows back to the report.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CellStatus {
    /// Every row the cell produced passed.
    Pass,
    /// At least one row the cell produced failed.
    Fail,
    /// The cell named no resolvable rows (nothing to judge).
    Missing,
}

/// One resolved grid cell: its status, best score, and the failing rows'
/// messages for drill-down.
struct Cell {
    status: CellStatus,
    score: Option<f64>,
    /// `(row name, failure message)` for each failing row in the cell.
    failures: Vec<(String, String)>,
}

/// A pivoted matrix grid ready to render: row labels (agent tests), column
/// labels (models/prompts), and a cell per `(row, column)`.
struct Grid {
    rows: Vec<String>,
    cols: Vec<String>,
    /// `cells[r][c]` is the cell for row `r`, column `c`.
    cells: Vec<Vec<Cell>>,
}

impl Grid {
    /// Build the grid from a report's matrix section, or a degraded
    /// single-column grid when the report carries no matrix.
    fn build(report: &Report) -> Grid {
        match &report.matrix {
            Some(matrix) => Grid::from_cells(report, matrix),
            None => Grid::degraded(report),
        }
    }

    /// Pivot the recorded matrix cells into the grid, joining each cell's row
    /// names back to [`Report::tests`] for status and score.
    fn from_cells(report: &Report, matrix: &crate::report::MatrixReport) -> Grid {
        let rows = unique_in_order(matrix.cells.iter().map(|c| c.agent.clone()));
        let cols = unique_in_order(matrix.cells.iter().map(|c| c.label.clone()));
        let mut cells = Vec::with_capacity(rows.len());
        for agent in &rows {
            let mut grid_row = Vec::with_capacity(cols.len());
            for label in &cols {
                let cell = matrix
                    .cells
                    .iter()
                    .find(|c| &c.agent == agent && &c.label == label);
                grid_row.push(match cell {
                    Some(c) => resolve_cell(report, &c.rows),
                    None => Cell {
                        status: CellStatus::Missing,
                        score: None,
                        failures: Vec::new(),
                    },
                });
            }
            cells.push(grid_row);
        }
        Grid { rows, cols, cells }
    }

    /// One row per test, a single `result` column. Used when there is no
    /// matrix to pivot.
    fn degraded(report: &Report) -> Grid {
        let rows: Vec<String> = report.tests.iter().map(|t| t.name.clone()).collect();
        let cols = vec!["result".to_string()];
        let cells = report
            .tests
            .iter()
            .map(|t| vec![resolve_cell(report, std::slice::from_ref(&t.name))])
            .collect();
        Grid { rows, cols, cells }
    }

    /// Pass rate down one column (per model), as a `(passed, total-judged)`
    /// pair. Missing cells do not count toward the total.
    fn col_pass_rate(&self, col: usize) -> (usize, usize) {
        let mut passed = 0;
        let mut total = 0;
        for row in &self.cells {
            match row[col].status {
                CellStatus::Pass => {
                    passed += 1;
                    total += 1;
                }
                CellStatus::Fail => total += 1,
                CellStatus::Missing => {}
            }
        }
        (passed, total)
    }

    /// Pass rate across one row (per test), as a `(passed, total-judged)` pair.
    fn row_pass_rate(&self, row: usize) -> (usize, usize) {
        let mut passed = 0;
        let mut total = 0;
        for cell in &self.cells[row] {
            match cell.status {
                CellStatus::Pass => {
                    passed += 1;
                    total += 1;
                }
                CellStatus::Fail => total += 1,
                CellStatus::Missing => {}
            }
        }
        (passed, total)
    }
}

/// Resolve a cell's status, best score, and failures by joining its row names
/// to [`Report::tests`]. The cell passes when every resolved row passed; its
/// score is the best (highest) row score, matching the matrix best-cell rule.
fn resolve_cell(report: &Report, row_names: &[String]) -> Cell {
    let mut any = false;
    let mut all_pass = true;
    let mut best: Option<f64> = None;
    let mut failures = Vec::new();
    for name in row_names {
        let Some(test) = report.tests.iter().find(|t| &t.name == name) else {
            continue;
        };
        any = true;
        if let Some(s) = test.score {
            best = Some(best.map_or(s, |b: f64| b.max(s)));
        }
        if matches!(test.status, TestStatus::Fail) {
            all_pass = false;
            let msg = test
                .failure
                .as_ref()
                .map(|f| f.message.clone())
                .unwrap_or_else(|| "failed".to_string());
            failures.push((name.clone(), msg));
        }
    }
    let status = if !any {
        CellStatus::Missing
    } else if all_pass {
        CellStatus::Pass
    } else {
        CellStatus::Fail
    };
    Cell {
        status,
        score: best,
        failures,
    }
}

/// Collect values into a vector keeping first-seen order and dropping
/// duplicates. Used so the grid axes read in declaration order.
fn unique_in_order<I: Iterator<Item = String>>(items: I) -> Vec<String> {
    let mut seen = std::collections::HashSet::new();
    let mut out = Vec::new();
    for item in items {
        if seen.insert(item.clone()) {
            out.push(item);
        }
    }
    out
}

/// Format a cell's score as a short fixed-point string, or an empty string.
fn score_str(score: Option<f64>) -> String {
    score.map(|s| format!("{s:.2}")).unwrap_or_default()
}

/// Self-contained HTML comparison-matrix reporter (see module docs).
#[derive(Default)]
pub struct MatrixReporter;

impl Reporter for MatrixReporter {
    fn name(&self) -> &'static str {
        "matrix"
    }

    fn render(&self, report: &Report, out: &mut dyn Write) -> std::io::Result<()> {
        let grid = Grid::build(report);
        out.write_all(render_html(report, &grid).as_bytes())
    }
}

/// GitHub-flavored Markdown comparison-matrix reporter (see module docs).
#[derive(Default)]
pub struct MatrixMarkdownReporter;

impl Reporter for MatrixMarkdownReporter {
    fn name(&self) -> &'static str {
        "matrix-md"
    }

    fn render(&self, report: &Report, out: &mut dyn Write) -> std::io::Result<()> {
        let grid = Grid::build(report);
        out.write_all(render_markdown(report, &grid).as_bytes())
    }
}

/// Render the grid as a self-contained HTML document with inline styling so
/// the single file opens in any browser with no assets.
fn render_html(report: &Report, grid: &Grid) -> String {
    let mut s = String::new();
    let _ = write!(
        s,
        "<!doctype html>\n<html lang=\"en\">\n<head>\n<meta charset=\"utf-8\">\n\
<title>mcptest comparison matrix</title>\n<style>{CSS}</style>\n</head>\n<body>\n"
    );
    let _ = writeln!(
        s,
        "<h1>Comparison matrix</h1>\n<p class=\"run\">run {}</p>",
        html_escape(&report.run_id)
    );
    let _ = writeln!(s, "<table>\n<thead><tr><th>test \\ model</th>");
    for col in &grid.cols {
        let _ = write!(s, "<th>{}</th>", html_escape(col));
    }
    let _ = writeln!(s, "<th>row pass</th></tr></thead>\n<tbody>");
    for (r, row) in grid.rows.iter().enumerate() {
        let _ = write!(s, "<tr><th>{}</th>", html_escape(row));
        for cell in &grid.cells[r] {
            write_html_cell(&mut s, cell);
        }
        let (p, t) = grid.row_pass_rate(r);
        let _ = write!(s, "<td class=\"rate\">{}</td>", pass_rate_str(p, t));
        let _ = writeln!(s, "</tr>");
    }
    let _ = write!(s, "<tr class=\"summary\"><th>model pass</th>");
    for c in 0..grid.cols.len() {
        let (p, t) = grid.col_pass_rate(c);
        let _ = write!(s, "<td class=\"rate\">{}</td>", pass_rate_str(p, t));
    }
    let _ = writeln!(s, "<td></td></tr>\n</tbody></table>");
    write_html_rollups(&mut s, report);
    let _ = writeln!(s, "</body>\n</html>");
    s
}

/// Write the suite-level metric rollup table beneath the grid, when any test
/// reported a derived metric.
fn write_html_rollups(s: &mut String, report: &Report) {
    let rollups = report.metric_rollups();
    if rollups.is_empty() {
        return;
    }
    let _ = writeln!(s, "<h2>Metric rollups</h2>");
    let _ = writeln!(
        s,
        "<table><thead><tr><th>Metric</th><th>Mean</th><th>Tests</th></tr></thead><tbody>"
    );
    for r in rollups {
        let _ = writeln!(
            s,
            "<tr><td>{}</td><td class=\"rate\">{:.2}</td><td class=\"rate\">{}</td></tr>",
            html_escape(&r.name),
            r.mean,
            r.count
        );
    }
    let _ = writeln!(s, "</tbody></table>");
}

/// Write one HTML grid cell: a colored pass/fail/missing box with the score
/// and, for a failing cell, a `<details>` drill-down of the failing rows.
fn write_html_cell(s: &mut String, cell: &Cell) {
    let (class, glyph) = match cell.status {
        CellStatus::Pass => ("pass", "PASS"),
        CellStatus::Fail => ("fail", "FAIL"),
        CellStatus::Missing => ("missing", "-"),
    };
    let _ = write!(s, "<td class=\"{class}\">");
    let _ = write!(s, "<span class=\"v\">{glyph}</span>");
    let score = score_str(cell.score);
    if !score.is_empty() {
        let _ = write!(s, " <span class=\"s\">{score}</span>");
    }
    if !cell.failures.is_empty() {
        let _ = write!(s, "<details><summary>why</summary><ul>");
        for (row, msg) in &cell.failures {
            let _ = write!(
                s,
                "<li><code>{}</code>: {}</li>",
                html_escape(row),
                html_escape(msg)
            );
        }
        let _ = write!(s, "</ul></details>");
    }
    let _ = write!(s, "</td>");
}

/// Render the grid as a GFM table plus per-cell drill-down blocks below it
/// (GFM table cells cannot hold block content, so failures move below).
fn render_markdown(report: &Report, grid: &Grid) -> String {
    let mut s = String::new();
    let _ = writeln!(s, "## Comparison matrix");
    let _ = writeln!(s);
    let _ = writeln!(s, "run {}", md_escape(&report.run_id));
    let _ = writeln!(s);
    // Header: test \ model | col... | row pass
    let _ = write!(s, "| test \\ model |");
    for col in &grid.cols {
        let _ = write!(s, " {} |", md_escape(col));
    }
    let _ = writeln!(s, " row pass |");
    let _ = write!(s, "| --- |");
    for _ in &grid.cols {
        let _ = write!(s, " --- |");
    }
    let _ = writeln!(s, " --- |");
    for (r, row) in grid.rows.iter().enumerate() {
        let _ = write!(s, "| {} |", md_escape(row));
        for cell in &grid.cells[r] {
            let _ = write!(s, " {} |", md_cell(cell));
        }
        let (p, t) = grid.row_pass_rate(r);
        let _ = writeln!(s, " {} |", pass_rate_str(p, t));
    }
    // Summary row: per-model pass rate.
    let _ = write!(s, "| **model pass** |");
    for c in 0..grid.cols.len() {
        let (p, t) = grid.col_pass_rate(c);
        let _ = write!(s, " {} |", pass_rate_str(p, t));
    }
    let _ = writeln!(s, "  |");
    render_markdown_rollups(&mut s, report);
    render_markdown_failures(&mut s, grid);
    s
}

/// Append the suite-level metric rollup table beneath the grid, when any test
/// reported a derived metric.
fn render_markdown_rollups(s: &mut String, report: &Report) {
    let rollups = report.metric_rollups();
    if rollups.is_empty() {
        return;
    }
    let _ = writeln!(s, "\n### Metric rollups\n");
    let _ = writeln!(s, "| Metric | Mean | Tests |");
    let _ = writeln!(s, "| --- | --- | --- |");
    for r in rollups {
        let _ = writeln!(
            s,
            "| {} | {:.2} | {} |",
            md_escape(&r.name),
            r.mean,
            r.count
        );
    }
}

/// One Markdown grid cell: a glyph plus the score, no block content.
fn md_cell(cell: &Cell) -> String {
    let glyph = match cell.status {
        CellStatus::Pass => "PASS",
        CellStatus::Fail => "FAIL",
        CellStatus::Missing => "-",
    };
    let score = score_str(cell.score);
    if score.is_empty() {
        glyph.to_string()
    } else {
        format!("{glyph} {score}")
    }
}

/// Emit a `<details>` drill-down per failing cell beneath the table.
fn render_markdown_failures(s: &mut String, grid: &Grid) {
    let mut wrote_header = false;
    for (r, row) in grid.rows.iter().enumerate() {
        for (c, cell) in grid.cells[r].iter().enumerate() {
            if cell.failures.is_empty() {
                continue;
            }
            if !wrote_header {
                let _ = writeln!(s, "\n### Failures");
                wrote_header = true;
            }
            let _ = writeln!(
                s,
                "\n<details><summary>{} @ {}</summary>\n",
                md_escape(row),
                md_escape(&grid.cols[c])
            );
            for (name, msg) in &cell.failures {
                let _ = writeln!(s, "- `{}`: {}", md_escape(name), md_escape(msg));
            }
            let _ = writeln!(s, "\n</details>");
        }
    }
}

/// Format a `(passed, total)` pair as `p/t (pct%)`, or a dash when nothing was
/// judged in that row or column.
fn pass_rate_str(passed: usize, total: usize) -> String {
    if total == 0 {
        return "-".to_string();
    }
    let pct = (passed as f64 / total as f64) * 100.0;
    format!("{passed}/{total} ({pct:.0}%)")
}

/// Minimal HTML escaping for text rendered into the document.
fn html_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for ch in s.chars() {
        match ch {
            '&' => out.push_str("&amp;"),
            '<' => out.push_str("&lt;"),
            '>' => out.push_str("&gt;"),
            '"' => out.push_str("&quot;"),
            '\'' => out.push_str("&#39;"),
            _ => out.push(ch),
        }
    }
    out
}

/// Escape the GFM inline characters that would break a table cell.
fn md_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for ch in s.chars() {
        match ch {
            '|' | '<' | '>' | '*' | '_' => {
                out.push('\\');
                out.push(ch);
            }
            _ => out.push(ch),
        }
    }
    out
}

/// Inline stylesheet for the HTML matrix: colored cells, readable grid.
const CSS: &str = "body{font-family:system-ui,sans-serif;margin:2rem;color:#1a1a1a}\
h1{font-size:1.3rem}.run{color:#666;font-size:.9rem}\
table{border-collapse:collapse}th,td{border:1px solid #ddd;padding:.4rem .6rem;text-align:center;vertical-align:top}\
thead th{background:#f4f4f4}td.pass{background:#e6f4ea}td.fail{background:#fce8e6}td.missing{background:#f4f4f4;color:#999}\
.v{font-weight:600;font-size:.8rem}.s{color:#444}.rate{font-variant-numeric:tabular-nums}\
tr.summary th,tr.summary td{background:#f4f4f4;font-weight:600}\
details{margin-top:.3rem;text-align:left}summary{cursor:pointer;color:#1a73e8;font-size:.8rem}\
details ul{margin:.3rem 0;padding-left:1rem;font-size:.8rem}code{background:#eee;padding:0 .2rem;border-radius:3px}";

#[cfg(test)]
mod tests {
    use super::*;
    use crate::report::Report;

    #[test]
    fn html_renders_grid_with_models_and_pass_fail() {
        let report = Report::sample_matrix();
        let mut buf = Vec::new();
        MatrixReporter.render(&report, &mut buf).unwrap();
        let s = String::from_utf8(buf).unwrap();
        assert!(s.contains("Comparison matrix"));
        assert!(s.contains("gpt-4o"));
        assert!(s.contains("claude"));
        assert!(s.contains("class=\"pass\""));
        assert!(s.contains("class=\"fail\""));
        // Drill-down on the failing cell.
        assert!(s.contains("<details>"));
    }

    #[test]
    fn markdown_renders_grid_table_and_failures() {
        let report = Report::sample_matrix();
        let mut buf = Vec::new();
        MatrixMarkdownReporter.render(&report, &mut buf).unwrap();
        let s = String::from_utf8(buf).unwrap();
        assert!(s.contains("## Comparison matrix"));
        assert!(s.contains("| test \\ model |"));
        assert!(s.contains("PASS"));
        assert!(s.contains("FAIL"));
        assert!(s.contains("### Failures"));
        assert!(s.contains("model pass"));
    }

    #[test]
    fn degraded_grid_when_no_matrix_section() {
        let mut report = Report::sample();
        report.matrix = None;
        let mut buf = Vec::new();
        MatrixReporter.render(&report, &mut buf).unwrap();
        let s = String::from_utf8(buf).unwrap();
        // One `result` column, a row per test.
        assert!(s.contains("result"));
        assert!(s.contains("test \\ model"));
    }

    #[test]
    fn col_and_row_pass_rates_count_only_judged_cells() {
        let report = Report::sample_matrix();
        let grid = Grid::build(&report);
        // The sample has two models; the column pass rate denominators reflect
        // judged (non-missing) cells only.
        let (_, t0) = grid.col_pass_rate(0);
        assert!(t0 > 0);
    }

    #[test]
    fn build_uses_matrix_section_axes() {
        let report = Report::sample_matrix();
        let grid = Grid::build(&report);
        assert_eq!(grid.cols.len(), 2, "two model columns");
        assert!(!grid.rows.is_empty());
    }

    #[test]
    fn unique_in_order_dedupes_and_preserves_order() {
        let v = unique_in_order(["b", "a", "b", "c", "a"].into_iter().map(str::to_string));
        assert_eq!(v, vec!["b", "a", "c"]);
    }
}
