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

import json
import os
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image

from flow_grpo.radiomics_reward import select_mask_roi


DEFAULT_TEST_JSONL = "/home/wenting/zr/wt_dataset/LIDC_IDRI/anno/cxr_synth_anno_mask_test.jsonl"


def _read_one_jsonl(path: Path) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                return json.loads(line)
    raise RuntimeError(f"No samples found in {path}")


def _load_mask_npz(path: Path, key: str = "mask") -> np.ndarray:
    with np.load(path) as data:
        selected_key = key if key in data.files else data.files[0]
        return data[selected_key]


def _image_summary(path: Path) -> dict[str, Any]:
    with Image.open(path) as image:
        return {"path": str(path), "mode": image.mode, "size": list(image.size)}


def main() -> int:
    jsonl_path = Path(os.environ.get("TEST_JSONL") or os.environ.get("TRAIN_JSONL") or DEFAULT_TEST_JSONL)
    sample = _read_one_jsonl(jsonl_path)
    output_image = Path(sample["output_image"])
    output_mask = Path(sample["output_mask"])
    mask = _load_mask_npz(output_mask, os.environ.get("MASK_NPZ_KEY", "mask"))

    roi_lung, lung_channels = select_mask_roi(mask, mask_channels=[0, 1], mode="union")
    roi_all, all_channels = select_mask_roi(mask, mask_channels="all", mode="union")
    result = {
        "jsonl": str(jsonl_path),
        "sample_id": sample.get("sample_id"),
        "patient_id": sample.get("patient_id"),
        "input_images": sample.get("input_images"),
        "gt_output_image": _image_summary(output_image),
        "gt_output_mask": str(output_mask),
        "mask_array_shape": list(mask.shape),
        "mask_array_dtype": str(mask.dtype),
        "channels_0_1": lung_channels,
        "roi_pixels_channels_0_1": int(roi_lung.sum()),
        "channels_all": all_channels,
        "roi_pixels_all": int(roi_all.sum()),
        "generated_mask_required": False,
    }
    print(json.dumps(result, indent=2, sort_keys=True))
    if int(roi_lung.sum()) <= 0:
        raise SystemExit("Selected [0, 1] ROI is empty.")
    if int(roi_all.sum()) <= 0:
        raise SystemExit("Selected all-channel ROI is empty.")
    return 0


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