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

import argparse
import hashlib
import importlib.metadata
import importlib.util
import inspect
import json
import os
import shutil
import subprocess
import sys
from collections import Counter
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image, ImageDraw
import torch


REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT_DIR = REPO_ROOT / "analysis_outputs" / "a5500_eval_golden_reference"
DEFAULT_CONFIG_ENTRY = "config/grpo.py:general_radiomics_omnigen_4gpu_kl_eval"
DEFAULT_SFT_LORA = Path("/home/wenting/gen_joint/results_new/scratch_15k")
DEFAULT_RL_LORA = REPO_ROOT / "logs/radiomics/img-only-r32-a64-bs32-evalbs24-kl-beta0p005-scratch-15k/checkpoints/checkpoint-190/lora"
DEFAULT_OMNIGEN_CODE_ROOT = Path("/home/wenting/gen_joint")

KEY_FLOW_FILES = [
    "scripts/single_node/eval_4gpu.sh",
    "scripts/single_node/eval_4gpu_scratch15k_image_only.sh",
    "scripts/eval_omnigen.py",
    "scripts/train_omnigen.py",
    "config/grpo.py",
    "flow_grpo/omnigen_patch/omnigen_pipeline_with_logprob.py",
    "flow_grpo/omnigen_patch/joint_model_loader.py",
    "flow_grpo/omnigen_patch/__init__.py",
]

KEY_GEN_FILES = [
    "OmniGen/__init__.py",
    "OmniGen/pipeline.py",
    "OmniGen/scheduler.py",
    "OmniGen/model.py",
    "OmniGen/processor.py",
    "OmniGen/transformer.py",
]

FLOW_DIFF_TARGETS = [
    "scripts/single_node/eval_4gpu.sh",
    "scripts/single_node/eval_4gpu_scratch15k_image_only.sh",
    "scripts/eval_omnigen.py",
    "scripts/train_omnigen.py",
    "config/grpo.py",
    "flow_grpo/omnigen_patch",
]

GEN_DIFF_TARGETS = [
    "OmniGen",
]

PACKAGES = {
    "torch": "torch",
    "xformers": "xformers",
    "diffusers": "diffusers",
    "transformers": "transformers",
    "accelerate": "accelerate",
    "peft": "peft",
    "safetensors": "safetensors",
    "ml_collections": "ml-collections",
    "huggingface_hub": "huggingface-hub",
    "numpy": "numpy",
    "Pillow": "Pillow",
}


def _sha256(path: Path) -> str | None:
    if not path.is_file():
        return None
    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 _file_record(path: Path) -> dict[str, Any]:
    return {
        "path": str(path),
        "exists": path.is_file(),
        "size": path.stat().st_size if path.is_file() else None,
        "sha256": _sha256(path),
    }


def _run(cmd: list[str], cwd: Path) -> dict[str, Any]:
    try:
        proc = subprocess.run(cmd, cwd=str(cwd), check=False, text=True, capture_output=True)
        return {"cmd": cmd, "returncode": proc.returncode, "stdout": proc.stdout.strip(), "stderr": proc.stderr.strip()}
    except OSError as exc:
        return {"cmd": cmd, "error": repr(exc)}


def _package_versions() -> dict[str, Any]:
    versions = {}
    for label, package in PACKAGES.items():
        try:
            versions[label] = importlib.metadata.version(package)
        except importlib.metadata.PackageNotFoundError:
            versions[label] = None
    return versions


def _plain(value: Any) -> Any:
    if hasattr(value, "to_dict"):
        return _plain(value.to_dict())
    if hasattr(value, "items"):
        return {str(key): _plain(item) for key, item in value.items()}
    if isinstance(value, tuple):
        return [_plain(item) for item in value]
    if isinstance(value, list):
        return [_plain(item) for item in value]
    return value


def _load_config(config_entry: str):
    module_path, function_name = config_entry.split(":", 1)
    module_file = (REPO_ROOT / module_path).resolve() if not Path(module_path).is_absolute() else Path(module_path)
    spec = importlib.util.spec_from_file_location("a5500_eval_ref_config", module_file)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Could not load config from {module_file}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return getattr(module, function_name)()


def _relocate_path(value: Any, replacements: list[Any]) -> Any:
    if not isinstance(value, str):
        return value
    for old_root, new_root in replacements:
        old_root = str(old_root).rstrip("/")
        new_root = str(new_root).rstrip("/")
        if value == old_root:
            return new_root
        if value.startswith(f"{old_root}/"):
            return f"{new_root}/{value[len(old_root) + 1:]}"
    return value


def _relocate_metadata(metadata: dict[str, Any], replacements: list[Any]) -> dict[str, Any]:
    metadata = dict(metadata)
    for key in ("output_image", "gt_image", "output_mask", "gt_mask", "mask"):
        if key in metadata:
            metadata[key] = _relocate_path(metadata[key], replacements)
    if "input_images" in metadata:
        metadata["input_images"] = [_relocate_path(path, replacements) for path in metadata["input_images"]]
    return metadata


def _dataset_file(dataset: Any, split: str) -> Path:
    if hasattr(dataset, "get"):
        file_path = dataset.get(f"{split}_jsonl") or dataset.get("jsonl")
        if file_path is None and dataset.get("root"):
            file_path = Path(dataset.get("root")) / f"{split}_metadata.jsonl"
        if file_path is None:
            raise ValueError(f"Dataset config is missing {split}_jsonl/jsonl/root")
        return Path(file_path).expanduser().resolve()
    return Path(dataset).expanduser().resolve() / f"{split}_metadata.jsonl"


def _first_sample(config, sample_id: str | None = None) -> dict[str, Any]:
    dataset = config.dataset
    replacements = list(dataset.get("path_replacements") or []) if hasattr(dataset, "get") else []
    file_path = _dataset_file(dataset, "test")
    with file_path.open("r", encoding="utf-8") as handle:
        for line in handle:
            if not line.strip():
                continue
            metadata = _relocate_metadata(json.loads(line), replacements)
            if sample_id and metadata.get("sample_id") != sample_id:
                continue
            input_images = metadata.get("input_images") or []
            gt_path = metadata.get("gt_image") or metadata.get("output_image")
            if input_images and gt_path and Path(input_images[0]).exists() and Path(gt_path).exists():
                return {"metadata": metadata, "dataset_file": str(file_path)}
    raise RuntimeError(f"No usable sample found in {file_path} for sample_id={sample_id!r}")


def _image_stats(path: Path) -> dict[str, Any]:
    image = Image.open(path).convert("RGB")
    arr = np.asarray(image)
    flat = arr.reshape(-1)
    counts = Counter(flat.tolist())
    mode_value, mode_count = counts.most_common(1)[0]
    return {
        "path": str(path),
        "sha256": _sha256(path),
        "size": list(image.size),
        "mode": image.mode,
        "min": int(arr.min()),
        "max": int(arr.max()),
        "mean": float(arr.mean()),
        "std": float(arr.std()),
        "pixel_mode_value": int(mode_value),
        "pixel_mode_count": int(mode_count),
    }


def _to_rgb_pil(image):
    from scripts.train_omnigen import _to_rgb_pil as train_to_rgb_pil

    return train_to_rgb_pil(image)


def _save_contact_sheet(input_path: Path, output_path: Path, gt_path: Path, sheet_path: Path) -> None:
    panels = [
        ("Input", Image.open(input_path).convert("RGB")),
        ("Output", Image.open(output_path).convert("RGB")),
        ("GT", Image.open(gt_path).convert("RGB")),
    ]
    target_w = max(image.width for _, image in panels)
    target_h = max(image.height for _, image in panels)
    gap = 12
    title_h = 24
    canvas = Image.new("RGB", (target_w * 3 + gap * 2, target_h + title_h), "white")
    draw = ImageDraw.Draw(canvas)
    x = 0
    for label, image in panels:
        draw.text((x, 4), label, fill=(0, 0, 0))
        canvas.paste(image.resize((target_w, target_h), Image.Resampling.BILINEAR), (x, title_h))
        x += target_w + gap
    canvas.save(sheet_path)


def _lora_fingerprint(path: Path) -> dict[str, Any]:
    from safetensors.torch import load_file

    model_path = path / "adapter_model.safetensors"
    config_path = path / "adapter_config.json"
    result = {
        "path": str(path),
        "adapter_model": _file_record(model_path),
        "adapter_config": _file_record(config_path),
        "num_keys": None,
        "first_10_tensors": [],
    }
    if model_path.is_file():
        tensors = load_file(str(model_path), device="cpu")
        result["num_keys"] = len(tensors)
        for name in sorted(tensors)[:10]:
            tensor = tensors[name]
            result["first_10_tensors"].append(
                {"name": name, "shape": list(tensor.shape), "dtype": str(tensor.dtype)}
            )
    return result


def _hf_snapshot_fingerprint(model_root: Path) -> dict[str, Any]:
    cache_root = model_root.parent.parent if model_root.parent.name == "snapshots" else None
    snapshots = []
    if cache_root is not None:
        snapshots_dir = cache_root / "snapshots"
        if snapshots_dir.is_dir():
            snapshots = sorted(path.name for path in snapshots_dir.iterdir() if path.is_dir())
    key_names = [
        "config.json",
        "model.safetensors.index.json",
        "special_tokens_map.json",
        "tokenizer.json",
        "tokenizer_config.json",
        "vae/config.json",
        "vae/diffusion_pytorch_model.safetensors",
    ]
    files = []
    for name in key_names:
        path = model_root / name
        if path.exists():
            files.append(_file_record(path))
    return {
        "resolved_snapshot_path": str(model_root),
        "snapshot_commit_id": model_root.name if model_root.parent.name == "snapshots" else None,
        "all_snapshots": snapshots,
        "multiple_snapshots": len(snapshots) > 1,
        "key_files": files,
    }


def _import_paths() -> dict[str, Any]:
    import diffusers
    import transformers
    import OmniGen
    from OmniGen import OmniGenPipeline, OmniGenScheduler
    from flow_grpo.omnigen_patch import omnigen_pipeline_with_logprob

    return {
        "OmniGen.__file__": getattr(OmniGen, "__file__", None),
        "OmniGenPipeline": inspect.getfile(OmniGenPipeline),
        "OmniGenScheduler": inspect.getfile(OmniGenScheduler),
        "pipeline_with_logprob": inspect.getfile(omnigen_pipeline_with_logprob.pipeline_with_logprob),
        "pipeline_with_logprob_unwrapped": inspect.getfile(inspect.unwrap(omnigen_pipeline_with_logprob.pipeline_with_logprob)),
        "diffusers.__file__": getattr(diffusers, "__file__", None),
        "transformers.__file__": getattr(transformers, "__file__", None),
    }


def _runtime_env() -> dict[str, Any]:
    return {
        "python_executable": sys.executable,
        "python_version": sys.version,
        "conda_default_env": os.environ.get("CONDA_DEFAULT_ENV"),
        "conda_prefix": os.environ.get("CONDA_PREFIX"),
        "torch_version": torch.__version__,
        "cuda_available": torch.cuda.is_available(),
        "torch_cuda_version": torch.version.cuda,
        "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
        "gpu_model": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        "device_capability": list(torch.cuda.get_device_capability(0)) if torch.cuda.is_available() else None,
        "allow_tf32_matmul": torch.backends.cuda.matmul.allow_tf32,
        "allow_tf32_cudnn": torch.backends.cudnn.allow_tf32,
        "cudnn_benchmark": torch.backends.cudnn.benchmark,
        "bf16_supported": torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
        "packages": _package_versions(),
        "env": {
            key: os.environ.get(key)
            for key in [
                "HF_HOME",
                "HF_HUB_CACHE",
                "PYTHONPATH",
                "OMNIGEN_CODE_ROOT",
                "SFT_LORA_PATH",
                "EVAL_LORA_PATH",
                "DATASET_ROOT",
                "TRAIN_JSONL",
                "TEST_JSONL",
            ]
        },
    }


def _repo_state(gen_root: Path) -> dict[str, Any]:
    return {
        "flow_grpo_cxr": {
            "cwd": str(REPO_ROOT),
            "head": _run(["git", "rev-parse", "HEAD"], REPO_ROOT),
            "status": _run(["git", "status", "--short"], REPO_ROOT),
            "diff_name_only": _run(["git", "diff", "--name-only"], REPO_ROOT),
            "targeted_diff": _run(["git", "diff", "--", *FLOW_DIFF_TARGETS], REPO_ROOT),
        },
        "gen_joint": {
            "cwd": str(gen_root),
            "head": _run(["git", "rev-parse", "HEAD"], gen_root),
            "status": _run(["git", "status", "--short"], gen_root),
            "diff_name_only": _run(["git", "diff", "--name-only"], gen_root),
            "targeted_diff": _run(["git", "diff", "--", *GEN_DIFF_TARGETS], gen_root),
        },
    }


def _key_hashes(gen_root: Path) -> dict[str, Any]:
    hashes = {}
    for rel in KEY_FLOW_FILES:
        hashes[f"flow_grpo_cxr/{rel}"] = _file_record(REPO_ROOT / rel)
    for rel in KEY_GEN_FILES:
        hashes[f"gen_joint/{rel}"] = _file_record(gen_root / rel)
    return hashes


def _copy_image(src: Path, dst: Path) -> dict[str, Any]:
    dst.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(src, dst)
    return _image_stats(dst)


def _generate_one(config, metadata: dict[str, Any], output_path: Path, *, eval_lora_path: Path | None) -> dict[str, Any]:
    from peft import PeftModel
    from scripts.train_omnigen import load_omnigen_components, merge_lora_into_base_model
    from flow_grpo.omnigen_patch.omnigen_pipeline_with_logprob import pipeline_with_logprob

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available; refusing to generate a golden A5500 eval image on CPU.")
    device = torch.device("cuda")
    weight_dtype = torch.bfloat16
    model, vae, processor = load_omnigen_components(config, device, weight_dtype)
    merge_lora_path = getattr(config.train, "merge_lora_path", None)
    if merge_lora_path:
        model = merge_lora_into_base_model(model, merge_lora_path, weight_dtype, trainable=False)
    if eval_lora_path is not None:
        model = PeftModel.from_pretrained(model, str(eval_lora_path), is_trainable=False)
        if hasattr(model, "set_adapter"):
            model.set_adapter("default")
        model.to(dtype=weight_dtype)
    model.eval()
    input_images = metadata.get("input_images") or []
    instruction = metadata.get("instruction")
    if not instruction:
        instruction = f"<img><|image_1|></img> {metadata['prompt']}"
    with torch.no_grad():
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            collected = pipeline_with_logprob(
                model,
                vae,
                processor,
                [instruction],
                [input_images],
                height=config.resolution,
                width=config.resolution,
                num_inference_steps=config.sample.eval_num_steps,
                guidance_scale=config.sample.eval_guidance_scale,
                img_guidance_scale=config.sample.eval_img_guidance_scale,
                max_input_image_size=config.sample.max_input_image_size,
                use_img_guidance=config.sample.use_img_guidance,
                use_input_image_size_as_output=config.sample.use_input_image_size_as_output,
                dtype=weight_dtype,
                output_type="pt",
                noise_level=getattr(config.sample, "noise_level", 0.0),
                sde_type=config.sample.sde_type,
            )
    image = collected["images"].float().cpu().numpy()[0]
    pil = _to_rgb_pil(image)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    pil.save(output_path, format="PNG")
    return _image_stats(output_path)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", default=str(DEFAULT_OUTPUT_DIR))
    parser.add_argument("--config", default=DEFAULT_CONFIG_ENTRY)
    parser.add_argument("--omnigen-code-root", default=str(DEFAULT_OMNIGEN_CODE_ROOT))
    parser.add_argument("--sft-lora-path", default=str(DEFAULT_SFT_LORA))
    parser.add_argument("--eval-lora-path", default=str(DEFAULT_RL_LORA))
    parser.add_argument("--sample-id", default=None)
    parser.add_argument("--skip-rl", action="store_true")
    parser.add_argument("--skip-generation", action="store_true", help="Collect static fingerprint only; do not generate images.")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    gen_root = Path(args.omnigen_code_root).expanduser().resolve()
    os.environ.setdefault("OMNIGEN_CODE_ROOT", str(gen_root))
    os.environ.setdefault("SFT_LORA_PATH", str(Path(args.sft_lora_path).expanduser().resolve()))
    if str(gen_root) not in sys.path:
        sys.path.insert(0, str(gen_root))

    config = _load_config(args.config)
    config.train.merge_lora_path = str(Path(args.sft_lora_path).expanduser().resolve())
    sample_bundle = _first_sample(config, args.sample_id)
    metadata = sample_bundle["metadata"]
    input_path = Path((metadata.get("input_images") or [])[0]).expanduser().resolve()
    gt_path = Path(metadata.get("gt_image") or metadata.get("output_image")).expanduser().resolve()

    from scripts.train_omnigen import resolve_model_root

    model_root = Path(resolve_model_root(config.pretrained.model)).resolve()
    imports = _import_paths()

    sample_id = metadata.get("sample_id") or gt_path.stem
    sample_dir = output_dir / str(sample_id)
    input_copy = sample_dir / "input.png"
    gt_copy = sample_dir / "gt.png"
    sft_output = sample_dir / "generated_sft_only.png"
    rl_output = sample_dir / "generated_checkpoint_190.png"
    sft_sheet = sample_dir / "contact_sheet_sft_only.png"
    rl_sheet = sample_dir / "contact_sheet_checkpoint_190.png"

    input_stats = _copy_image(input_path, input_copy)
    gt_stats = _copy_image(gt_path, gt_copy)
    sft_stats = None
    if not args.skip_generation:
        sft_stats = _generate_one(config, metadata, sft_output, eval_lora_path=None)
        _save_contact_sheet(input_copy, sft_output, gt_copy, sft_sheet)

    rl_stats = None
    eval_lora_path = Path(args.eval_lora_path).expanduser().resolve() if args.eval_lora_path else None
    if not args.skip_generation and not args.skip_rl and eval_lora_path is not None and eval_lora_path.exists():
        rl_stats = _generate_one(config, metadata, rl_output, eval_lora_path=eval_lora_path)
        _save_contact_sheet(input_copy, rl_output, gt_copy, rl_sheet)

    fingerprint = {
        "config_entry": args.config,
        "repo_state": _repo_state(gen_root),
        "key_file_hashes": _key_hashes(gen_root),
        "runtime_env": _runtime_env(),
        "import_paths": imports,
        "sft_lora": _lora_fingerprint(Path(args.sft_lora_path).expanduser().resolve()),
        "checkpoint_190_lora": _lora_fingerprint(eval_lora_path) if eval_lora_path else None,
        "hf_snapshot": _hf_snapshot_fingerprint(model_root),
        "dataset": {
            "config": _plain(config.dataset),
            "test_file": sample_bundle["dataset_file"],
        },
        "single_sample": {
            "sample_id": sample_id,
            "prompt": metadata.get("prompt"),
            "instruction": metadata.get("instruction"),
            "metadata": metadata,
            "resolved_input_path": str(input_path),
            "resolved_gt_path": str(gt_path),
            "input_copy": input_stats,
            "gt_copy": gt_stats,
            "noise_level": getattr(config.sample, "noise_level", None),
            "guidance_scale": getattr(config.sample, "eval_guidance_scale", None),
            "img_guidance_scale": getattr(config.sample, "eval_img_guidance_scale", None),
            "eval_num_steps": getattr(config.sample, "eval_num_steps", None),
            "sde_type": getattr(config.sample, "sde_type", None),
            "sft_only_output": sft_stats,
            "sft_contact_sheet": _file_record(sft_sheet) if sft_stats else None,
            "checkpoint_190_output": rl_stats,
            "checkpoint_190_contact_sheet": _file_record(rl_sheet) if rl_stats else None,
        },
        "generation_status": {
            "skip_generation": bool(args.skip_generation),
            "cuda_available_at_runtime": torch.cuda.is_available(),
            "note": "Generated image fields are null when skip_generation is true.",
        },
    }
    fingerprint_path = output_dir / "fingerprint.json"
    fingerprint_path.write_text(json.dumps(fingerprint, indent=2, sort_keys=True), encoding="utf-8")
    print(json.dumps(fingerprint, indent=2, sort_keys=True))
    print(f"Wrote {fingerprint_path}")
    return 0


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