"""Original-file LTF inference for Hugging Face Spaces and local use."""

from __future__ import annotations

import hashlib
import json
import logging
import math
import os
from pathlib import Path
import sys
import threading
import time

# Gradio probes its own local HTTP server during startup. Keep local traffic
# direct even on machines configured with outbound HTTP proxies.
for key in ("NO_PROXY", "no_proxy"):
    existing = os.environ.get(key, "")
    os.environ[key] = ",".join(
        filter(None, [existing, "localhost", "127.0.0.1", "::1"])
    )

import gradio as gr
from huggingface_hub import snapshot_download
from PIL import Image, UnidentifiedImageError
import torch

sys.path.insert(0, str(Path(__file__).absolute().parent))
from attention import (
    capture_prediction,
    draw_grid,
    head_view,
    parameter_rows,
    patch_view,
    public_attention,
    tile_boxes,
    tile_overlay,
    top_tiles,
)
from presentation import (
    EMPTY_RESULT,
    HEADER,
    patch_matrix,
    result_card,
    stats_card,
    weight_chart,
)

from sample_images import (
    gallery_items,
    sample_at,
    annotate_example,
    same_image_as_state,
)

from locale_settings import english_app_kwargs

LOG = logging.getLogger(__name__)
MAX_PIXELS = int(os.environ.get("LTF_MAX_PIXELS", "16000000"))
MAX_VIEWS = int(os.environ.get("LTF_MAX_VIEWS", "400"))
MODEL_LOCK = threading.Lock()
_detector = None
_model_revision = None


def validate_upload(path):
    if not path:
        raise gr.Error("Please upload an image first.")
    source = Path(path)
    try:
        with Image.open(source) as image:
            width, height = image.size
            if getattr(image, "n_frames", 1) != 1:
                raise gr.Error("Please upload a single-frame image.")
            views = math.ceil(width / 224) * math.ceil(height / 224) + 1
            if width * height > MAX_PIXELS or views > MAX_VIEWS:
                raise gr.Error(
                    f"Image exceeds demo limits: {MAX_PIXELS:,} pixels "
                    f"and {MAX_VIEWS} views. Please choose a smaller original image."
                )
            image.verify()
    except (UnidentifiedImageError, OSError, ValueError, Image.DecompressionBombError):
        raise gr.Error("This file is not a readable image.") from None
    return source, width, height, views


def resolve_model_dir():
    local_dir = os.environ.get("LTF_MODEL_DIR")
    if local_dir:
        root = Path(local_dir).expanduser().resolve()
        if not root.is_dir():
            raise RuntimeError("LTF_MODEL_DIR does not point to a model directory.")
        return root
    repo = os.environ.get("LTF_MODEL_REPO")
    if not repo:
        raise RuntimeError("Set LTF_MODEL_REPO or LTF_MODEL_DIR before inference.")
    return Path(
        snapshot_download(
            repo_id=repo,
            revision=os.environ.get("LTF_MODEL_REVISION", "main"),
            allow_patterns=[
                "*.json",
                "configs/*",
                "weights/*",
                "backbones/*",
                "ltf_final/*.py",
                "model/*.py",
                "mil_module/*.py",
                "utils/*.py",
            ],
        )
    )


def verify_model(root):
    manifest = json.loads((root / "release_manifest.json").read_text())
    for relative, expected in manifest["files"].items():
        relative_path = Path(relative)
        # Hub snapshot files are symlinks into the sibling blobs directory.
        # Validate the manifest path before following links, then hash the bytes.
        if relative_path.is_absolute() or ".." in relative_path.parts:
            raise RuntimeError("The model manifest contains an invalid path.")
        path = root / relative_path
        if not path.is_file():
            raise RuntimeError("The model release is incomplete.")
        digest = hashlib.sha256()
        with path.open("rb") as stream:
            for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
                digest.update(block)
        if digest.hexdigest() != expected["sha256"]:
            raise RuntimeError(f"Model release checksum failed: {relative}")
    return manifest


def get_detector():
    """Called under MODEL_LOCK, so initialization and inference are serialized."""
    global _detector, _model_revision
    if _detector is None:
        root = resolve_model_dir()
        verify_model(root)
        sys.path.insert(0, str(root))
        from ltf_final import LTFDetector

        torch.set_num_threads(int(os.environ.get("LTF_CPU_THREADS", "2")))
        _detector = LTFDetector(
            config=root / "configs/eval.yaml",
            weights=root / "weights/epoch_20.pth",
            device=os.environ.get("LTF_DEVICE", "auto"),
            chunk_size=int(os.environ.get("LTF_CHUNK_SIZE", "8")),
        )
        _model_revision = os.environ.get("LTF_MODEL_REVISION", "local")
    return _detector


def _predict_request(image_path, branch="fusion", threshold=0.5, with_attention=False):
    source, width, height, views = validate_upload(image_path)
    if branch not in ("fusion", "level1", "level2"):
        raise gr.Error("Unknown model branch.")
    if not math.isfinite(float(threshold)) or not 0 <= float(threshold) <= 1:
        raise gr.Error("Threshold must be between 0 and 1.")
    start = time.perf_counter()
    try:
        with MODEL_LOCK:
            detector = get_detector()
            if with_attention:
                result, attention = capture_prediction(
                    detector, source, branch, float(threshold)
                )
                state = dict(
                    path=str(source),
                    attention=attention,
                    boxes=tile_boxes(
                        width, height, int(detector.model.input_resolution)
                    ),
                    parameters=parameter_rows(
                        detector, branch, float(threshold), views
                    ),
                )
            else:
                result = detector.predict(source, branch, float(threshold))
                state = None
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        raise gr.Error(
            "GPU memory is insufficient. Please try a smaller image."
        ) from None
    except gr.Error:
        raise
    except Exception:
        LOG.exception("LTF inference failed")
        raise gr.Error(
            "Model inference failed. Check the Space runtime logs."
        ) from None
    # The preview never replaces the original image passed to the model.
    with Image.open(source) as original:
        preview = original.convert("RGB")
        preview.thumbnail((1024, 1024))
    result.pop("path", None)
    result.update(
        width=width,
        height=height,
        model_revision=_model_revision,
        elapsed_seconds=round(time.perf_counter() - start, 3),
    )
    assert result["views"] == views, "Preview and model tile counts disagree"
    scores = result["probabilities"][branch]
    summary = result_card(result)
    rows = [
        [
            name,
            result["probabilities"][name]["real"],
            result["probabilities"][name]["fake"],
        ]
        for name in ("fusion", "level1", "level2")
    ]
    if state:
        result["attention"] = public_attention(state)
        state["result"] = dict(result)
    return (
        preview,
        summary,
        {"Real": scores["real"], "AI-generated": scores["fake"]},
        rows,
        result,
    ), state


def predict(image_path, branch="fusion", threshold=0.5):
    """Keep the original five-output public prediction API compatible."""
    return _predict_request(image_path, branch, threshold)[0]


def explain(image_path, branch="fusion", threshold=0.5):
    base, state = _predict_request(image_path, branch, threshold, with_attention=True)
    attention = state["attention"]
    count = len(state["boxes"])
    if len(attention["l1"]) != count or len(attention["l2"]) != count + 1:
        raise gr.Error("Attention and original image tile counts disagree.")
    grid = draw_grid(state["path"], state["boxes"])
    l1_image, l1_description = head_view(state, "mean")
    l2_image = tile_overlay(
        state["path"], state["boxes"], attention["l2"][1:], reference_count=count + 1
    )
    description = (
        f'Global view weight **{attention["l2"][0]*100:.2f}%**; total local view weight '
        f'**{float(attention["l2"][1:].sum())*100:.2f}%**. '
        'The global view is listed separately; the overlay shows local tiles only.\n\n'
        f'Feature fusion coefficients: L1 × {attention["gate"]:.4f}, L2 × {1-attention["gate"]:.4f}. '
        'These coefficients scale branch embeddings, not branch probabilities.'
    )
    weight_rows = [["Global", "whole image", None, float(attention["l2"][0])]]
    weight_rows.extend(
        [
            f"T{i + 1}",
            str(list(box)),
            float(attention["l1"][i]),
            float(attention["l2"][i + 1]),
        ]
        for i, box in enumerate(state["boxes"])
    )
    selected = int(attention["l1"].argmax()) + 1
    tile, patch, patch_description = patch_view(state, str(selected))
    return (
        *base,
        grid,
        l1_image,
        l1_description,
        l2_image,
        description,
        top_tiles(state),
        state["parameters"],
        weight_rows,
        state,
        gr.Dropdown(choices=[str(i + 1) for i in range(count)], value=str(selected)),
        tile,
        patch,
        patch_description,
        gr.Dropdown(
            choices=[("Mean (8 Heads)", "mean")]
            + [
                (f"Head {i + 1}", f"head_{i + 1}")
                for i in range(attention["l1_heads"].shape[0])
            ],
            value="mean",
        ),
        weight_chart(state),
        weight_chart(state, level="l2"),
        patch_matrix(state, str(selected)),
        stats_card(state),
        gr.update(visible=True),
    )


def inspect_head(state, selection):
    image, description = head_view(state, selection)
    return image, description, weight_chart(state, selection)


def inspect_patch(state, tile_number):
    tile, image, description = patch_view(state, tile_number)
    return tile, image, description, patch_matrix(state, tile_number)


def select_top_tile(state, event: gr.SelectData):
    if not state:
        return gr.skip()
    index = event.index[0] if isinstance(event.index, (tuple, list)) else event.index
    order = sorted(
        range(len(state["boxes"])), key=lambda i: -float(state["attention"]["l1"][i])
    )[:4]
    return str(order[int(index)] + 1)


def update_decision(state, branch, threshold):
    """Threshold and branch selection use cached logits, without another forward."""
    if not state:
        return tuple(gr.skip() for _ in range(5))
    result = dict(state["result"])
    score = result["probabilities"][branch]["fake"]
    label = int(score >= float(threshold))
    result.update(
        branch=branch,
        threshold=float(threshold),
        score_fake=score,
        label=label,
        prediction="fake" if label else "real",
    )
    scores = result["probabilities"][branch]
    rows = [
        [
            name,
            result["probabilities"][name]["real"],
            result["probabilities"][name]["fake"],
        ]
        for name in ("fusion", "level1", "level2")
    ]
    parameters = [list(row) for row in state["parameters"]]
    parameters[0][2] = f"{branch} / {float(threshold):.2f}"
    return (
        result_card(result),
        {"Real": scores["real"], "AI-generated": scores["fake"]},
        rows,
        result,
        parameters,
    )


def reset_upload(path):
    image = None
    if path:
        try:
            source, *_ = validate_upload(path)
            with Image.open(source) as original:
                image = original.convert("RGB")
            image.thumbnail((1024, 1024))
        except gr.Error:
            pass
    return image, EMPTY_RESULT, None, [], None, [], None, gr.update(visible=False)


def run_example(branch, threshold, event: gr.SelectData):
    sample, path = sample_at(event.index)
    outputs = explain(str(path), branch, threshold)
    description = annotate_example(outputs, sample)
    return (*outputs, str(path), description)


def on_upload_change(path, current_state):
    if same_image_as_state(path, current_state):
        return tuple(gr.skip() for _ in range(9))
    return (*reset_upload(path), "")


def build_demo():
    with gr.Blocks(
        title="LTF · Image Detection and Attention Analysis",
        fill_width=True,
        delete_cache=(3600, 3600),
    ) as demo:
        gr.HTML(HEADER)
        with gr.Row():
            with gr.Column(scale=3, min_width=260, elem_id="upload-panel"):
                gr.Markdown(
                    "### 01 · Upload an image\nInference uses the original file. Resizing the preview does not change the model input."
                )
                upload = gr.File(
                    label="Image file",
                    type="filepath",
                    file_count="single",
                    file_types=["image"],
                    height=150,
                )
                clear_upload = gr.ClearButton([upload], value="Change image", size="sm")
                branch = gr.Dropdown(
                    [
                        ("Fusion · Combined branches", "fusion"),
                        ("L1 · Patch-level branch", "level1"),
                        ("L2 · View-level branch", "level2"),
                    ],
                    value="fusion",
                    label="Prediction branch",
                    elem_id="branch-select",
                )
                threshold = gr.Slider(
                    0,
                    1,
                    value=0.5,
                    step=0.01,
                    label="Decision threshold",
                    elem_id="threshold-select",
                )
                button = gr.Button(
                    "Run Analysis", variant="primary", elem_id="run-button"
                )
                gr.Markdown(
                    f"Limits: {MAX_PIXELS/1_000_000:g} MP · {MAX_VIEWS} views · 30 MB. Large images take longer on CPU."
                )
            with gr.Column(scale=4, min_width=300):
                preview = gr.Image(
                    label="Original image preview",
                    height=390,
                    interactive=False,
                    buttons=["download", "fullscreen"],
                    elem_id="original-preview",
                )
            with gr.Column(scale=4, min_width=300):
                summary = gr.HTML(EMPTY_RESULT, label="Analysis result")
                scores = gr.Label(
                    label="Model scores", num_top_classes=2, visible=False
                )
        with gr.Column(elem_id="example-section"):
            gr.Markdown(
                "### GenImage examples · Click an image to analyze\nReal / Generated captions are dataset labels. Model predictions and attention appear in the analysis panels."
            )
            example_gallery = gr.Gallery(
                value=gallery_items(),
                columns=4,
                rows=2,
                height="auto",
                label="GenImage validation examples",
                show_label=False,
                object_fit="contain",
                allow_preview=False,
                preview=False,
                interactive=False,
                buttons=[],
                elem_id="example-gallery",
            )
            example_info = gr.Markdown("", elem_id="example-info")
        state = gr.State(value=None, time_to_live=3600)
        with gr.Column(visible=False, elem_id="analysis-section") as analysis:
            statistics = gr.HTML()
            gr.Markdown(
                "## 02 · Attention Analysis\nAll heatmaps use the same uniform-attention reference. Nearly uniform weights have similar colors. Attention is not a forgery probability."
            )
            with gr.Tabs(selected="l1", elem_id="analysis-tabs"):
                with gr.Tab("L1 · Tile Attention", id="l1"):
                    with gr.Row():
                        with gr.Column(scale=7, min_width=320):
                            l1_image = gr.Image(
                                label="L1 tile attention overlay",
                                height=500,
                                interactive=False,
                                buttons=["download", "fullscreen"],
                                elem_classes="attention-image",
                            )
                            l1_description = gr.Markdown()
                        with gr.Column(scale=5, min_width=300):
                            head = gr.Dropdown(
                                [("Mean (8 Heads)", "mean")],
                                value="mean",
                                label="Aggregation head",
                                elem_id="head-select",
                            )
                            l1_chart = gr.HTML()
                with gr.Tab("L2 · View Weights", id="l2"):
                    with gr.Row():
                        with gr.Column(scale=7, min_width=320):
                            l2_image = gr.Image(
                                label="L2 local view attention overlay",
                                height=500,
                                interactive=False,
                                buttons=["download", "fullscreen"],
                                elem_classes="attention-image",
                            )
                            attention_description = gr.Markdown()
                        with gr.Column(scale=5, min_width=300):
                            l2_chart = gr.HTML()
                with gr.Tab("Original · Tile Layout", id="grid"):
                    grid = gr.Image(
                        label="Original 224 × 224 tile boundaries",
                        height=570,
                        interactive=False,
                        buttons=["download", "fullscreen"],
                        elem_classes="attention-image",
                    )
                    gr.Markdown(
                        "Tiles are numbered row by row. Full 224-pixel blocks do not overlap; only the final right/bottom boundary blocks may overlap. Some labels are omitted on dense grids; complete coordinates are listed below."
                    )
                with gr.Tab("Tile · Patch Details", id="patch"):
                    gr.Markdown(
                        "### Top 4 tiles by L1 attention\nClick a tile to inspect its original crop, attention overlay, and patch weights. The top-four selection is for display only and does not affect prediction."
                    )
                    gallery = gr.Gallery(
                        label="High-attention tiles",
                        columns=4,
                        object_fit="contain",
                        height=250,
                        elem_id="top-crops",
                        show_label=False,
                        preview=False,
                        allow_preview=False,
                        interactive=False,
                        buttons=["download_all"],
                    )
                    tile_number = gr.Dropdown(
                        choices=[],
                        label="Tile number",
                        allow_custom_value=False,
                        elem_id="tile-select",
                    )
                    with gr.Row():
                        with gr.Column(min_width=250):
                            selected_tile = gr.Image(
                                label="Original 224 × 224 crop",
                                height=330,
                                interactive=False,
                                buttons=["download", "fullscreen"],
                            )
                        with gr.Column(min_width=250):
                            selected_patch = gr.Image(
                                label="Patch attention overlay",
                                height=330,
                                interactive=False,
                                buttons=["download", "fullscreen"],
                            )
                        with gr.Column(min_width=280):
                            patch_chart = gr.HTML()
                    patch_description = gr.Markdown()
        with gr.Accordion(
            "Exact weights, parameters, and raw output",
            open=False,
            elem_id="technical-details",
        ):
            gr.Markdown("### Branch scores")
            table = gr.Dataframe(
                headers=["Branch", "Real score", "Generated score"],
                datatype=["str", "number", "number"],
                interactive=False,
            )
            gr.Markdown(
                "### View weights\nL1 is the mean across eight heads. The first L2 entry represents the separate global view."
            )
            weight_table = gr.Dataframe(
                headers=["View", "Original crop box", "L1 mean weight", "L2 weight"],
                datatype=["str", "str", "number", "number"],
                interactive=False,
            )
            gr.Markdown(
                "### Model and hyperparameters\nTraining-only rows record the historical configuration. Random augmentation, Cross-Shuffle, and dropout are disabled during inference."
            )
            parameters = gr.Dataframe(
                headers=["Scope", "Parameter", "Value"],
                datatype=["str", "str", "str"],
                interactive=False,
            )
            details = gr.JSON(label="Prediction JSON")
        analysis_outputs = [
            preview,
            summary,
            scores,
            table,
            details,
            grid,
            l1_image,
            l1_description,
            l2_image,
            attention_description,
            gallery,
            parameters,
            weight_table,
            state,
            tile_number,
            selected_tile,
            selected_patch,
            patch_description,
            head,
            l1_chart,
            l2_chart,
            patch_chart,
            statistics,
            analysis,
        ]
        example_event = example_gallery.select(
            run_example,
            [branch, threshold],
            analysis_outputs + [upload, example_info],
            api_name=False,
            concurrency_limit=1,
            concurrency_id="ltf-inference",
            trigger_mode="always_last",
            scroll_to_output=True,
        )
        run_event = button.click(
            explain,
            [upload, branch, threshold],
            analysis_outputs,
            api_name="explain",
            concurrency_limit=1,
            concurrency_id="ltf-inference",
            cancels=[example_event],
        )
        example_gallery.select(fn=None, cancels=[run_event], queue=False)
        upload.upload(fn=None, cancels=[example_event], queue=False)
        clear_upload.click(fn=None, cancels=[example_event], queue=False)
        head.change(
            inspect_head,
            [state, head],
            [l1_image, l1_description, l1_chart],
            api_name=False,
        )
        tile_number.change(
            inspect_patch,
            [state, tile_number],
            [selected_tile, selected_patch, patch_description, patch_chart],
            api_name=False,
        )
        gallery.select(select_top_tile, [state], [tile_number], api_name=False)
        for control in (branch, threshold):
            control.change(
                update_decision,
                [state, branch, threshold],
                [summary, scores, table, details, parameters],
                api_name=False,
                queue=False,
                trigger_mode="always_last",
            )
        upload.change(
            on_upload_change,
            [upload, state],
            [
                preview,
                summary,
                scores,
                table,
                details,
                parameters,
                state,
                analysis,
                example_info,
            ],
            api_name=False,
            queue=False,
            cancels=[run_event],
        )
        api_button = gr.Button(visible=False)
        api_button.click(
            predict,
            [upload, branch, threshold],
            [preview, summary, scores, table, details],
            api_name="predict",
            concurrency_limit=1,
            concurrency_id="ltf-inference",
        )
        gr.Markdown(
            "LTF · Original checkpoint and inference pipeline · Visualization preserves model scores",
            elem_id="footer-note",
        )
    return demo.queue(max_size=16, default_concurrency_limit=1)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    build_demo().launch(
        server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.environ.get("PORT", "7860")),
        share=False,
        max_file_size="30mb",
        app_kwargs=english_app_kwargs(),
        ssr_mode=False,
        footer_links=["api"],
        theme=gr.themes.Soft(
            primary_hue="teal",
            neutral_hue="slate",
            font=["ui-sans-serif", "system-ui", "sans-serif"],
        ).set(block_label_background_fill="#eef3f8", block_label_text_color="#445b72"),
        css_paths=Path(__file__).with_name("styles.css"),
    )