import os
import uuid
from pathlib import Path

os.environ.setdefault("nnUNet_raw", "/tmp")
os.environ.setdefault("nnUNet_preprocessed", "/tmp")
os.environ.setdefault("nnUNet_results", "/tmp")

import gradio as gr

import baseline_infer
from baseline_infer import DATASET_DISPATCH


APP_ROOT = Path(__file__).parent.resolve()
MODEL_ROOT = APP_ROOT / "baseline_models"
EXAMPLE_ROOT = APP_ROOT / "AutoMSC_examples"
OUTPUT_ROOT = APP_ROOT / "gradio_output"
MAX_CHANNELS = 6

# Keep baseline_infer aligned with the bundled Space repository layout.
baseline_infer.BASELINE_ROOT = str(MODEL_ROOT)


DATASET_CONFIGS = {
    "Dataset001_PETWB_Lung": {
        "label": "Dataset001 — PETWB Lung",
        "title": "PETWB Lung",
        "description": "Lung segmentation + cancer type classification (Other / Lung Cancer)",
        "modalities": ["CT", "PET"],
    },
    "Dataset002_BMLMPS_FLAIR": {
        "label": "Dataset002 — BMLMPS FLAIR",
        "title": "BMLMPS FLAIR",
        "description": "Whole-tumor segmentation + EGFR status classification (Wild-Type / Mutation)",
        "modalities": ["FLAIR"],
    },
    "Dataset003_BMLMPS_T1CE": {
        "label": "Dataset003 — BMLMPS T1CE",
        "title": "BMLMPS T1CE",
        "description": "Core-tumor segmentation + EGFR status classification (Wild-Type / Mutation)",
        "modalities": ["T1CE"],
    },
    "Dataset004_BrainMets": {
        "label": "Dataset004 — BrainMets",
        "title": "BrainMets",
        "description": "Brain metastasis segmentation (necrotic/enhancing/edema) + primary tumor origin classification",
        "modalities": ["T1", "T1c", "T2", "FLAIR", "CT", "RTP"],
    },
    "Dataset005_MU_Glioma_Post": {
        "label": "Dataset005 — MU Glioma Post",
        "title": "MU Glioma Post",
        "description": "Brain glioma segmentation (NCR/ED/ET/NET_RC) + primary diagnosis classification (GBM / Astrocytoma / Others)",
        "modalities": ["t1c", "t1n", "t2f", "t2w"],
    },
    "Dataset006_JSC_UCSD_PTGB": {
        "label": "Dataset006 — JSC UCSD Post-Tx GBM",
        "title": "JSC UCSD Post-Tx GBM",
        "description": "Tumor segmentation + IDH mutation status classification (Wild-Type / Mutant)",
        "modalities": ["T1post", "FLAIR", "ADC"],
    },
    "Dataset007_PICAI": {
        "label": "Dataset007 — PICAI Prostate",
        "title": "PICAI Prostate",
        "description": "Prostate segmentation + ISUP grade classification (Grade 0-5)",
        "modalities": ["T2W", "ADC", "HBV"],
    },
    "Dataset008_PETWB_Liver": {
        "label": "Dataset008 — PETWB Liver",
        "title": "PETWB Liver",
        "description": "Liver segmentation + cancer type classification (Other / Liver Cancer)",
        "modalities": ["CT", "PET"],
    },
    "Dataset009_LUNA25": {
        "label": "Dataset009 — LUNA25",
        "title": "LUNA25",
        "description": "Pulmonary nodule segmentation + malignancy classification (Benign / Malignant)",
        "modalities": ["CT"],
    },
}

LABEL_TO_DATASET = {config["label"]: name for name, config in DATASET_CONFIGS.items()}


try:
    import spaces  # type: ignore
except ImportError:
    spaces = None


def gpu_task(duration: int = 600):
    if spaces is not None:
        return spaces.GPU(duration=duration)

    def _identity(fn):
        return fn

    return _identity


def dataset_label(dataset_name: str) -> str:
    return DATASET_CONFIGS[dataset_name]["label"]


def dataset_name_from_label(label: str) -> str:
    return LABEL_TO_DATASET.get(label, label)


def available_dataset_names() -> list[str]:
    found = [
        dataset_name
        for dataset_name in DATASET_CONFIGS
        if (MODEL_ROOT / dataset_name).exists()
    ]
    return found or list(DATASET_CONFIGS)


def dataset_markdown(dataset_name: str) -> str:
    config = DATASET_CONFIGS[dataset_name]
    modalities = ", ".join(config["modalities"])
    return (
        f"# {config['label']}\n"
        f"{config['description']}  \n"
        f"Upload {len(config['modalities'])} channel(s): **{modalities}**"
    )


def channel_label(idx: int, modality: str) -> str:
    return f"Channel {idx} — {modality}  (.nii / .nii.gz)"


def file_input_updates(dataset_name: str):
    modalities = DATASET_CONFIGS[dataset_name]["modalities"]
    updates = []
    for idx in range(MAX_CHANNELS):
        if idx < len(modalities):
            modality = modalities[idx]
            updates.append(
                gr.update(
                    label=channel_label(idx, modality),
                    visible=True,
                    value=None,
                )
            )
        else:
            updates.append(gr.update(visible=False, value=None))
    return updates


def example_search_dirs(dataset_name: str) -> list[Path]:
    return [
        EXAMPLE_ROOT / dataset_name / "imagesTr",
        EXAMPLE_ROOT / dataset_name / "examples" / "imagesTr",
        MODEL_ROOT / dataset_name / "examples" / "imagesTr",
    ]


def strip_nifti_suffix(path: Path) -> str:
    name = path.name
    for suffix in (".nii.gz", ".nii"):
        if name.endswith(suffix):
            return name[: -len(suffix)]
    return path.stem


def discover_examples(dataset_name: str) -> list[list[str]]:
    n_channels = len(DATASET_CONFIGS[dataset_name]["modalities"])
    for folder in example_search_dirs(dataset_name):
        if not folder.exists():
            continue

        cases: dict[str, dict[int, Path]] = {}
        for file_path in sorted(folder.glob("*.nii*")):
            stem = strip_nifti_suffix(file_path)
            if len(stem) < 5 or stem[-5] != "_" or not stem[-4:].isdigit():
                continue
            channel = int(stem[-4:])
            if channel >= n_channels:
                continue
            case_id = stem[:-5]
            cases.setdefault(case_id, {})[channel] = file_path

        examples = []
        for case_id in sorted(cases):
            channels = cases[case_id]
            if all(idx in channels for idx in range(n_channels)):
                examples.append([str(channels[idx]) for idx in range(n_channels)])
        if examples:
            return examples

    return []


def example_button_updates(dataset_name: str):
    examples = discover_examples(dataset_name)
    return [
        gr.update(visible=len(examples) >= 1),
        gr.update(visible=len(examples) >= 2),
    ]


def on_dataset_change(dataset_label_value: str):
    dataset_name = dataset_name_from_label(dataset_label_value)
    return (
        gr.update(value=dataset_markdown(dataset_name)),
        *file_input_updates(dataset_name),
        *example_button_updates(dataset_name),
    )


def load_example(dataset_label_value: str, index: int):
    dataset_name = dataset_name_from_label(dataset_label_value)
    examples = discover_examples(dataset_name)
    if len(examples) < index:
        raise gr.Error(f"Example {index} not found for {dataset_label(dataset_name)}.")

    example_files = examples[index - 1]
    outputs = []
    for idx in range(MAX_CHANNELS):
        outputs.append(example_files[idx] if idx < len(example_files) else None)
    return outputs


def format_classification_results(cls_results: dict) -> str:
    if not cls_results:
        return "No classification results were returned."

    md_parts = []
    for task_name, probs in cls_results.items():
        md = f"### {task_name}\n\n| Class | Probability |\n|-------|-------------|\n"
        for cls_name, prob in sorted(probs.items(), key=lambda item: -item[1]):
            md += f"| {cls_name} | {prob:.4f} |\n"
        md_parts.append(md)
    return "\n\n".join(md_parts)


@gpu_task(duration=600)
def run_gpu_inference(dataset_name: str, files: list[str], output_dir: str):
    import torch

    device = "cuda" if torch.cuda.is_available() else "cpu"
    infer_fn = DATASET_DISPATCH[dataset_name]
    return infer_fn(image=files, output_dir=output_dir, device=device)


def run_inference(dataset_label_value: str, *inputs):
    dataset_name = dataset_name_from_label(dataset_label_value)
    config = DATASET_CONFIGS[dataset_name]
    modalities = config["modalities"]
    files = list(inputs[: len(modalities)])

    missing = [modality for modality, file_path in zip(modalities, files) if file_path is None]
    if missing:
        raise gr.Error(f"Missing files: {', '.join(missing)}")

    if not (MODEL_ROOT / dataset_name).exists():
        raise gr.Error(
            f"Model folder not found: {MODEL_ROOT / dataset_name}. "
            "Bundle the matching DatasetXXX_* folder under baseline_models/."
        )

    output_dir = OUTPUT_ROOT / dataset_name.lower() / str(uuid.uuid4())
    output_dir.mkdir(parents=True, exist_ok=True)

    seg_path, video_path, cls_results = run_gpu_inference(
        dataset_name,
        [str(file_path) for file_path in files],
        str(output_dir),
    )
    return video_path, seg_path, format_classification_results(cls_results)


def build_ui(available_names: list[str] | None = None):
    available_names = available_names or available_dataset_names()
    labels = [dataset_label(name) for name in available_names]
    single = len(labels) == 1
    default_name = available_names[0]
    default_label = dataset_label(default_name)

    with gr.Blocks(title="AutoMSC Baseline Inference") as demo:
        header = gr.Markdown(dataset_markdown(default_name))

        with gr.Row():
            with gr.Column(scale=1):
                dataset_dropdown = gr.Dropdown(
                    choices=labels,
                    value=default_label,
                    label="Model",
                    interactive=not single,
                )
                gr.Markdown("**Upload NIfTI file(s)**")

                file_inputs = []
                default_modalities = DATASET_CONFIGS[default_name]["modalities"]
                for idx in range(MAX_CHANNELS):
                    file_inputs.append(
                        gr.File(
                            label=channel_label(idx, default_modalities[idx]) if idx < len(default_modalities) else f"Channel {idx}",
                            file_types=[".gz", ".nii"],
                            type="filepath",
                            visible=idx < len(default_modalities),
                        )
                    )

                default_examples = discover_examples(default_name)
                with gr.Row():
                    ex1_btn = gr.Button("Load Example 1", variant="secondary", visible=len(default_examples) >= 1)
                    ex2_btn = gr.Button("Load Example 2", variant="secondary", visible=len(default_examples) >= 2)
                    run_btn = gr.Button("Run Inference", variant="primary")

            with gr.Column(scale=1):
                video_out = gr.Video(label="Overlay Video", format="mp4", height=420, width="100%")
                seg_out = gr.File(label="Download Segmentation Mask (.nii.gz)")
                cls_out = gr.Markdown(label="Classification Results")

        dataset_dropdown.change(
            fn=on_dataset_change,
            inputs=[dataset_dropdown],
            outputs=[header, *file_inputs, ex1_btn, ex2_btn],
        )
        ex1_btn.click(
            fn=lambda label: load_example(label, 1),
            inputs=[dataset_dropdown],
            outputs=file_inputs,
        )
        ex2_btn.click(
            fn=lambda label: load_example(label, 2),
            inputs=[dataset_dropdown],
            outputs=file_inputs,
        )
        run_btn.click(
            fn=run_inference,
            inputs=[dataset_dropdown, *file_inputs],
            outputs=[video_out, seg_out, cls_out],
        )

    return demo


if __name__ == "__main__":
    labels = [dataset_label(name) for name in available_dataset_names()]
    print(f"[startup] available AutoMSC models: {labels}")

    demo = build_ui()
    port = int(os.environ.get("GRADIO_SERVER_PORT", 7860))
    demo.launch(
        server_name="0.0.0.0",
        server_port=port,
        share=False,
        theme=gr.themes.Soft(),
        ssr_mode=False,
        show_error=True,
        mcp_server=True
    )