from __future__ import annotations

import argparse
import json
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_CHECKPOINT = REPO_ROOT / "checkpoints" / "29999"
ASSET_ID = "real_piper_x_6tasks"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Validate the released StreamPI Piper-X 20k "
            "deployment checkpoint."
        )
    )
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=DEFAULT_CHECKPOINT,
        help="Checkpoint directory. Default: checkpoints/29999",
    )
    return parser.parse_args()


def directory_size(path: Path) -> int:
    total = 0

    for item in path.rglob("*"):
        if item.is_file():
            total += item.stat().st_size

    return total


def main() -> None:
    args = parse_args()
    checkpoint = args.checkpoint.expanduser().resolve()

    required_files = (
        checkpoint / "params" / "_METADATA",
        checkpoint / "params" / "_sharding",
        checkpoint / "params" / "manifest.ocdbt",
        checkpoint
        / "assets"
        / ASSET_ID
        / "norm_stats.json",
    )

    missing = [
        path for path in required_files
        if not path.is_file()
    ]

    if missing:
        details = "\n".join(
            f"  - {path}" for path in missing
        )
        raise FileNotFoundError(
            "Missing deployment checkpoint files:\n"
            + details
        )

    forbidden_paths = (
        checkpoint / "train_state",
        checkpoint / "optimizer",
    )

    unexpected = [
        path for path in forbidden_paths
        if path.exists()
    ]

    if unexpected:
        details = "\n".join(
            f"  - {path}" for path in unexpected
        )
        raise RuntimeError(
            "Training-only artifacts remain in the "
            "deployment checkpoint:\n"
            + details
        )

    norm_path = (
        checkpoint
        / "assets"
        / ASSET_ID
        / "norm_stats.json"
    )

    with norm_path.open("r", encoding="utf-8") as stream:
        norm_stats = json.load(stream)

    if not isinstance(norm_stats, dict) or not norm_stats:
        raise ValueError(
            f"Invalid or empty norm stats: {norm_path}"
        )

    parameter_files = [
        path
        for path in (checkpoint / "params").rglob("*")
        if path.is_file()
    ]

    if not parameter_files:
        raise RuntimeError(
            f"No parameter files found under {checkpoint / 'params'}"
        )

    symlinks = [
        path
        for path in checkpoint.rglob("*")
        if path.is_symlink()
    ]

    if symlinks:
        details = "\n".join(
            f"  - {path}" for path in symlinks
        )
        raise RuntimeError(
            "Deployment checkpoint contains symbolic links:\n"
            + details
        )

    for required in required_files:
        print(
            "OK:",
            required.relative_to(checkpoint),
        )

    print(f"checkpoint={checkpoint}")
    print(f"parameter_files={len(parameter_files)}")
    print(f"checkpoint_bytes={directory_size(checkpoint)}")
    print(f"norm_stat_keys={sorted(norm_stats)}")
    print("train_state_present=false")
    print("symlink_count=0")
    print("CHECKPOINT_STRUCTURE_OK")


if __name__ == "__main__":
    main()