#!/usr/bin/env python3
"""Evaluate sampled DeepSeek-V2-Lite-Chat completions safely.

Stopping, task-terminal state, evaluator coverage, and correctness remain
separate ledgers.  HumanEval candidates execute only in fresh networkless,
read-only Docker containers.  An exact candidate+task+harness cache avoids
re-executing duplicate code while preserving one result row per sampled cell.
"""

from __future__ import annotations

import argparse
import copy
import hashlib
import json
from collections import Counter
from pathlib import Path
from statistics import mean
from typing import Any

import v2_lite_chat_completion_evaluator as completion


COMPARISONS = (
    ("system_eos", "s0_eos", "s1_eos"),
    ("system_bos", "s0_bos", "s1_bos"),
    ("system_x", "s0_x", "s1_x"),
    ("system_period", "s0_period", "s1_period"),
    ("bos_at_s0", "s0_eos", "s0_bos"),
    ("bos_at_s1", "s1_eos", "s1_bos"),
    ("x_at_s0", "s0_eos", "s0_x"),
    ("x_at_s1", "s1_eos", "s1_x"),
    ("period_at_s0", "s0_eos", "s0_period"),
    ("period_at_s1", "s1_eos", "s1_period"),
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--sampling-json", type=Path, required=True)
    parser.add_argument("--human-eval", type=Path, required=True)
    parser.add_argument("--gsm8k", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--sandbox-image",
        required=True,
        help="Pinned image reference including @sha256 digest.",
    )
    parser.add_argument("--timeout-seconds", type=float, default=5.0)
    parser.add_argument("--skip-code-execution", action="store_true")
    return parser.parse_args()


def evaluate_code_cached(
    *,
    text: str,
    hit_eos: bool,
    task: dict[str, Any],
    image: str,
    timeout_seconds: float,
    skip_execution: bool,
    cache: dict[tuple[str, str, str], dict[str, Any]],
) -> tuple[dict[str, Any], bool]:
    extracted = completion.extract_code(
        text,
        task["prompt"],
        task["entry_point"],
    )
    harness = completion.sandbox_harness(
        extracted["candidate"],
        task["test"],
        task["entry_point"],
    )
    cache_key = (
        extracted["candidate_sha256"],
        hashlib.sha256(task["test"].encode()).hexdigest(),
        hashlib.sha256(harness.encode()).hexdigest(),
    )
    cache_hit = cache_key in cache
    if cache_hit:
        execution = copy.deepcopy(cache[cache_key])
    else:
        execution = {
            "status": "not_run",
            "return_code": None,
            "runtime_ms": None,
            "harness_sha256": cache_key[2],
        }
        if extracted["python_ast_parse"] and not skip_execution:
            execution = completion.execute_code(
                extracted["candidate"],
                task,
                image,
                timeout_seconds,
            )
        cache[cache_key] = copy.deepcopy(execution)
    passed = execution["status"] == "passed"
    semantic_terminal = bool(
        hit_eos
        or extracted["closed_code_fence"]
        or passed
    )
    result = {
        key: value
        for key, value in extracted.items()
        if key != "candidate"
    } | {
        "semantic_terminal": semantic_terminal,
        "evaluator_covered": (
            extracted["python_ast_parse"]
            and execution["status"] != "not_run"
        ),
        "execution": execution,
        "fixed_budget_tests_pass": passed,
        "strict_complete_tests_pass": (
            passed and semantic_terminal
        ),
        "execution_cache_hit": cache_hit,
    }
    return result, cache_hit


def output_success(row: dict[str, Any]) -> dict[str, int | None]:
    evaluation = row["task_evaluation"]
    if row["domain"] == "math":
        return {
            "fixed_budget": int(
                evaluation["fixed_budget_numeric_exact"]
            ),
            "strict_complete": int(
                evaluation["strict_complete_numeric_exact"]
            ),
        }
    if row["domain"] == "code":
        return {
            "fixed_budget": int(
                evaluation["fixed_budget_tests_pass"]
            ),
            "strict_complete": int(
                evaluation["strict_complete_tests_pass"]
            ),
        }
    return {
        "fixed_budget": None,
        "strict_complete": None,
    }


def summarize_group(rows: list[dict[str, Any]]) -> dict[str, Any]:
    math_rows = [
        row for row in rows if row["domain"] == "math"
    ]
    code_rows = [
        row for row in rows if row["domain"] == "code"
    ]
    math_source_ids = sorted(
        {row["source_id"] for row in math_rows}
    )
    single_math_source = len(math_source_ids) == 1
    answers = (
        Counter(
            row["task_evaluation"]["predicted_final"]
            for row in math_rows
            if row["task_evaluation"]["predicted_final"] is not None
        )
        if single_math_source
        else Counter()
    )
    max_count = max(answers.values(), default=0)
    modes = sorted(
        answer
        for answer, count in answers.items()
        if count == max_count
    )
    gold = (
        math_rows[0]["task_evaluation"]["gold_final"]
        if single_math_source
        else None
    )
    return {
        "outputs": len(rows),
        "natural_eos": sum(row["hit_eos"] for row in rows),
        "budget_truncated": sum(
            row["stopped_at_max_new_tokens"] for row in rows
        ),
        "mean_generated_tokens": (
            mean(row["generated_tokens"] for row in rows)
            if rows
            else None
        ),
        "completion_classes": dict(
            Counter(row["completion_class"] for row in rows)
        ),
        "math": {
            "outputs": len(math_rows),
            "sources": len(math_source_ids),
            "source_ids": math_source_ids,
            "answer_aggregation_scope": (
                "single_source"
                if single_math_source
                else "disabled_across_distinct_gold_answers"
            ),
            "evaluator_covered": sum(
                row["task_evaluation"]["evaluator_covered"]
                for row in math_rows
            ),
            "fixed_budget_exact": sum(
                row["task_evaluation"][
                    "fixed_budget_numeric_exact"
                ]
                for row in math_rows
            ),
            "strict_complete_exact": sum(
                row["task_evaluation"][
                    "strict_complete_numeric_exact"
                ]
                for row in math_rows
            ),
            "answer_frequencies": dict(answers),
            "modal_answers": modes,
            "modal_count": max_count,
            "absolute_majority_exists": (
                max_count > len(math_rows) / 2
                if single_math_source
                else None
            ),
            "unique_absolute_majority": (
                modes[0]
                if (
                    single_math_source
                    and len(modes) == 1
                    and max_count > len(math_rows) / 2
                )
                else None
            ),
            "gold": gold,
            "unique_absolute_majority_matches_gold": (
                len(modes) == 1
                and max_count > len(math_rows) / 2
                and modes[0] == gold
                if single_math_source
                else None
            ),
        },
        "code": {
            "outputs": len(code_rows),
            "sources": len(
                {row["source_id"] for row in code_rows}
            ),
            "ast_parse": sum(
                row["task_evaluation"]["python_ast_parse"]
                for row in code_rows
            ),
            "executed": sum(
                row["task_evaluation"]["execution"]["status"]
                != "not_run"
                for row in code_rows
            ),
            "execution_cache_hits": sum(
                row["task_evaluation"]["execution_cache_hit"]
                for row in code_rows
            ),
            "tests_pass": sum(
                row["task_evaluation"][
                    "fixed_budget_tests_pass"
                ]
                for row in code_rows
            ),
            "strict_complete_tests_pass": sum(
                row["task_evaluation"][
                    "strict_complete_tests_pass"
                ]
                for row in code_rows
            ),
            "execution_statuses": dict(
                Counter(
                    row["task_evaluation"]["execution"]["status"]
                    for row in code_rows
                )
            ),
        },
    }


def source_condition_summary(
    rows: list[dict[str, Any]],
) -> dict[str, Any]:
    result: dict[str, dict[str, Any]] = {}
    source_ids = sorted({row["source_id"] for row in rows})
    conditions = tuple(
        condition
        for _, left, right in COMPARISONS
        for condition in (left, right)
    )
    ordered_conditions = tuple(dict.fromkeys(conditions))
    for source_id in source_ids:
        result[source_id] = {}
        for condition in ordered_conditions:
            subset = [
                row
                for row in rows
                if row["source_id"] == source_id
                and row["condition"] == condition
            ]
            result[source_id][condition] = summarize_group(subset)
    return result


def edge_summary(rows: list[dict[str, Any]]) -> dict[str, Any]:
    result: dict[str, dict[str, Any]] = {}
    source_ids = sorted({row["source_id"] for row in rows})
    for source_id in source_ids:
        result[source_id] = {}
        source_rows = [
            row for row in rows if row["source_id"] == source_id
        ]
        for name, left, right in COMPARISONS:
            left_rows = [
                row
                for row in source_rows
                if row["condition"] == left
            ]
            right_rows = [
                row
                for row in source_rows
                if row["condition"] == right
            ]
            left_success = [
                output_success(row) for row in left_rows
            ]
            right_success = [
                output_success(row) for row in right_rows
            ]
            result[source_id][name] = {
                "left": left,
                "right": right,
                "samples_per_side": len(left_rows),
                "completion_classes": {
                    "left": dict(
                        Counter(
                            row["completion_class"]
                            for row in left_rows
                        )
                    ),
                    "right": dict(
                        Counter(
                            row["completion_class"]
                            for row in right_rows
                        )
                    ),
                },
                "natural_eos_difference_right_minus_left": (
                    sum(row["hit_eos"] for row in right_rows)
                    - sum(row["hit_eos"] for row in left_rows)
                ),
                "fixed_budget_success_difference_right_minus_left": (
                    sum(
                        row["fixed_budget"]
                        for row in right_success
                        if row["fixed_budget"] is not None
                    )
                    - sum(
                        row["fixed_budget"]
                        for row in left_success
                        if row["fixed_budget"] is not None
                    )
                    if any(
                        row["fixed_budget"] is not None
                        for row in [*left_success, *right_success]
                    )
                    else None
                ),
                "strict_success_difference_right_minus_left": (
                    sum(
                        row["strict_complete"]
                        for row in right_success
                        if row["strict_complete"] is not None
                    )
                    - sum(
                        row["strict_complete"]
                        for row in left_success
                        if row["strict_complete"] is not None
                    )
                    if any(
                        row["strict_complete"] is not None
                        for row in [*left_success, *right_success]
                    )
                    else None
                ),
            }
    return result


def main() -> None:
    args = parse_args()
    for path in (
        args.sampling_json,
        args.human_eval,
        args.gsm8k,
    ):
        if not path.is_file():
            raise FileNotFoundError(path)
    if "@sha256:" not in args.sandbox_image:
        raise ValueError("--sandbox-image must include an immutable digest")
    if args.timeout_seconds <= 0:
        raise ValueError("--timeout-seconds must be positive")

    sampling = json.loads(
        args.sampling_json.read_text(encoding="utf-8")
    )
    human_eval, gsm8k = completion.load_tasks(
        args.human_eval,
        args.gsm8k,
    )
    code_cache: dict[
        tuple[str, str, str],
        dict[str, Any],
    ] = {}
    rows = []
    for source in sampling["sources"]:
        for run in source["runs"]:
            for output in run["outputs"]:
                evaluation = None
                cache_hit = False
                if source["domain"] == "math":
                    evaluation = completion.evaluate_math(
                        output["text"],
                        output["hit_eos"],
                        gsm8k[source["id"]]["answer"],
                    )
                elif source["domain"] == "code":
                    evaluation, cache_hit = evaluate_code_cached(
                        text=output["text"],
                        hit_eos=output["hit_eos"],
                        task=human_eval[source["id"]],
                        image=args.sandbox_image,
                        timeout_seconds=args.timeout_seconds,
                        skip_execution=args.skip_code_execution,
                        cache=code_cache,
                    )
                rows.append(
                    {
                        "source_id": source["id"],
                        "domain": source["domain"],
                        "replicate_index": run["replicate_index"],
                        "replicate_label": run["replicate_label"],
                        "base_seed": run["base_seed"],
                        "run_seed": run["run_seed"],
                        "condition": output["condition"],
                        "generated_tokens": output[
                            "generated_tokens"
                        ],
                        "hit_eos": output["hit_eos"],
                        "stopped_at_max_new_tokens": output[
                            "stopped_at_max_new_tokens"
                        ],
                        "prompt_token_ids_sha256": output[
                            "prompt_token_ids_sha256"
                        ],
                        "generated_token_ids_sha256": output[
                            "generated_token_ids_sha256"
                        ],
                        "text_sha256": output["text_sha256"],
                        "task_evaluation": evaluation,
                        "completion_class": (
                            completion.completion_class(
                                output,
                                evaluation,
                            )
                        ),
                        "code_execution_cache_hit": cache_hit,
                    }
                )

    by_condition = {
        condition: summarize_group(
            [
                row
                for row in rows
                if row["condition"] == condition
            ]
        )
        for condition in sampling["seed_contract"][
            "condition_row_order"
        ]
    }
    result = {
        "schema_version": 1,
        "protocol_id": sampling["protocol_id"],
        "input": {
            "sampling_path": str(args.sampling_json),
            "sampling_sha256": completion.sha256_file(
                args.sampling_json
            ),
            "sampling_content_hash": sampling["content_hash"],
            "human_eval_sha256": completion.sha256_file(
                args.human_eval
            ),
            "gsm8k_sha256": completion.sha256_file(args.gsm8k),
            "model_revision": sampling["model"]["revision"],
            "base_seeds": sampling["seed_contract"][
                "executed_base_seeds"
            ],
            "max_new_tokens": sampling["generation_contract"][
                "max_new_tokens"
            ],
        },
        "sandbox": {
            "image": args.sandbox_image,
            "timeout_seconds": args.timeout_seconds,
            "code_execution_skipped": args.skip_code_execution,
            "network": "none",
            "filesystem": "read-only",
            "user": "65534:65534",
            "capabilities": "ALL dropped",
            "no_new_privileges": True,
            "memory": "256m",
            "memory_swap": "256m",
            "pids_limit": 64,
            "cpus": 0.5,
            "tmpfs": "/tmp:rw,noexec,nosuid,size=16m",
            "host_mounts": 0,
            "cache_key": (
                "candidate SHA-256 + task-test SHA-256 + "
                "sandbox-harness SHA-256"
            ),
            "unique_code_cache_entries": len(code_cache),
        },
        "rows": rows,
        "summary": {
            **summarize_group(rows),
            "by_condition": by_condition,
            "by_source_condition": source_condition_summary(rows),
            "by_source_edge": edge_summary(rows),
        },
        "claim_boundary": [
            (
                f"{len(sampling['sources'])} sources and "
                f"{len(sampling['seed_contract']['executed_base_seeds'])} "
                "seeds are not full benchmark estimates."
            ),
            "Completion-conditioned metrics are selection-biased diagnostics.",
            "A passing HumanEval test is functional evidence, not code-safety evidence.",
            "A modal sampled math answer is not standard self-consistency.",
            "Counterfactual token sequences are not official-valid chats.",
            "Repeated samples from one source are not independent tasks.",
        ],
    }
    result["content_hash"] = completion.canonical_hash(rows)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(result, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    payload = args.output.read_bytes()
    print(
        json.dumps(
            {
                "output": str(args.output),
                "sha256": completion.sha256_bytes(payload),
                "bytes": len(payload),
                "outputs": len(rows),
                "summary": {
                    key: result["summary"][key]
                    for key in (
                        "natural_eos",
                        "budget_truncated",
                        "completion_classes",
                        "math",
                        "code",
                    )
                },
                "unique_code_cache_entries": len(code_cache),
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()