from __future__ import annotations

import json
import os
from pathlib import Path

os.environ.setdefault("SPACE_DEFAULT_HF_ONLY", "1")
os.environ.setdefault("DEFAULT_MODEL_NAME", "page_vithplus")
os.environ.setdefault("PRELOAD_DEFAULT_MODEL", "0")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("NO_PROXY", "127.0.0.1,localhost")
os.environ.setdefault("no_proxy", "127.0.0.1,localhost")
os.environ.setdefault("RELEASE_GPU_AFTER_INFERENCE", "1")

import gradio as gr
import numpy as np
import torch
from PIL import Image

try:
    import spaces
except Exception:
    class _SpacesFallback:
        @staticmethod
        def GPU(*_args, **_kwargs):
            def decorator(fn):
                return fn

            return decorator

    spaces = _SpacesFallback()

from app import (
    BBox,
    DEFAULT_MODEL_NAME,
    MODEL_OPTIONS,
    OUT_THRESHOLD,
    PersonSpec,
    _clamp_bbox,
    _downsample_for_inference,
    _draw_combined_visualization,
    _draw_multi_visualization,
    _find_peak,
    _make_head_inputs,
    model_manager,
)


ROOT = Path(__file__).resolve().parent
EXAMPLES_DIR = ROOT / "static" / "examples"
MANIFEST_PATH = EXAMPLES_DIR / "manifest.json"


def _parse_bbox_line(line: str, width: int, height: int, default_color: str) -> PersonSpec:
    fields = [field.strip() for field in line.replace("\t", ",").split(",") if field.strip()]
    if len(fields) < 4:
        raise ValueError(f"bbox needs at least 4 values: {line}")

    values = [float(fields[i]) for i in range(4)]
    color = fields[4] if len(fields) >= 5 else default_color
    if all(0 <= value <= 1 for value in values):
        x1, y1, x2, y2 = values
        values = [x1 * width, y1 * height, x2 * width, y2 * height]

    return PersonSpec(
        bbox=BBox(x1=values[0], y1=values[1], x2=values[2], y2=values[3]),
        color=color,
    )


def _parse_people(bbox_text: str, width: int, height: int) -> list[PersonSpec]:
    colors = [
        "#00ff00",
        "#ff7a00",
        "#00a7ff",
        "#ff4fd8",
        "#ffd400",
        "#8b5cff",
    ]
    lines = [
        line.split("# ", 1)[0].strip()
        for line in bbox_text.replace(";", "\n").splitlines()
    ]
    lines = [line for line in lines if line]
    if not lines:
        raise ValueError("Please enter at least one bbox: x1,y1,x2,y2. Normalized coordinates in [0,1] are supported.")
    return [_parse_bbox_line(line, width, height, colors[idx % len(colors)]) for idx, line in enumerate(lines)]


def _to_pil_image(value) -> Image.Image | None:
    if value is None:
        return None
    if isinstance(value, Image.Image):
        return value.convert("RGB")
    if isinstance(value, np.ndarray):
        return Image.fromarray(value).convert("RGB")
    if isinstance(value, (str, Path)):
        return Image.open(value).convert("RGB")
    return None


def _layer_to_bbox(layer, width: int, height: int, color: str) -> PersonSpec | None:
    if layer is None:
        return None
    if isinstance(layer, Image.Image):
        arr = np.array(layer.convert("RGBA"))
    elif isinstance(layer, np.ndarray):
        arr = layer
        if arr.ndim == 2:
            mask = arr > 0
            ys, xs = np.where(mask)
            if len(xs) == 0:
                return None
            return PersonSpec(bbox=BBox(x1=float(xs.min()), y1=float(ys.min()), x2=float(xs.max() + 1), y2=float(ys.max() + 1)), color=color)
        if arr.shape[-1] == 3:
            mask = np.any(arr[..., :3] > 0, axis=-1)
            ys, xs = np.where(mask)
            if len(xs) == 0:
                return None
            return PersonSpec(bbox=BBox(x1=float(xs.min()), y1=float(ys.min()), x2=float(xs.max() + 1), y2=float(ys.max() + 1)), color=color)
    else:
        return None

    mask = arr[..., 3] > 0 if arr.shape[-1] >= 4 else np.any(arr[..., :3] > 0, axis=-1)
    ys, xs = np.where(mask)
    if len(xs) == 0:
        return None

    x1 = max(0, min(width - 1, int(xs.min())))
    y1 = max(0, min(height - 1, int(ys.min())))
    x2 = max(x1 + 1, min(width, int(xs.max()) + 1))
    y2 = max(y1 + 1, min(height, int(ys.max()) + 1))
    return PersonSpec(bbox=BBox(x1=x1, y1=y1, x2=x2, y2=y2), color=color)


def _image_and_people_from_editor(editor_value, bbox_text: str) -> tuple[Image.Image, list[PersonSpec], str]:
    colors = ["#00ff00", "#ff7a00", "#00a7ff", "#ff4fd8", "#ffd400", "#8b5cff"]

    if isinstance(editor_value, dict):
        image = _to_pil_image(editor_value.get("background")) or _to_pil_image(editor_value.get("composite"))
        if image is None:
            raise gr.Error("Please upload or choose an image first.")

        people: list[PersonSpec] = []
        for idx, layer in enumerate(editor_value.get("layers") or []):
            spec = _layer_to_bbox(layer, image.width, image.height, colors[idx % len(colors)])
            if spec is not None:
                people.append(spec)

        if people:
            return image, people, "editor"
        return image, _parse_people(bbox_text, image.width, image.height), "text"

    image = _to_pil_image(editor_value)
    if image is None:
        raise gr.Error("Please upload or choose an image first.")
    return image, _parse_people(bbox_text, image.width, image.height), "text"


def _run_one(
    pil_image: Image.Image,
    model_name: str,
    spec: PersonSpec,
    show_heatmap: bool,
    show_gaze: bool,
):
    original_width, original_height = pil_image.size
    pil_image, image_scale = _downsample_for_inference(pil_image)
    width, height = pil_image.size

    bbox_obj = spec.bbox
    if image_scale != 1.0:
        bbox_obj = BBox(
            x1=bbox_obj.x1 * image_scale,
            y1=bbox_obj.y1 * image_scale,
            x2=bbox_obj.x2 * image_scale,
            y2=bbox_obj.y2 * image_scale,
        )
    bbox_px = _clamp_bbox(bbox_obj, width, height)
    bbox_norm = [bbox_px[0] / width, bbox_px[1] / height, bbox_px[2] / width, bbox_px[3] / height]

    model, scene_transforms, head_transforms = model_manager.get(model_name)
    scene_inputs = [transform(pil_image).unsqueeze(0).to(model_manager.device) for transform in scene_transforms]
    head_inputs = _make_head_inputs(pil_image, bbox_px, head_transforms)
    if head_inputs is not None:
        head_inputs = [tensor.to(model_manager.device) for tensor in head_inputs]

    with torch.inference_mode():
        output = model({
            "images": scene_inputs,
            "head_images": head_inputs,
            "bboxes": [[tuple(bbox_norm)]],
        })

    heatmap = output["heatmap"][0][0]
    inout_score = float(output["inout"][0][0].detach().cpu().item()) if output.get("inout") is not None else 1.0
    is_out = inout_score < OUT_THRESHOLD
    target_x, target_y, peak_score = _find_peak(heatmap, width, height)
    combined = _draw_combined_visualization(
        pil_image,
        bbox_px,
        bbox_norm,
        heatmap,
        (target_x, target_y),
        is_out,
        show_heatmap=show_heatmap,
        show_gaze=show_gaze,
    )
    return {
        "bbox_px": bbox_px,
        "bbox_norm": bbox_norm,
        "heatmap": heatmap,
        "target_xy": (target_x, target_y),
        "is_out": is_out,
        "result": {
            "bbox": {"x1": bbox_px[0], "y1": bbox_px[1], "x2": bbox_px[2], "y2": bbox_px[3]},
            "target": None if is_out else {"x": target_x, "y": target_y, "score": peak_score},
            "inout_score": inout_score,
            "is_out": is_out,
            "image_size": {"width": width, "height": height},
            "original_image_size": {"width": original_width, "height": original_height},
        },
        "combined": combined,
    }


@spaces.GPU(duration=180)
def run_inference(
    image_editor,
    bbox_text: str,
    model_name: str,
    show_heatmap: bool,
    show_gaze: bool,
):
    try:
        pil_image, people, bbox_source = _image_and_people_from_editor(image_editor, bbox_text)
        if model_name not in MODEL_OPTIONS:
            raise gr.Error(f"Unsupported model: {model_name}")

        if len(people) == 1:
            item = _run_one(pil_image, model_name, people[0], show_heatmap, show_gaze)
            return item["combined"], json.dumps({
                "model_name": model_name,
                "device": str(model_manager.device),
                "bbox_source": bbox_source,
                **item["result"],
            }, ensure_ascii=False, indent=2)

        downsampled, _ = _downsample_for_inference(pil_image)
        width, height = downsampled.size

        bbox_px_list = []
        bbox_norm_list = []
        heatmaps = []
        target_xy_list = []
        is_out_list = []
        colors = []
        results = []

        for idx, spec in enumerate(people):
            item = _run_one(pil_image, model_name, spec, False, show_gaze)
            bbox_px_list.append(item["bbox_px"])
            bbox_norm_list.append(item["bbox_norm"])
            heatmaps.append(item["heatmap"])
            target_xy_list.append(item["target_xy"])
            is_out_list.append(item["is_out"])
            colors.append(spec.color)
            results.append({"index": idx, "color": spec.color, **item["result"]})

        combined = _draw_multi_visualization(
            downsampled,
            bbox_px_list,
            bbox_norm_list,
            heatmaps,
            target_xy_list,
            is_out_list,
            colors,
            show_heatmap=show_heatmap,
            show_gaze=show_gaze,
        )
        return combined, json.dumps({
            "model_name": model_name,
            "device": str(model_manager.device),
            "bbox_source": bbox_source,
            "people": results,
            "image_size": {"width": width, "height": height},
            "original_image_size": {"width": pil_image.width, "height": pil_image.height},
        }, ensure_ascii=False, indent=2)
    finally:
        if os.environ.get("RELEASE_GPU_AFTER_INFERENCE", "1").lower() in {"1", "true", "yes"}:
            model_manager.move_to_cpu()


def _load_examples() -> list[list[str]]:
    if not MANIFEST_PATH.exists():
        return []
    items = json.loads(MANIFEST_PATH.read_text())
    examples = []
    for item in items[:8]:
        image_path = ROOT / item["image"].lstrip("/")
        if image_path.exists():
            examples.append([
                str(image_path),
                "0.35,0.15,0.60,0.55",
                DEFAULT_MODEL_NAME,
                True,
                True,
            ])
    return examples


def build_demo() -> gr.Blocks:
    model_choices = [
        (cfg.get("display_name", name), name)
        for name, cfg in MODEL_OPTIONS.items()
    ]

    with gr.Blocks(title="CrossGaze ZeroGPU Demo") as demo:
        gr.Markdown(
            """
            # CrossGaze / PaGE ZeroGPU Demo

            Upload an image, enter one or more head/face bounding boxes, and run gaze target inference.
            BBox format: `x1,y1,x2,y2` per line. If all four values are in `[0,1]`, they are treated as normalized coordinates.
            Optional fifth value sets color, e.g. `0.1,0.2,0.3,0.5,#ff7a00`.
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.ImageEditor(
                    type="pil",
                    label="Input image / draw head region",
                    image_mode="RGBA",
                    sources=["upload", "clipboard"],
                    brush=gr.Brush(default_size=24),
                    eraser=gr.Eraser(default_size=24),
                    layers=True,
                )
                bbox_input = gr.Textbox(
                    label="BBox coordinates (fallback / precise input)",
                    value="0.35,0.15,0.60,0.55",
                    lines=5,
                    placeholder="0.35,0.15,0.60,0.55\n0.12,0.20,0.30,0.48,#ff7a00",
                )
                model_input = gr.Dropdown(
                    choices=model_choices,
                    value=DEFAULT_MODEL_NAME,
                    label="Model",
                )
                show_heatmap = gr.Checkbox(value=True, label="Show heatmap")
                show_gaze = gr.Checkbox(value=True, label="Show gaze point and line")
                run_button = gr.Button("Run inference", variant="primary")

            with gr.Column(scale=1):
                output_image = gr.Image(type="pil", label="Result")
                output_json = gr.Code(label="Result JSON", language="json")

        gr.Examples(
            examples=_load_examples(),
            inputs=[image_input, bbox_input, model_input, show_heatmap, show_gaze],
            label="Examples",
        )

        run_button.click(
            run_inference,
            inputs=[image_input, bbox_input, model_input, show_heatmap, show_gaze],
            outputs=[output_image, output_json],
        )

    return demo


demo = build_demo()


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