from __future__ import annotations

import argparse
import glob
import json
from collections import defaultdict
from pathlib import Path
from typing import Any

EVAL_RECORD_SCHEMA_VERSION = 2
REQUIRED_TOP_LEVEL_FIELDS = (
    "schema_version",
    "run_id",
    "case_id",
    "condition",
    "attempt",
    "success",
    "category",
    "difficulty",
    "snapshot_dependency",
    "run_metadata",
    "attempt_metadata",
    "model_metadata",
)
REQUIRED_RUN_METADATA_FIELDS = (
    "run_id",
    "started_at_utc",
    "deterministic",
    "seed",
    "python_version",
    "platform",
    "llmdebug_version",
    "patcher_name",
    "patcher_config_fingerprint",
)
REQUIRED_ATTEMPT_METADATA_FIELDS = (
    "attempt",
    "prompt_sha256",
    "prompt_path",
    "evidence_sha256",
    "evidence_path",
)
REQUIRED_MODEL_METADATA_FIELDS = ("backend", "name")
REQUIRED_OPENAI_MODEL_METADATA_FIELDS = (
    "base_url",
    "model",
    "temperature",
    "max_tokens",
    "response_mode",
    "api_compat_seed_requested",
    "api_compat_seed_acknowledged",
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Analyze llmdebug eval JSONL results.")
    parser.add_argument("paths", nargs="+", help="One or more JSONL result files (globs ok).")
    parser.add_argument(
        "--out",
        default="",
        help="Optional path to write machine-readable summary JSON.",
    )
    parser.add_argument(
        "--strict-input",
        action="store_true",
        help="Fail on first unreadable input file or malformed JSONL record.",
    )
    parser.add_argument(
        "--stats",
        action="store_true",
        help="Run statistical significance tests (McNemar, Cohen's h, bootstrap CI).",
    )
    parser.add_argument(
        "--aggregation",
        choices=["per_run", "majority_vote", "any_success"],
        default="per_run",
        help="Aggregation strategy for multi-run (k>1) analysis.",
    )
    args = parser.parse_args(argv)

    try:
        grouped, input_errors = _load_and_group_records(args.paths, strict_input=args.strict_input)
    except ValueError as exc:
        print(f"Input error: {exc}")
        return 2
    if not grouped:
        _print_input_errors(input_errors)
        print("No records found.")
        return 1

    rows = _derive_case_condition_rows(grouped)
    if not rows:
        print("No valid case-condition records found.")
        return 1

    by_category: dict[str, list[dict[str, Any]]] = defaultdict(list)
    by_difficulty: dict[str, list[dict[str, Any]]] = defaultdict(list)
    by_snapshot_dependency: dict[str, list[dict[str, Any]]] = defaultdict(list)
    by_bug_source: dict[str, list[dict[str, Any]]] = defaultdict(list)
    by_contamination_risk: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        by_category[row["category"]].append(row)
        by_difficulty[row["difficulty"]].append(row)
        by_snapshot_dependency[row["snapshot_dependency"]].append(row)
        by_bug_source[row.get("bug_source") or "hand_crafted"].append(row)
        by_contamination_risk[row.get("contamination_risk") or "unknown"].append(row)

    summary = {
        "overall": _summarize_bucket(rows),
        "by_category": {
            key: _summarize_bucket(bucket_rows) for key, bucket_rows in sorted(by_category.items())
        },
        "by_difficulty": {
            key: _summarize_bucket(bucket_rows)
            for key, bucket_rows in sorted(by_difficulty.items())
        },
        "by_snapshot_dependency": {
            key: _summarize_bucket(bucket_rows)
            for key, bucket_rows in sorted(by_snapshot_dependency.items())
        },
        "by_bug_source": {
            key: _summarize_bucket(bucket_rows)
            for key, bucket_rows in sorted(by_bug_source.items())
        },
        "by_contamination_risk": {
            key: _summarize_bucket(bucket_rows)
            for key, bucket_rows in sorted(by_contamination_risk.items())
        },
        "input_errors": input_errors,
    }

    _print_summary(summary)

    if args.stats:
        stats_result = _run_statistical_analysis(rows, aggregation=args.aggregation)
        summary["statistics"] = stats_result
        _print_stats(stats_result)

    if args.out:
        out_path = Path(args.out)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8")
        print(f"\nWrote summary JSON: {out_path}")

    return 0


def _load_and_group_records(
    paths: list[str],
    *,
    strict_input: bool = False,
) -> tuple[dict[tuple[str, str, str], list[dict[str, Any]]], dict[str, int]]:
    grouped: dict[tuple[str, str, str], list[dict[str, Any]]] = defaultdict(list)
    input_errors = {
        "files_unreadable": 0,
        "json_lines_invalid": 0,
        "records_skipped": 0,
    }
    for src in _expand_paths(paths):
        src_path = Path(src)
        try:
            handle = src_path.open(encoding="utf-8")
        except OSError as exc:
            input_errors["files_unreadable"] += 1
            if strict_input:
                raise ValueError(f"{src}: unreadable file ({type(exc).__name__}: {exc})") from exc
            continue
        with handle:
            for line_no, line in enumerate(handle, start=1):
                if not line.strip():
                    continue
                try:
                    record = json.loads(line)
                except json.JSONDecodeError as exc:
                    input_errors["json_lines_invalid"] += 1
                    if strict_input:
                        raise ValueError(f"{src}:{line_no}: invalid JSON ({exc})") from exc
                    continue
                if not isinstance(record, dict):
                    input_errors["records_skipped"] += 1
                    if strict_input:
                        raise ValueError(f"{src}:{line_no}: record is not a JSON object")
                    continue
                run_id = str(record.get("run_id") or src)
                case_id = str(record.get("case_id") or "")
                condition = str(record.get("condition") or "")
                if not case_id or not condition:
                    input_errors["records_skipped"] += 1
                    if strict_input:
                        raise ValueError(f"{src}:{line_no}: missing case_id/condition")
                    continue
                _validate_eval_record_schema(record, src=src, line_no=line_no)
                grouped[(run_id, case_id, condition)].append(record)
    return grouped, input_errors


def _validate_eval_record_schema(record: dict[str, Any], *, src: str, line_no: int) -> None:
    for key in REQUIRED_TOP_LEVEL_FIELDS:
        if key not in record:
            raise ValueError(f"{src}:{line_no}: missing required field '{key}'")

    schema_version = record.get("schema_version")
    if schema_version != EVAL_RECORD_SCHEMA_VERSION:
        raise ValueError(
            f"{src}:{line_no}: unsupported schema_version={schema_version!r}, "
            f"expected {EVAL_RECORD_SCHEMA_VERSION}"
        )

    if not isinstance(record.get("run_id"), str) or not str(record.get("run_id")):
        raise ValueError(f"{src}:{line_no}: run_id must be a non-empty string")
    if not isinstance(record.get("case_id"), str) or not str(record.get("case_id")):
        raise ValueError(f"{src}:{line_no}: case_id must be a non-empty string")
    if not isinstance(record.get("condition"), str) or not str(record.get("condition")):
        raise ValueError(f"{src}:{line_no}: condition must be a non-empty string")
    if not isinstance(record.get("success"), bool):
        raise ValueError(f"{src}:{line_no}: success must be a boolean")

    attempt = record.get("attempt")
    if not isinstance(attempt, int) or attempt < 0:
        raise ValueError(f"{src}:{line_no}: attempt must be a non-negative integer")

    run_metadata = record.get("run_metadata")
    if not isinstance(run_metadata, dict):
        raise ValueError(f"{src}:{line_no}: run_metadata must be an object")
    for key in REQUIRED_RUN_METADATA_FIELDS:
        if key not in run_metadata:
            raise ValueError(f"{src}:{line_no}: run_metadata missing '{key}'")
    if run_metadata.get("run_id") != record.get("run_id"):
        raise ValueError(f"{src}:{line_no}: run_metadata.run_id must match run_id")
    if not isinstance(run_metadata.get("deterministic"), bool):
        raise ValueError(f"{src}:{line_no}: run_metadata.deterministic must be a boolean")
    if run_metadata.get("seed") is not None and not isinstance(run_metadata.get("seed"), int):
        raise ValueError(f"{src}:{line_no}: run_metadata.seed must be int|null")

    attempt_metadata = record.get("attempt_metadata")
    if not isinstance(attempt_metadata, dict):
        raise ValueError(f"{src}:{line_no}: attempt_metadata must be an object")
    for key in REQUIRED_ATTEMPT_METADATA_FIELDS:
        if key not in attempt_metadata:
            raise ValueError(f"{src}:{line_no}: attempt_metadata missing '{key}'")
    if attempt_metadata.get("attempt") != attempt:
        raise ValueError(f"{src}:{line_no}: attempt_metadata.attempt must match attempt")
    _validate_optional_string_field(
        attempt_metadata.get("prompt_sha256"),
        path="attempt_metadata.prompt_sha256",
        src=src,
        line_no=line_no,
    )
    _validate_optional_string_field(
        attempt_metadata.get("prompt_path"),
        path="attempt_metadata.prompt_path",
        src=src,
        line_no=line_no,
    )
    _validate_optional_string_field(
        attempt_metadata.get("evidence_sha256"),
        path="attempt_metadata.evidence_sha256",
        src=src,
        line_no=line_no,
    )
    _validate_optional_string_field(
        attempt_metadata.get("evidence_path"),
        path="attempt_metadata.evidence_path",
        src=src,
        line_no=line_no,
    )

    model_metadata = record.get("model_metadata")
    if not isinstance(model_metadata, dict):
        raise ValueError(f"{src}:{line_no}: model_metadata must be an object")
    for key in REQUIRED_MODEL_METADATA_FIELDS:
        if key not in model_metadata:
            raise ValueError(f"{src}:{line_no}: model_metadata missing '{key}'")
    backend = model_metadata.get("backend")
    if not isinstance(backend, str) or not backend:
        raise ValueError(f"{src}:{line_no}: model_metadata.backend must be a non-empty string")

    if backend == "openai":
        for key in REQUIRED_OPENAI_MODEL_METADATA_FIELDS:
            if key not in model_metadata:
                raise ValueError(
                    f"{src}:{line_no}: model_metadata missing '{key}' for openai backend"
                )
        requested_seed = model_metadata.get("api_compat_seed_requested")
        if requested_seed is not None and not isinstance(requested_seed, int):
            raise ValueError(
                f"{src}:{line_no}: model_metadata.api_compat_seed_requested must be int|null"
            )
        acknowledged = model_metadata.get("api_compat_seed_acknowledged")
        if not isinstance(acknowledged, bool):
            raise ValueError(
                f"{src}:{line_no}: model_metadata.api_compat_seed_acknowledged must be a boolean"
            )


def _validate_optional_string_field(value: Any, *, path: str, src: str, line_no: int) -> None:
    if value is not None and not isinstance(value, str):
        raise ValueError(f"{src}:{line_no}: {path} must be string|null")


def _derive_case_condition_rows(
    grouped: dict[tuple[str, str, str], list[dict[str, Any]]],
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for (run_id, case_id, condition), records in sorted(grouped.items()):
        ordered = sorted(records, key=lambda item: _as_int(item.get("attempt")))
        success = any(bool(item.get("success")) for item in ordered)
        attempts_to_success = _attempts_to_success(ordered)
        tokens = _tokens_used(ordered)
        runtime = _runtime_used(ordered)
        exemplar = ordered[0]
        rows.append(
            {
                "run_id": run_id,
                "case_id": case_id,
                "condition": condition,
                "category": str(exemplar.get("category") or "unknown"),
                "difficulty": str(exemplar.get("difficulty") or "unknown"),
                "snapshot_dependency": str(exemplar.get("snapshot_dependency") or "unknown"),
                "bug_source": exemplar.get("bug_source"),
                "contamination_risk": exemplar.get("contamination_risk"),
                "success": success,
                "attempts_to_success": attempts_to_success,
                "token_usage": tokens,
                "runtime_sec": runtime,
            }
        )
    return rows


def _summarize_bucket(rows: list[dict[str, Any]]) -> dict[str, Any]:
    by_condition: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        by_condition[str(row["condition"])].append(row)

    out: dict[str, Any] = {}
    for condition, condition_rows in sorted(by_condition.items()):
        successes = sum(1 for row in condition_rows if bool(row["success"]))
        cases = len(condition_rows)
        success_rate = (successes / cases) if cases else 0.0
        successful_attempts = [
            int(row["attempts_to_success"])
            for row in condition_rows
            if row["attempts_to_success"] is not None and bool(row["success"])
        ]
        token_values = [
            float(row["token_usage"]) for row in condition_rows if row["token_usage"] is not None
        ]
        runtime_values = [
            float(row["runtime_sec"]) for row in condition_rows if row["runtime_sec"] is not None
        ]
        out[condition] = {
            "cases": cases,
            "successes": successes,
            "success_rate": success_rate,
            "median_attempts_to_success": _median(successful_attempts),
            "mean_token_usage": _mean(token_values),
            "mean_runtime_sec": _mean(runtime_values),
        }

    out["uplift_success_rate"] = _uplift_success_rate(out)
    return out


def _attempts_to_success(records: list[dict[str, Any]]) -> int | None:
    for record in records:
        if bool(record.get("success")):
            return _as_int(record.get("attempt"))
    return None


def _tokens_used(records: list[dict[str, Any]]) -> float | None:
    total = 0.0
    found = False
    for record in records:
        patch_meta = record.get("patch_meta")
        if not isinstance(patch_meta, dict):
            continue
        usage = patch_meta.get("usage")
        if not isinstance(usage, dict):
            continue
        value = usage.get("total_tokens")
        if isinstance(value, (int, float)):
            found = True
            total += float(value)
    return total if found else None


def _runtime_used(records: list[dict[str, Any]]) -> float | None:
    if not records:
        return None
    cumulative_values = [
        float(record["cumulative_runtime_sec"])
        for record in records
        if isinstance(record.get("cumulative_runtime_sec"), (int, float))
    ]
    if cumulative_values:
        return max(cumulative_values)
    total = 0.0
    found = False
    first_before = records[0].get("before_duration_sec")
    if isinstance(first_before, (int, float)):
        total += float(first_before)
        found = True
    for record in records:
        after = record.get("after_duration_sec")
        if isinstance(after, (int, float)):
            total += float(after)
            found = True
    return total if found else None


def _uplift_success_rate(summary: dict[str, Any]) -> float | None:
    a = summary.get("traceback_only")
    b = summary.get("with_snapshot")
    if not isinstance(a, dict) or not isinstance(b, dict):
        return None
    ar = a.get("success_rate")
    br = b.get("success_rate")
    if not isinstance(ar, (int, float)) or not isinstance(br, (int, float)):
        return None
    return float(br - ar)


def _print_summary(summary: dict[str, Any]) -> None:
    print("Overall:")
    _print_bucket(summary["overall"])

    print("\nBy category:")
    by_category = summary["by_category"]
    for category, bucket in by_category.items():
        print(f"- {category}:")
        _print_bucket(bucket, prefix="  ")

    print("\nBy difficulty:")
    by_difficulty = summary["by_difficulty"]
    for difficulty, bucket in by_difficulty.items():
        print(f"- {difficulty}:")
        _print_bucket(bucket, prefix="  ")

    print("\nBy snapshot dependency:")
    by_snapshot_dependency = summary["by_snapshot_dependency"]
    for dependency, bucket in by_snapshot_dependency.items():
        print(f"- {dependency}:")
        _print_bucket(bucket, prefix="  ")

    if "by_bug_source" in summary:
        print("\nBy bug source:")
        for source, bucket in summary["by_bug_source"].items():
            print(f"- {source}:")
            _print_bucket(bucket, prefix="  ")

    if "by_contamination_risk" in summary:
        print("\nBy contamination risk:")
        for risk, bucket in summary["by_contamination_risk"].items():
            print(f"- {risk}:")
            _print_bucket(bucket, prefix="  ")

    _print_input_errors(summary.get("input_errors"))


def _print_bucket(bucket: dict[str, Any], prefix: str = "") -> None:
    for condition in ("traceback_only", "with_snapshot"):
        stats = bucket.get(condition)
        if not isinstance(stats, dict):
            continue
        cases = int(stats.get("cases", 0))
        successes = int(stats.get("successes", 0))
        rate = float(stats.get("success_rate", 0.0))
        med_attempts = stats.get("median_attempts_to_success")
        mean_tokens = stats.get("mean_token_usage")
        mean_runtime = stats.get("mean_runtime_sec")
        print(
            f"{prefix}- {condition}: {successes}/{cases} ({rate:.1%}), "
            f"median attempts={_fmt_num(med_attempts)}, "
            f"mean tokens={_fmt_num(mean_tokens)}, mean runtime={_fmt_num(mean_runtime)}s"
        )
    uplift = bucket.get("uplift_success_rate")
    print(f"{prefix}- uplift_success_rate: {_fmt_signed_percent(uplift)}")


def _print_input_errors(input_errors: Any) -> None:
    if not isinstance(input_errors, dict):
        return
    files_unreadable = int(input_errors.get("files_unreadable", 0))
    json_lines_invalid = int(input_errors.get("json_lines_invalid", 0))
    records_skipped = int(input_errors.get("records_skipped", 0))
    if files_unreadable == 0 and json_lines_invalid == 0 and records_skipped == 0:
        return
    print("\nInput diagnostics:")
    print(
        "- errors: "
        f"files_unreadable={files_unreadable}, "
        f"json_lines_invalid={json_lines_invalid}, "
        f"records_skipped={records_skipped}"
    )


def _fmt_num(value: Any) -> str:
    if value is None:
        return "n/a"
    if isinstance(value, float) and value.is_integer():
        return str(int(value))
    if isinstance(value, (int, float)):
        return f"{value:.2f}"
    return "n/a"


def _fmt_signed_percent(value: Any) -> str:
    if value is None:
        return "n/a"
    if not isinstance(value, (int, float)):
        return "n/a"
    return f"{value:+.1%}"


def _as_int(value: Any) -> int:
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        return int(value)
    if isinstance(value, str):
        try:
            return int(value)
        except ValueError:
            return 0
    return 0


def _mean(xs: list[float]) -> float | None:
    if not xs:
        return None
    return float(sum(xs) / len(xs))


def _median(xs: list[int]) -> float | None:
    ys = sorted(xs)
    n = len(ys)
    if n == 0:
        return None
    mid = n // 2
    if n % 2 == 1:
        return float(ys[mid])
    return (ys[mid - 1] + ys[mid]) / 2.0


def _expand_paths(items: list[str]) -> list[str]:
    out: list[str] = []
    for item in items:
        matches = sorted(glob.glob(item))
        out.extend(matches if matches else [item])
    return out


def _run_statistical_analysis(
    rows: list[dict[str, Any]], *, aggregation: str = "per_run"
) -> dict[str, Any]:
    """Run statistical analysis on paired rows using evals.stats module."""
    from evals.stats import (
        analyze_paired_results,
        any_success_aggregation,
        majority_vote_aggregation,
    )

    if aggregation in ("majority_vote", "any_success"):
        # Group rows by (case_id, condition) across runs.
        from evals.stats import BinaryPair

        by_case: dict[str, list[BinaryPair]] = {}
        case_meta: dict[str, dict[str, str]] = {}
        # First, group by (run_id, case_id) to get matched pairs per run.
        run_case_map: dict[str, dict[str, dict[str, Any]]] = defaultdict(lambda: defaultdict(dict))
        for row in rows:
            run_id = str(row.get("run_id", ""))
            case_id = str(row.get("case_id", ""))
            condition = str(row.get("condition", ""))
            run_case_map[run_id][case_id][condition] = row
            if case_id not in case_meta:
                case_meta[case_id] = {
                    "category": str(row.get("category", "unknown")),
                    "run_id": run_id,
                }

        # Build per-case lists of (A, B) pairs across runs.
        for cases in run_case_map.values():
            for case_id, conditions in cases.items():
                a = conditions.get("traceback_only")
                b = conditions.get("with_snapshot")
                if a is not None and b is not None:
                    pair: BinaryPair = (bool(a.get("success")), bool(b.get("success")))
                    by_case.setdefault(case_id, []).append(pair)

        if aggregation == "majority_vote":
            aggregated_pairs = majority_vote_aggregation(by_case)
        else:
            aggregated_pairs = any_success_aggregation(by_case)

        # Rebuild rows from aggregated pairs for analyze_paired_results.
        aggregated_rows: list[dict[str, Any]] = []
        for case_id, pair_val in zip(sorted(by_case), aggregated_pairs, strict=True):
            meta = case_meta.get(case_id, {})
            aggregated_rows.append(
                {
                    "run_id": "aggregated",
                    "case_id": case_id,
                    "condition": "traceback_only",
                    "category": meta.get("category", "unknown"),
                    "success": pair_val[0],
                }
            )
            aggregated_rows.append(
                {
                    "run_id": "aggregated",
                    "case_id": case_id,
                    "condition": "with_snapshot",
                    "category": meta.get("category", "unknown"),
                    "success": pair_val[1],
                }
            )
        result = analyze_paired_results(aggregated_rows)
        result["aggregation"] = aggregation
        return result

    result = analyze_paired_results(rows)
    result["aggregation"] = "per_run"
    return result


def _print_stats(stats: dict[str, Any]) -> None:
    """Print statistical analysis results."""
    if "error" in stats:
        print(f"\nStatistics: {stats['error']}")
        return

    print(f"\nStatistical analysis (aggregation={stats.get('aggregation', 'per_run')}):")
    print(f"  Matched pairs: {stats.get('n_pairs', 0)}")
    print(f"  A success rate: {stats.get('a_success_rate', 0):.1%}")
    print(f"  B success rate: {stats.get('b_success_rate', 0):.1%}")

    mcn = stats.get("overall_mcnemar", {})
    print(
        f"  McNemar's test: n_01={mcn.get('n_01', 0)}, n_10={mcn.get('n_10', 0)}, "
        f"p={mcn.get('p_value', 1.0):.4f} ({mcn.get('method', '?')})"
        f"{' *' if mcn.get('significant') else ''}"
    )

    ch = stats.get("overall_cohens_h", {})
    print(f"  Cohen's h: {ch.get('h', 0):.3f} ({ch.get('magnitude', '?')})")

    bci = stats.get("overall_bootstrap_ci", {})
    print(
        f"  Bootstrap CI ({bci.get('confidence', 0.95):.0%}): "
        f"[{bci.get('ci_lower', 0):.3f}, {bci.get('ci_upper', 0):.3f}]"
    )

    pwr = stats.get("overall_power", {})
    print(f"  Power: {pwr.get('approximate_power', 0):.2f} — {pwr.get('recommendation', '')}")

    per_cat = stats.get("per_category", {})
    if per_cat:
        hb = stats.get("holm_bonferroni", {})
        print("  Per-category:")
        for cat, cat_stats in sorted(per_cat.items()):
            cat_mcn = cat_stats.get("mcnemar", {})
            cat_ch = cat_stats.get("cohens_h", {})
            hb_entry = hb.get(cat, {})
            sig_marker = " *" if hb_entry.get("rejected") else ""
            adj_p = hb_entry.get("adjusted_p")
            adj_str = f", adj_p={adj_p:.4f}" if adj_p is not None else ""
            print(
                f"    {cat}: n={cat_stats.get('n_pairs', 0)}, "
                f"p={cat_mcn.get('p_value', 1.0):.4f}{adj_str}, "
                f"h={cat_ch.get('h', 0):.3f}{sig_marker}"
            )


if __name__ == "__main__":
    raise SystemExit(main())
