from __future__ import annotations

import argparse
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate a hackathon-ready submission bundle from training artifacts")
    parser.add_argument(
        "--eval-summary",
        default="artifacts/ppo_tuned_v3/evaluation_summary.json",
        help="Path to evaluation_summary.json",
    )
    parser.add_argument(
        "--output-dir",
        default="artifacts/submission_bundle",
        help="Directory to write submission bundle files",
    )
    parser.add_argument(
        "--run-command",
        default=(
            "python scripts/train_sb3_agent.py --stage-steps 8000,8000,8000 "
            "--entropy-schedule 0.02,0.01,0.003 --eval-freq 100000000 --output-dir artifacts/ppo_tuned_v3"
        ),
        help="Training command used to produce the metrics",
    )
    return parser.parse_args()


def _load_json(path: Path) -> Dict[str, Any]:
    if not path.exists():
        raise FileNotFoundError(f"Missing file: {path}")
    return json.loads(path.read_text(encoding="utf-8"))


def _task_row(task: str, summary: Dict[str, Any]) -> Dict[str, Any]:
    payload = summary.get(task, {})
    scores = payload.get("scores", [])
    return {
        "task": task,
        "avg": float(payload.get("avg_score", 0.0) or 0.0),
        "min": float(payload.get("min_score", 0.0) or 0.0),
        "max": float(payload.get("max_score", 0.0) or 0.0),
        "scores": [float(value) for value in scores],
    }


def main() -> None:
    args = parse_args()
    project_root = Path(__file__).resolve().parents[1]

    eval_summary_path = (project_root / args.eval_summary).resolve()
    output_dir = (project_root / args.output_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    summary = _load_json(eval_summary_path)

    tasks = ["clean_sweep", "tangled_lake", "legal_minefield", "cascade_collapse"]
    rows = [_task_row(task, summary) for task in tasks]

    legal_row = next(row for row in rows if row["task"] == "legal_minefield")
    hard_gate = 0.75

    bundle = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "project": "gdpr-eraseops",
        "training_command": args.run_command,
        "evaluation_summary_path": str(eval_summary_path),
        "tasks": rows,
        "headline": {
            "metric": "legal_minefield_avg",
            "value": legal_row["avg"],
            "target": hard_gate,
            "pass": legal_row["avg"] >= hard_gate,
            "margin": round(legal_row["avg"] - hard_gate, 4),
        },
    }

    json_path = output_dir / "submission_report.json"
    json_path.write_text(json.dumps(bundle, indent=2), encoding="utf-8")

    md_lines = [
        "# GDPR-EraseOps Submission Bundle",
        "",
        f"Generated at: {bundle['generated_at']}",
        "",
        "## Headline",
        "",
        f"- legal_minefield avg: {legal_row['avg']:.3f}",
        f"- target: {hard_gate:.2f}",
        f"- pass: {str(bundle['headline']['pass']).lower()}",
        f"- margin: {bundle['headline']['margin']:+.3f}",
        "",
        "## Task Scores",
        "",
        "| Task | Avg | Min | Max | Seeds |",
        "|---|---:|---:|---:|---|",
    ]

    for row in rows:
        seeds = ", ".join(f"{value:.2f}" for value in row["scores"])
        md_lines.append(
            f"| {row['task']} | {row['avg']:.3f} | {row['min']:.3f} | {row['max']:.3f} | [{seeds}] |"
        )

    md_lines.extend(
        [
            "",
            "## Reproduce",
            "",
            "```bash",
            args.run_command,
            "```",
            "",
            "## Primary Sources",
            "",
            f"- {eval_summary_path}",
            f"- {json_path}",
        ]
    )

    md_path = output_dir / "submission_report.md"
    md_path.write_text("\n".join(md_lines) + "\n", encoding="utf-8")

    print(json.dumps({"json": str(json_path), "markdown": str(md_path)}, indent=2))


if __name__ == "__main__":
    main()