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

import argparse
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.dataset_paths import DatasetPathResolver
from flow_grpo.server_profiles import apply_server_profile_defaults


apply_server_profile_defaults()


def load_jsonl(path: Path, num_samples: int) -> list[dict[str, Any]]:
    rows = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                rows.append(json.loads(line))
            if len(rows) >= num_samples:
                break
    return rows


def path_values(sample: dict[str, Any]) -> list[tuple[str, str]]:
    values = []
    for index, value in enumerate(sample.get("input_images") or []):
        values.append((f"input_images[{index}]", value))
    for key in ("output_image", "gt_image", "output_mask", "gt_mask", "mask", "mask_path", "gt_mask_path"):
        if sample.get(key):
            values.append((key, sample[key]))
    return values


def inspect_jsonl(path: Path, resolver: DatasetPathResolver, num_samples: int) -> int:
    print(f"[path-debug] jsonl={path}")
    rows = load_jsonl(path, num_samples)
    if not rows:
        raise RuntimeError(f"No rows found in {path}")

    missing = 0
    for row_index, sample in enumerate(rows):
        print(f"[path-debug] sample={row_index}")
        for label, original in path_values(sample):
            try:
                resolved = resolver.resolve(original, label=label)
                exists = resolved.exists()
            except FileNotFoundError as exc:
                resolved = resolver.resolve(original, required=False, label=label)
                exists = False
                missing += 1
                print(f"  {label}:")
                print(f"    original={original}")
                print(f"    resolved={resolved}")
                print(f"    exists={exists}")
                print(f"    error={exc}")
                continue
            print(f"  {label}:")
            print(f"    original={original}")
            print(f"    resolved={resolved}")
            print(f"    exists={exists}")
            if not exists:
                missing += 1

    print(f"[path-debug] inspected={len(rows)} missing={missing}")
    return missing


def main() -> int:
    parser = argparse.ArgumentParser(description="Debug H20 Ceph dataset JSONL path resolution.")
    parser.add_argument("--jsonl", default=None, help="Single JSONL to inspect. Defaults to TRAIN_JSONL and TEST_JSONL.")
    parser.add_argument("--dataset_root", default=os.environ.get("DATASET_ROOT"))
    parser.add_argument("--remap_from", default=os.environ.get("DATASET_PATH_REMAP_FROM"))
    parser.add_argument("--remap_to", default=os.environ.get("DATASET_PATH_REMAP_TO"))
    parser.add_argument("--num_samples", type=int, default=5)
    args = parser.parse_args()

    resolver = DatasetPathResolver(args.dataset_root, args.remap_from, args.remap_to)
    print(f"[path-debug] SERVER_PROFILE={os.environ.get('SERVER_PROFILE', '')}")
    print(f"[path-debug] DATASET_ROOT={resolver.dataset_root}")
    print(f"[path-debug] DATASET_PATH_REMAP_FROM={resolver.remap_from}")
    print(f"[path-debug] DATASET_PATH_REMAP_TO={resolver.remap_to}")

    jsonl_paths = []
    if args.jsonl:
        jsonl_paths.append(Path(args.jsonl).expanduser().resolve())
    else:
        for env_name in ("TRAIN_JSONL", "TEST_JSONL"):
            value = os.environ.get(env_name)
            if not value:
                raise ValueError(f"{env_name} is not set; pass --jsonl or set SERVER_PROFILE=h20_ceph")
            jsonl_paths.append(Path(value).expanduser().resolve())

    total_missing = 0
    for jsonl_path in jsonl_paths:
        if not jsonl_path.exists():
            raise FileNotFoundError(f"JSONL does not exist: {jsonl_path}")
        total_missing += inspect_jsonl(jsonl_path, resolver, args.num_samples)

    if total_missing:
        raise RuntimeError(f"Dataset path debug found {total_missing} missing required paths.")
    print("[path-debug] dataset path resolution OK")
    return 0


if __name__ == "__main__":
    sys.exit(main())