#!/usr/bin/env python3
import importlib.metadata as md
import importlib.util
import os
import sys
from pathlib import Path

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()

omnigen_code_root = os.environ.get("OMNIGEN_CODE_ROOT")
if omnigen_code_root and os.path.exists(omnigen_code_root) and omnigen_code_root not in sys.path:
    sys.path.insert(0, omnigen_code_root)

print("[preflight] python:", sys.executable)
print("[preflight] cwd:", os.getcwd())
print("[preflight] SERVER_PROFILE:", os.environ.get("SERVER_PROFILE"))
print("[preflight] OMNIGEN_CODE_ROOT:", os.environ.get("OMNIGEN_CODE_ROOT"))
print("[preflight] SFT_LORA_PATH:", os.environ.get("SFT_LORA_PATH"))
print("[preflight] DATASET_ROOT:", os.environ.get("DATASET_ROOT"))
print("[preflight] TRAIN_JSONL:", os.environ.get("TRAIN_JSONL"))
print("[preflight] TEST_JSONL:", os.environ.get("TEST_JSONL"))
print("[preflight] DATASET_PATH_REMAP_FROM:", os.environ.get("DATASET_PATH_REMAP_FROM"))
print("[preflight] DATASET_PATH_REMAP_TO:", os.environ.get("DATASET_PATH_REMAP_TO"))
print("[preflight] OmniGen spec:", importlib.util.find_spec("OmniGen"))

for package in [
    "torch",
    "xformers",
    "transformers",
    "diffusers",
    "accelerate",
    "datasets",
    "peft",
    "timm",
    "tokenizers",
    "safetensors",
    "itk",
    "pydicom",
]:
    try:
        print(f"[preflight] {package}: {md.version(package)}")
    except Exception:
        print(f"[preflight] {package}: MISSING")

errors = []

try:
    from diffusers.models import AutoencoderKL
    print("[preflight] AutoencoderKL import OK")
except Exception as exc:
    errors.append(f"AutoencoderKL import failed: {exc!r}")
    print("[preflight] AutoencoderKL import FAILED:", repr(exc))

try:
    from transformers import Phi3Config, Phi3Model
    print("[preflight] Phi3Config/Phi3Model import OK")
except Exception as exc:
    errors.append(f"Phi3Config/Phi3Model import failed: {exc!r}")
    print("[preflight] Phi3Config/Phi3Model import FAILED:", repr(exc))

try:
    from OmniGen import OmniGen, OmniGenProcessor
    print("[preflight] OmniGen import OK")
except Exception as exc:
    errors.append(f"OmniGen import failed: {exc!r}")
    print("[preflight] OmniGen import FAILED:", repr(exc))

try:
    import itk
    print("[preflight] itk import OK:", itk.Version.GetITKVersion())
except Exception as exc:
    errors.append(f"itk import failed: {exc!r}")
    print("[preflight] itk import FAILED:", repr(exc))

try:
    import pydicom
    print("[preflight] pydicom import OK")
except Exception as exc:
    errors.append(f"pydicom import failed: {exc!r}")
    print("[preflight] pydicom import FAILED:", repr(exc))

mirp_repo_root = os.environ.get("MIRP_REPO_ROOT", "/data/wenting/mirp")
if os.path.exists(mirp_repo_root) and mirp_repo_root not in sys.path:
    sys.path.insert(0, mirp_repo_root)
print("[preflight] MIRP_REPO_ROOT:", mirp_repo_root)

try:
    from mirp import extract_features
    from mirp._images.digital_xray_image import DXImage
    from mirp._masks.base_mask import BaseMask
except Exception as exc:
    errors.append(f"MIRP import failed: {exc!r}")
    print("[preflight] MIRP import FAILED:", repr(exc))
    print(f"[preflight] Try: pip install -e {mirp_repo_root}")
else:
    print("[preflight] MIRP import OK")

if errors:
    print("[preflight] FAILED checks:")
    for error in errors:
        print(f"[preflight] - {error}")
    raise SystemExit(1)