from __future__ import annotations

import os
import threading
import time
from functools import lru_cache
from typing import Any

import gradio as gr
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

try:
    import spaces
except ImportError:
    class _LocalSpaces:
        @staticmethod
        def GPU(*_args: Any, **_kwargs: Any):
            def decorator(function):
                return function

            return decorator

    spaces = _LocalSpaces()

from histagent import load_pretrained, predict_ranked_genes


MODEL_REPO = "wli13/HistAgent"
BASE_MODEL_REPO = "prov-gigapath/prov-gigapath"
DATA_REPO = "wli13/HistAgent-data"
MODEL_COMMIT = "f93e130"

ORGANS = [
    "Unknown",
    "b16f10 syngeneic tumor",
    "bone",
    "brain",
    "breast",
    "cervix",
    "colon",
    "digit",
    "embryo",
    "endometrium",
    "glioblastoma",
    "glioma",
    "heart",
    "joint",
    "kidney",
    "lacrimal gland",
    "leiomyosarcoma",
    "liver",
    "lung",
    "lymph node",
    "melanoma",
    "mouth",
    "muscle",
    "ovary",
    "pancreas",
    "prostate",
    "skin",
    "spleen",
    "stomach",
    "tendon",
    "thymus",
    "undifferentiated pleomorphic sarcoma",
]

_MODEL_BUNDLE: tuple[Any, Any, Any] | None = None
_MODEL_LOCK = threading.Lock()


@lru_cache(maxsize=1)
def example_images() -> tuple[str | None, str | None]:
    try:
        local_path = hf_hub_download(
            DATA_REPO,
            "tutorials/figure5_he_query_brain_local.png",
            repo_type="dataset",
        )
        context_path = hf_hub_download(
            DATA_REPO,
            "tutorials/figure5_he_query_brain_context.png",
            repo_type="dataset",
        )
        return local_path, context_path
    except Exception:
        return None, None


def _load_model() -> tuple[Any, Any, Any]:
    global _MODEL_BUNDLE
    if _MODEL_BUNDLE is not None:
        return _MODEL_BUNDLE

    with _MODEL_LOCK:
        if _MODEL_BUNDLE is not None:
            return _MODEL_BUNDLE
        token = os.getenv("HF_TOKEN")
        if not token:
            raise RuntimeError(
                "The Space owner must add an HF_TOKEN secret with access to the gated "
                "Prov-GigaPath repository."
            )
        if not torch.cuda.is_available():
            raise RuntimeError("A GPU worker is required for HistAgent inference.")

        torch.set_float32_matmul_precision("high")
        _MODEL_BUNDLE = load_pretrained(
            MODEL_REPO,
            token=token,
            device="cuda",
        )
        return _MODEL_BUNDLE


def _friendly_error(error: Exception) -> str:
    message = str(error).lower()
    if "hf_token" in message or "gated" in message or "403" in message:
        return (
            "The demo cannot access the gated Prov-GigaPath base encoder. "
            "The Space owner needs to enable access to public gated repositories "
            "for the `HF_TOKEN` secret."
        )
    if "cuda" in message or "gpu" in message:
        return "No GPU worker is currently available. Please retry after a short wait."
    return f"Inference failed with {type(error).__name__}. Please retry or check the Space logs."


@spaces.GPU(duration=180)
def generate_ranked_readout(
    local_image: Image.Image | None,
    context_image: Image.Image | None,
    species: str,
    organ: str,
    top_k: int,
    progress=gr.Progress(),
):
    if local_image is None or context_image is None:
        return [], "", {}, "Please provide both a local H&E view and a context H&E view."

    started = time.perf_counter()
    try:
        progress(0.1, desc="Loading HistAgent")
        model, tokenizer, config = _load_model()
        progress(0.55, desc="Generating ranked molecular readout")
        genes = predict_ranked_genes(
            model,
            tokenizer,
            local_image,
            context_image,
            species=species,
            organ=organ,
            top_k=int(top_k),
            device="cuda",
        )
    except Exception as error:
        return [], "", {}, _friendly_error(error)

    elapsed = time.perf_counter() - started
    ranked_rows = [[rank, gene] for rank, gene in enumerate(genes, start=1)]
    metadata = {
        "model": MODEL_REPO,
        "base_encoder": BASE_MODEL_REPO,
        "species": species,
        "organ": organ,
        "genes_generated": len(genes),
        "elapsed_seconds": round(elapsed, 2),
        "input_views": ["local", "context"],
        "input_size_after_preprocessing": "224 × 224 pixels per view",
    }
    sentence = " ".join(genes)
    return (
        ranked_rows,
        sentence,
        metadata,
        f"Generated {len(genes)} ranked genes in {elapsed:.1f} seconds.",
    )


CSS = """
.gradio-container {
  max-width: 1240px !important;
  color: #18312b;
}
.module-note {
  background: #f2f8f6;
  border: 1px solid #d7e5e0;
  border-radius: 12px;
  color: #526b63;
  margin-bottom: 12px;
  padding: 12px 14px;
}
.module-note strong {color: #1f5d52;}
.research-note {
  border-left: 4px solid #2e8578;
  padding: 10px 14px;
  background: #f2f8f6;
  border-radius: 6px;
}
"""


with gr.Blocks(
    title="HistAgent · H&E to ranked molecular readout",
    theme=gr.themes.Soft(
        primary_hue="indigo",
        secondary_hue="orange",
        neutral_hue="slate",
    ),
    css=CSS,
) as demo:
    with gr.Tab("1 · Ranked molecular readout"):
        gr.HTML(
            """
            <div class="module-note">
              <strong>Visual-omics foundation model.</strong>
              Supply paired H&amp;E views centered on the same tissue location.
              HistAgent returns an ordered gene list rather than a continuous
              expression matrix.
            </div>
            """
        )
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                local_input = gr.Image(
                    type="pil",
                    label="Local H&E view",
                    height=300,
                )
                context_input = gr.Image(
                    type="pil",
                    label="Context H&E view",
                    height=300,
                )
            with gr.Column(scale=1):
                with gr.Row():
                    species_input = gr.Dropdown(
                        ["human", "mouse", "unknown"],
                        value="human",
                        label="Species",
                    )
                    organ_input = gr.Dropdown(
                        ORGANS,
                        value="brain",
                        label="Organ",
                        allow_custom_value=False,
                    )
                top_k_input = gr.Slider(
                    minimum=10,
                    maximum=50,
                    step=5,
                    value=50,
                    label="Number of ranked genes",
                )
                run_button = gr.Button(
                    "Generate ranked molecular readout",
                    variant="primary",
                    size="lg",
                )
                status_output = gr.Markdown(
                    "Upload paired views or load the example below.",
                    elem_classes=["research-note"],
                )
                metadata_output = gr.JSON(label="Run information")

        example_local, example_context = example_images()
        if example_local and example_context:
            gr.Examples(
                examples=[[example_local, example_context, "human", "brain", 50]],
                inputs=[
                    local_input,
                    context_input,
                    species_input,
                    organ_input,
                    top_k_input,
                ],
                label="Example: human brain",
                cache_examples=False,
            )

        with gr.Row():
            ranked_output = gr.Dataframe(
                headers=["Rank", "Gene"],
                datatype=["number", "str"],
                label="Ranked genes",
                interactive=False,
                wrap=True,
            )
            sentence_output = gr.Textbox(
                label="Ordered gene sentence",
                lines=12,
                show_copy_button=True,
            )

        run_button.click(
            fn=generate_ranked_readout,
            inputs=[
                local_input,
                context_input,
                species_input,
                organ_input,
                top_k_input,
            ],
            outputs=[
                ranked_output,
                sentence_output,
                metadata_output,
                status_output,
            ],
        )

    with gr.Tab("2 · Evidence-grounded reasoning"):
        gr.HTML(
            """
            <iframe
              src="https://wli13-histagent-chat.hf.space/?view=chat"
              title="HistAgent Chat"
              style="width: 100%; height: 900px; border: 0; background: white;"
              loading="lazy">
            </iframe>
            """
        )

    with gr.Tab("About"):
        gr.Markdown(
            f"""
            ### What this demo runs

            HistAgent uses local and surrounding H&E morphology to autoregressively
            generate an ordered list of genes. The demo loads the released
            [`{MODEL_REPO}`](https://huggingface.co/{MODEL_REPO}) checkpoint and the
            official gated
            [`{BASE_MODEL_REPO}`](https://huggingface.co/{BASE_MODEL_REPO}) encoder.

            ### Input

            - A spot-centred local H&E crop.
            - A broader context crop centred on the same tissue location.
            - Species and organ labels.

            Both images are center-cropped to 224 × 224 pixels during preprocessing.

            ### Output

            The output is an ordered gene list, not a continuous expression matrix.
            Generated readouts are intended for research use and must not be used for
            clinical decision-making without independent validation.

            [GitHub repository](https://github.com/zipging/HistAgent) ·
            [Model card](https://huggingface.co/{MODEL_REPO}) ·
            [Tutorial data](https://huggingface.co/datasets/{DATA_REPO})
            """
        )

demo.queue(default_concurrency_limit=1, max_size=8)


if __name__ == "__main__":
    demo.launch()