#!/usr/bin/env python3
from __future__ import annotations

import hashlib
import importlib.util
import json
import os
import sys
from pathlib import Path
from typing import Any

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from flow_grpo.server_profiles import apply_server_profile_defaults


apply_server_profile_defaults()

OUT_DIR = REPO_ROOT / "analysis_outputs" / "h20_eval_corruption"
DEFAULT_OLD_RL_LORA = REPO_ROOT / "logs/radiomics/img-only-r32-a64-bs32-evalbs24-kl-beta0p005-scratch-15k/checkpoints/checkpoint-190/lora"


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def tensor_summary(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    if path.suffix == ".safetensors":
        try:
            from safetensors import safe_open
        except Exception as exc:
            return {"error": f"safetensors import failed: {exc!r}"}
        by_dtype: dict[str, int] = {}
        first = []
        with safe_open(path, framework="pt", device="cpu") as handle:
            keys = list(handle.keys())
            for key in keys[:20]:
                tensor = handle.get_tensor(key)
                by_dtype[str(tensor.dtype)] = by_dtype.get(str(tensor.dtype), 0) + 1
                first.append({"key": key, "shape": list(tensor.shape), "dtype": str(tensor.dtype)})
        return {"key_count": len(keys), "dtype_counts_first20": by_dtype, "first_tensors": first}
    if path.suffix in {".bin", ".pt"}:
        try:
            import torch
            obj = torch.load(path, map_location="cpu")
        except Exception as exc:
            return {"error": f"torch load failed: {exc!r}"}
        state = obj if isinstance(obj, dict) else {}
        first = []
        by_dtype: dict[str, int] = {}
        for key, value in list(state.items())[:20]:
            if hasattr(value, "shape"):
                by_dtype[str(value.dtype)] = by_dtype.get(str(value.dtype), 0) + 1
                first.append({"key": key, "shape": list(value.shape), "dtype": str(value.dtype)})
        return {"key_count": len(state), "dtype_counts_first20": by_dtype, "first_tensors": first}
    return {}


def inspect_path(label: str, path: str | Path) -> dict[str, Any]:
    p = Path(path).expanduser()
    record: dict[str, Any] = {"label": label, "path": str(p), "exists": p.exists()}
    if p.exists() and p.is_file():
        record.update({"size_bytes": p.stat().st_size, "sha256": sha256_file(p)})
        if p.name.endswith((".json", ".md", ".txt")):
            try:
                record["json_keys"] = sorted(json.loads(p.read_text(encoding="utf-8")).keys()) if p.suffix == ".json" else None
            except Exception as exc:
                record["text_read_error"] = repr(exc)
        record["tensor_summary"] = tensor_summary(p)
    elif p.exists() and p.is_dir():
        try:
            record["entries"] = sorted(child.name for child in p.iterdir())[:100]
        except Exception as exc:
            record["entries_error"] = repr(exc)
    return record


def main() -> int:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    sft = Path(os.environ.get("SFT_LORA_PATH", ""))
    old_rl = Path(os.environ.get("OLD_RL_LORA_PATH") or os.environ.get("EVAL_LORA_PATH") or DEFAULT_OLD_RL_LORA)
    hf_home = Path(os.environ.get("HF_HOME", Path.home() / ".cache/huggingface"))
    hf_hub = Path(os.environ.get("HF_HUB_CACHE", hf_home / "hub"))
    snapshot_root = hf_hub / "models--Shitao--OmniGen-v1" / "snapshots"

    records = {
        "environment": {
            "python": sys.executable,
            "server_profile": os.environ.get("SERVER_PROFILE"),
            "hf_home": str(hf_home),
            "hf_hub_cache": str(hf_hub),
            "omnigen_code_root": os.environ.get("OMNIGEN_CODE_ROOT"),
            "sft_lora_path": str(sft),
            "old_rl_lora_path": str(old_rl),
        },
        "paths": [
            inspect_path("sft_lora_dir", sft),
            inspect_path("sft_adapter_model", sft / "adapter_model.safetensors"),
            inspect_path("sft_adapter_config", sft / "adapter_config.json"),
            inspect_path("old_rl_lora_dir", old_rl),
            inspect_path("old_rl_adapter_model", old_rl / "adapter_model.safetensors"),
            inspect_path("old_rl_adapter_config", old_rl / "adapter_config.json"),
            inspect_path("omnigen_code_root", os.environ.get("OMNIGEN_CODE_ROOT", "")),
            inspect_path("hf_snapshot_root", snapshot_root),
        ],
    }

    json_path = OUT_DIR / "model_file_integrity.json"
    md_path = OUT_DIR / "model_file_integrity.md"
    json_path.write_text(json.dumps(records, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    lines = ["# H20 Model File Integrity", ""]
    for item in records["paths"]:
        lines.extend([
            f"## {item['label']}",
            f"- path: `{item['path']}`",
            f"- exists: `{item['exists']}`",
            f"- size_bytes: `{item.get('size_bytes')}`",
            f"- sha256: `{item.get('sha256')}`",
            f"- tensor_summary: `{json.dumps(item.get('tensor_summary', {}), sort_keys=True)[:2000]}`",
            "",
        ])
    md_path.write_text("\n".join(lines), encoding="utf-8")
    print(json.dumps(records, indent=2, sort_keys=True))
    print(f"[integrity] wrote {json_path} and {md_path}")
    return 0


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