#!/usr/bin/env python3
"""Check required artifacts before a training run starts."""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Any

import yaml


DEFAULT_BASELINE_REPORT = Path("reports/cybergym/qwen36_27b_base/cybergym_base_level1_sample.md")
DEFAULT_FROZEN_TASKS = Path("reports/cybergym/frozen_level1_baseline_tasks.txt")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True, help="Training config YAML.")
    parser.add_argument(
        "--allow-missing-cybergym-baseline",
        action="store_true",
        help="Allow local dry-runs before the remote CyberGym baseline exists.",
    )
    return parser.parse_args()


def read_yaml(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as fh:
        payload = yaml.safe_load(fh) or {}
    if not isinstance(payload, dict):
        raise TypeError(f"Expected a YAML mapping in {path}")
    return payload


def existing_path(path: Path) -> Path | None:
    candidates = [path]
    if not path.is_absolute():
        candidates.append(Path("/workspace") / path)
        candidates.append(Path("/workspace/infosec") / path)
    for candidate in candidates:
        if candidate.is_file():
            return candidate
    return None


def main() -> int:
    args = parse_args()
    config_path = Path(args.config)
    config = read_yaml(config_path)

    missing: list[str] = []
    warnings: list[str] = []

    run_stage = str(config.get("run", {}).get("stage", "unknown"))
    baseline_required = run_stage in {"pilot_qlora", "stage1_lora_sft", "stage1_dpo"}

    gates = config.get("gates", {}) or {}
    baseline_report = Path(gates.get("cybergym_baseline_report", DEFAULT_BASELINE_REPORT))
    frozen_tasks = Path(gates.get("frozen_cybergym_tasks", DEFAULT_FROZEN_TASKS))

    if baseline_required and not args.allow_missing_cybergym_baseline:
        if existing_path(baseline_report) is None:
            missing.append(str(baseline_report))
        if existing_path(frozen_tasks) is None:
            missing.append(str(frozen_tasks))

    data = config.get("data", {}) or {}
    for key in ("train_jsonl", "validation_jsonl"):
        value = data.get(key)
        if value and existing_path(Path(value)) is None:
            missing.append(str(value))

    if data.get("decontamination_required"):
        report = data.get("decontamination_report")
        if not report:
            missing.append("data.decontamination_report (decontamination_required=true but unset)")
        elif existing_path(Path(report)) is None:
            missing.append(str(report))

    print(f"Config: {config_path}")
    print(f"Stage: {run_stage}")

    if warnings:
        print("Warnings:")
        for warning in warnings:
            print(f"  - {warning}")

    if missing:
        print("Missing required files:")
        for path in missing:
            print(f"  - {path}")
        return 2

    print("All configured training gates passed.")
    return 0


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