import os
import sys
from pathlib import Path

APP_DIR = Path(__file__).resolve().parent
DATA_DIR = APP_DIR / "urbex_data"
CITY_ROOT = DATA_DIR / "cityrefer"
CACHE_ROOT = CITY_ROOT / "verify_region_cache" / "cityrefer"
# The released environment reads these at import time (see notebooks/urbex_colab.ipynb).
os.environ["CITY_AGENT_DATA"] = str(CITY_ROOT)
os.environ["CITY_AGENT_DATA_STPLS3D"] = str(DATA_DIR / "stpls3d")
os.environ["CITY_TOOL_REQUIRE_CACHE"] = "1"
sys.path.insert(0, str(APP_DIR))

import spaces  # noqa: E402  (must precede torch)
import torch  # noqa: E402
import gradio as gr  # noqa: E402

import json  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
from concurrent.futures import ThreadPoolExecutor  # noqa: E402

from PIL import Image, ImageDraw  # noqa: E402
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration  # noqa: E402

from scripts.data.hub_assets import fetch, load_manifest, unpack  # noqa: E402

MODEL_ID = "vvangfaye/Urbex-4B"
DATASET_ID = "vvangfaye/Urbex-Data"
MAX_NEW_TOKENS = 1024
CONTEXT_LIMIT = 16384

# --------------------------------------------------------------------------
# 1. Download the released CityRefer test scenes (BEV + cached oblique views)
# --------------------------------------------------------------------------
t0 = time.time()
manifest, revision = load_manifest(DATASET_ID, DATA_DIR)
city = manifest["datasets"]["cityrefer"]
for entry in city["files"]:
    fetch(DATASET_ID, entry, DATA_DIR, revision)

SCENES = sorted(city["scenes"])


def _download_scene(scene):
    entry = city["scenes"][scene]
    for file in entry["files"]:
        fetch(DATASET_ID, file, DATA_DIR, revision)
    if (DATA_DIR / entry["extract_prefix"] / ".urbex-download.json").is_file():
        return scene
    archive = fetch(DATASET_ID, entry["archive"], DATA_DIR, revision, install=False)
    unpack(archive, DATA_DIR, entry["extract_prefix"], entry["archive"])
    try:  # the extracted copy is all we need
        Path(archive).resolve().unlink()
    except OSError:
        pass
    return scene


with ThreadPoolExecutor(max_workers=6) as pool:
    for s in pool.map(_download_scene, SCENES):
        print("scene ready:", s, flush=True)
print(f"Data ready in {time.time() - t0:.0f}s", flush=True)

from roll.pipeline.agentic.env.city_grounding import CityGroundingEnv  # noqa: E402
from roll.pipeline.agentic.env.city_grounding.render import _load_bev_cached  # noqa: E402
from scripts.demo.view_3d import rebuild_scene, viewer_html  # noqa: E402

# Released test annotations: used for scene-wise example queries and to report
# BEV IoU against ground truth when the query matches a benchmark sample.
GT = {}
for split in ("test_ND", "test_NO"):
    for line in (CITY_ROOT / "meta" / f"cityrefer_{split}.jsonl").read_text().splitlines():
        if line.strip():
            s = json.loads(line)
            GT[(s["scene_id"], s["query"].strip())] = s["gt_bbox"]


def scene_label(scene):
    town, _, idx = scene.partition("_block_")
    return f"{town.capitalize()} · block {idx}"


SCENE_CHOICES = [(scene_label(s), s) for s in SCENES]

# Small previews of each global BEV (the full PNGs are ~100 MB each).
PREVIEW_DIR = DATA_DIR / "previews"
PREVIEW_DIR.mkdir(exist_ok=True)
MAPPINGS = {}
for s in SCENES:
    png = CITY_ROOT / "bev_cityrefer" / f"{s}.png"
    MAPPINGS[s] = json.loads(png.with_suffix(".json").read_text())
    out = PREVIEW_DIR / f"{s}.jpg"
    if not out.exists():
        with Image.open(png) as im:
            im.draft("RGB", (2048, 2048))
            im = im.convert("RGB")
            im.thumbnail((2048, 2048), Image.Resampling.LANCZOS)
            im.save(out, quality=88)
print("previews ready", flush=True)


def preview(scene):
    return Image.open(PREVIEW_DIR / f"{scene}.jpg").convert("RGB")


# --------------------------------------------------------------------------
# 2. Load the model (bf16, full precision, eagerly on GPU)
# --------------------------------------------------------------------------
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = Qwen3VLForConditionalGeneration.from_pretrained(
    MODEL_ID, dtype=torch.bfloat16, attn_implementation="sdpa"
).to("cuda").eval()


def collect_images(conversation):
    return [item["image"] for message in conversation
            if isinstance(message.get("content"), list)
            for item in message["content"]
            if isinstance(item, dict) and item.get("type") == "image"]


def _fence(text):
    return "```text\n" + text.replace("```", "'''").strip() + "\n```"


def _draw_norm_box(img, bbox, color=(255, 30, 30), width=4):
    img = img.copy()
    w, h = img.size
    x0, y0, x1, y1 = [float(v) / 1000.0 for v in bbox]
    ImageDraw.Draw(img).rectangle([x0 * w, y0 * h, x1 * w, y1 * h], outline=color, width=width)
    return img


def _draw_world_boxes(img, mapping, boxes):
    """Draw world-XY boxes [(box, colour, width), ...] onto a (downscaled) global BEV."""
    img = img.copy()
    W, H = img.size
    sx = W / (mapping["x_max"] - mapping["x_min"])
    sy = H / (mapping["y_max"] - mapping["y_min"])
    d = ImageDraw.Draw(img)
    for (x0, y0, x1, y1), colour, width in boxes:
        d.rectangle([(x0 - mapping["x_min"]) * sx, (mapping["y_max"] - y1) * sy,
                     (x1 - mapping["x_min"]) * sx, (mapping["y_max"] - y0) * sy],
                    outline=colour, width=width)
    return img


# --------------------------------------------------------------------------
# 3. Agent loop (matches the reference Colab: greedy, 8 turns, 1024 tokens)
# --------------------------------------------------------------------------
def _gpu_duration(scene_id, query, max_turns, *args, **kwargs):
    # Measured on the live Space: ~4-5 s per agent turn (4-turn runs took 14-19 s).
    return 10 + 6 * int(max_turns)


@spaces.GPU(duration=_gpu_duration)
def _agent_gpu(scene_id, query, max_turns):
    sample = {"dataset": "cityrefer", "scene_id": scene_id, "query": query,
              "bev_png": f"bev_cityrefer/{scene_id}.png"}
    gt = GT.get((scene_id, query.strip()))
    if gt is not None:
        sample["gt_bbox"] = gt
    with tempfile.NamedTemporaryFile("w", suffix=".jsonl", delete=False) as f:
        f.write(json.dumps(sample) + "\n")
        data_path = f.name
    env = CityGroundingEnv(
        data_path=data_path, mode="val", seed=42, multiview_cache_dir=str(CACHE_ROOT),
        max_tool_calls=max_turns - 1, max_assistant_turns=max_turns,
        display_max_side=768, bbox_norm_max=1000.0,
    )
    messages, reset_info = env.reset(seed=0)
    os.unlink(data_path)
    tools = reset_info["tools"]
    torch.manual_seed(42)

    log = [f"**Query:** {query}\n\nThe agent starts from the global bird's-eye view of `{scene_id}`."]
    gallery = [(env._env.view_image("global"), "global BEV (initial observation)")]
    final_info = {}
    t_start = time.perf_counter()
    yield {"log": "\n\n".join(log), "gallery": list(gallery), "done": False}

    for turn in range(1, max_turns + 1):
        prompt = processor.apply_chat_template(
            messages, tools=tools, tokenize=False, add_generation_prompt=True)
        images = collect_images(messages)
        inputs = processor(text=[prompt], images=images or None, padding=True,
                           return_tensors="pt").to(model.device)
        n_prompt = inputs.input_ids.shape[1]
        if n_prompt + MAX_NEW_TOKENS > CONTEXT_LIMIT:
            log.append("⚠️ Conversation exceeded the released 16K context; stopping.")
            break
        with torch.inference_mode():
            out = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS,
                                 do_sample=False, use_cache=True)
        action = processor.tokenizer.decode(out[0, n_prompt:], skip_special_tokens=False)
        del inputs, out
        shown = action.replace("<|im_end|>", "").replace("<|endoftext|>", "")
        log.append(f"### Turn {turn} · model\n" + _fence(shown))

        prior = len(messages)
        messages, _reward, terminated, truncated, final_info = env.step(action)
        for m in messages[prior:]:
            if m.get("role") != "tool":
                continue
            content = m.get("content", [])
            name = m.get("name", "tool")
            if isinstance(content, str):
                log.append(f"**Tool ({name}):**\n" + _fence(content))
                continue
            for item in content:
                if item.get("type") == "text":
                    log.append(f"**Tool ({name}):**\n" + _fence(item["text"]))
                elif item.get("type") == "image":
                    caption = f"turn {turn} · {name}"
                    if final_info.get("view_id") and name == "zoom_in":
                        caption += f" → view {final_info['view_id']}"
                    elif final_info.get("direction"):
                        caption += f" ({final_info['direction']})"
                    gallery.append((item["image"], caption))
        yield {"log": "\n\n".join(log), "gallery": list(gallery), "done": False}
        if terminated or truncated:
            break

    elapsed = time.perf_counter() - t_start
    answer = final_info.get("city_answer")
    world_bbox = final_info.get("city_world_bbox")
    result = {"scene_id": scene_id, "query": query, "answer": answer,
              "world_bbox_xyxy_m": world_bbox, "turns": env._assistant_turns,
              "agent_seconds": round(elapsed, 1)}
    zooms = []
    if answer and world_bbox:
        view_id = answer.get("view_id", "global")
        view = env._env.view_registry.get(view_id)
        if view is not None:
            gallery.append((_draw_norm_box(view["image"], answer["bbox_2d"]),
                            f"final answer on view {view_id}"))
        vid = view_id
        while vid and vid != "global" and vid in env._env.view_registry:
            v = env._env.view_registry[vid]
            zooms.insert(0, v["world_bounds"])
            vid = v["parent_id"]
        if gt is not None:
            result["bev_iou_vs_ground_truth"] = round(float(final_info.get("city_iou_2d", 0.0)), 4)
    else:
        result["reason"] = final_info.get("city_reason", "turn limit reached")
    print(f"[urbex] scene={scene_id} turns={result['turns']} agent_seconds={elapsed:.1f}", flush=True)
    yield {"log": "\n\n".join(log), "gallery": list(gallery), "done": True,
           "result": result, "zooms": zooms}


def ground(scene_id: str, query: str, max_turns: int = 8):
    """Ground a natural-language description to an object in a city-scale 3D scene.

    The Urbex-4B agent looks at the scene's bird's-eye view, then actively searches it by
    looking up landmarks, zooming into regions and inspecting cached oblique 3D views, and
    finally commits a box.

    Args:
        scene_id: CityRefer/SensatUrban scene id, e.g. "cambridge_block_3".
        query: description of the object (car, building, parking lot, ground area...).
        max_turns: maximum assistant turns (the last turn is reserved for the answer).

    Returns:
        Search trajectory (markdown), gallery of observed views, global BEV with the
        prediction, interactive 3D view (HTML) and a JSON result with the world-space box.
    """
    if scene_id not in MAPPINGS:
        raise gr.Error(f"Unknown scene {scene_id!r}")
    query = (query or "").strip()
    if not query:
        raise gr.Error("Please describe the object to find.")
    max_turns = int(max(2, min(8, max_turns)))
    bev = preview(scene_id)
    placeholder = "<div style='padding:1em;opacity:.7'>The 3D view appears once the agent commits an answer.</div>"
    state = None
    for state in _agent_gpu(scene_id, query, max_turns):
        if not state["done"]:
            yield state["log"], state["gallery"], bev, placeholder, {}
    result = state["result"]
    wb = result.get("world_bbox_xyxy_m")
    overview, html3d = bev, "<div style='padding:1em'>No valid final prediction — try rephrasing the query or using more turns.</div>"
    if wb:
        boxes = [(z, (255, 200, 0), 3) for z in state["zooms"]] + [(wb, (255, 30, 30), 6)]
        overview = _draw_world_boxes(bev, MAPPINGS[scene_id], boxes)
        try:
            scene = rebuild_scene(CACHE_ROOT, scene_id, wb, max_points=200_000)
            html3d = viewer_html(scene, title="Urbex prediction", subtitle=scene_label(scene_id),
                                 query=query, bev=(bev, MAPPINGS[scene_id]), zooms=state["zooms"],
                                 height=520)
        except Exception as e:  # e.g. a box outside the cached area
            html3d = f"<div style='padding:1em'>3D view unavailable: {e}</div>"
    yield state["log"], state["gallery"], overview, html3d, result


# --------------------------------------------------------------------------
# 4. UI
# --------------------------------------------------------------------------
EXAMPLES = [
    ["cambridge_block_3", "A red van that has a black car behind it, on Castle Road that is a cross from a bus stop marking on the other side of that road"],
    ["birmingham_block_12", "A square building with a black flat roof surrounded by cars. It has a big car parked on its left, right and front. It is to the right of Holford drive road."],
    ["birmingham_block_9", "The large gray triangular parking lot next to the large green yard on Aldridge Road."],
]

CSS = """
#col-container { max-width: 1280px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Urbex: Agentic Spatial Grounding") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🏙️ Urbex: Agentic Spatial Grounding for City-Scale 3D Scenes\n"
            "Describe an object in a city-scale 3D scene. The **Urbex-4B** agent searches the scene "
            "actively — landmark lookup, BEV zoom-ins and oblique 3D inspections — before committing a box.  \n"
            "[Paper](https://huggingface.co/papers/2610.18356) · [Code](https://github.com/AMAP-ML/Urbex) · "
            "[Model](https://huggingface.co/vvangfaye/Urbex-4B) · [Data](https://huggingface.co/datasets/vvangfaye/Urbex-Data)"
        )
        with gr.Row():
            with gr.Column(scale=2):
                scene = gr.Dropdown(SCENE_CHOICES, value="cambridge_block_3", label="Scene (CityRefer / SensatUrban)")
                query = gr.Textbox(label="What should the agent find?", lines=3,
                                   placeholder="e.g. The white car parked next to the red-roofed house on Castle Road")
                run = gr.Button("Find it", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    max_turns = gr.Slider(2, 8, value=8, step=1, label="Max assistant turns",
                                          info="The last turn is reserved for the final answer (paper default: 8).")
            with gr.Column(scale=3):
                overview = gr.Image(value=preview("cambridge_block_3"), label="Global BEV (red = prediction, yellow = zooms)",
                                    type="pil", height=420)
        view3d = gr.HTML(label="3D view of the prediction")
        with gr.Row():
            with gr.Column(scale=3):
                gallery = gr.Gallery(label="What the agent saw", columns=4, height=360, object_fit="contain")
            with gr.Column(scale=2):
                result = gr.JSON(label="Result")
        with gr.Accordion("Search trajectory", open=True):
            log = gr.Markdown()

        gr.Examples(examples=EXAMPLES, inputs=[scene, query],
                    outputs=[log, gallery, overview, view3d, result],
                    fn=ground, cache_examples=True, cache_mode="lazy")

    scene.change(preview, inputs=scene, outputs=overview, api_visibility="private")
    run.click(ground, inputs=[scene, query, max_turns],
              outputs=[log, gallery, overview, view3d, result], api_name="ground")
    query.submit(ground, inputs=[scene, query, max_turns],
                 outputs=[log, gallery, overview, view3d, result], api_visibility="private")

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)