#!/usr/bin/env python3
"""Convert result JSON files to CSV and validate data completeness.

Based on load_dicts_to_df from:
https://github.com/erikerlandson/paired-comparison-ranking/blob/main/nb/paired-comparison-ranking.ipynb
"""

import ast
import json
import sys
from pathlib import Path
from typing import Any

import pandas as pd


REQUIRED_METRICS = [
    "metrics.n_tasks",
    "metrics.n_errors",
    "metrics.score",
    "metrics.cost_usd",
    "metrics.n_input_tokens",
    "metrics.n_output_tokens",
    "metrics.agent_time_seconds",
    "metrics.total_time_seconds",
]


def load_dicts_to_df(
    directory: str | Path,
    pattern: str = "*.json",
) -> pd.DataFrame:
    directory = Path(directory)
    if not directory.is_dir():
        raise NotADirectoryError(directory)

    rows: list[dict[str, Any]] = []
    for path in sorted(directory.glob(pattern)):
        if not path.is_file() or path.name.startswith("."):
            continue
        suffix = path.suffix.lower()
        if suffix == ".json":
            with path.open(encoding="utf-8") as f:
                data = json.load(f)
        else:
            text = path.read_text(encoding="utf-8")
            data = ast.literal_eval(text)

        if not isinstance(data, dict):
            raise TypeError(f"{path} does not contain a dict (got {type(data).__name__})")
        rows.append(data)

    return pd.json_normalize(rows)


def validate(df: pd.DataFrame) -> bool:
    issues = []
    for _, row in df.iterrows():
        label = f"{row.get('benchmark.name', '?')} / {row.get('model.name', '?')} / {row.get('harness.name', '?')}"
        missing = [col for col in REQUIRED_METRICS if col not in df.columns or pd.isna(row.get(col))]
        if missing:
            fields = ", ".join(c.replace("metrics.", "") for c in missing)
            issues.append(f"  {label}: missing {fields}")

    if issues:
        print(f"Validation: {len(issues)} result(s) with missing data:")
        for issue in issues:
            print(issue)
        return False

    print("Validation: all results complete")
    return True


def main():
    results_dir = Path(__file__).parent.parent / "results"
    output_path = Path(__file__).parent.parent / "results.csv"

    if len(sys.argv) > 1:
        results_dir = Path(sys.argv[1])
    if len(sys.argv) > 2:
        output_path = Path(sys.argv[2])

    df = load_dicts_to_df(results_dir)

    drop_cols = [c for c in df.columns if c.startswith("environment.config.")]
    df = df.drop(columns=drop_cols)

    validate(df)

    df.to_csv(output_path, index=False)
    print(f"Wrote {len(df)} rows to {output_path}")


if __name__ == "__main__":
    main()