#!/usr/bin/env python3
"""HyperView Space runtime for core-claims top jaguar ReID models."""

from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Any

import numpy as np
from datasets import Dataset as HFDataset
from datasets import DatasetDict as HFDatasetDict
from datasets import load_dataset, load_from_disk
import hyperview as hv
from hyperview.core.sample import Sample

SPACE_HOST = os.environ.get("SPACE_HOST", "0.0.0.0")
LOCAL_BIND_HOSTS = {"0.0.0.0", "127.0.0.1", "localhost", "::", "::1"}

DATASET_NAME = os.environ.get("HYPERVIEW_DATASET_NAME", "jaguar_core_claims_demo")
HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "hyper3labs/jaguar-hyperview-demo")
HF_DATASET_CONFIG = os.environ.get("HF_DATASET_CONFIG", "default")
HF_DATASET_SPLIT = os.environ.get("HF_DATASET_SPLIT", "train")

EMBEDDING_ASSET_DIR = Path(
    os.environ.get(
        "EMBEDDING_ASSET_DIR",
        str((Path(__file__).resolve().parent / "assets").resolve()),
    )
)
ASSET_MANIFEST_PATH = Path(
    os.environ.get("EMBEDDING_ASSET_MANIFEST", str((EMBEDDING_ASSET_DIR / "manifest.json").resolve()))
)


def _patch_hyperview_default_panel() -> None:
    """Patch HyperView 0.3.1 frontend for default panel and dock cache-key migration.

    HyperView currently has no public API for these behaviors. This runtime patch is
    intentionally narrow and idempotent, targeting the known bundled chunk for v0.3.1.
    """
    default_panel = os.environ.get("HYPERVIEW_DEFAULT_PANEL", "spherical3d").strip().lower()
    apply_default_panel_patch = default_panel in {"spherical3d", "sphere3d"}
    if not apply_default_panel_patch:
        print(f"Skipping frontend default-panel patch (HYPERVIEW_DEFAULT_PANEL={default_panel!r}).")

    cache_version = os.environ.get("HYPERVIEW_LAYOUT_CACHE_VERSION", "v6").strip() or "v6"
    target_layout_key = f"hyperview:dockview-layout:{cache_version}"
    legacy_layout_key = "hyperview:dockview-layout:v5"
    layout_key_pattern = r"hyperview:dockview-layout:v\d+"

    chunk_path = (
        Path(hv.__file__).resolve().parent
        / "server"
        / "static"
        / "_next"
        / "static"
        / "chunks"
        / "077b38561d6ea80d.js"
    )
    if not chunk_path.exists():
        print(f"Default-panel patch skipped: chunk not found at {chunk_path}")
        return

    marker_before = 'v||(v=n)};if(f&&l&&w({id:dr,title:"Euclidean"'
    marker_after = 'v||(v=n),t.id===dd&&n.api.setActive()};if(f&&l&&w({id:dr,title:"Euclidean"'

    try:
        payload = chunk_path.read_text(encoding="utf-8")
    except OSError as exc:
        print(f"Default-panel patch skipped: failed reading chunk ({exc})")
        return

    patched = payload
    changed = False

    if apply_default_panel_patch:
        if marker_after in patched:
            print("HyperView frontend already patched for Sphere 3D default panel.")
        elif marker_before in patched:
            patched = patched.replace(marker_before, marker_after, 1)
            changed = True
            print("Patched HyperView frontend: Sphere 3D will open as default scatter panel.")
        else:
            print("Default-panel patch skipped: expected marker not found in HyperView chunk.")

    if target_layout_key in patched:
        print(f"HyperView frontend already uses dock cache key '{target_layout_key}'.")
    elif legacy_layout_key in patched:
        patched = patched.replace(legacy_layout_key, target_layout_key, 1)
        changed = True
        print(f"Patched HyperView frontend: dock cache key {legacy_layout_key} -> {target_layout_key}.")
    else:
        discovered = re.search(layout_key_pattern, patched)
        if discovered:
            source_key = discovered.group(0)
            if source_key == target_layout_key:
                print(f"HyperView frontend already uses dock cache key '{target_layout_key}'.")
            else:
                print(
                    f"Dock cache patch notice: expected legacy key '{legacy_layout_key}' not found; "
                    f"migrating detected key '{source_key}' -> '{target_layout_key}'."
                )
                patched = patched.replace(source_key, target_layout_key, 1)
                changed = True
        else:
            print(
                "Dock cache patch warning: expected layout cache key marker "
                f"'{legacy_layout_key}' not found in HyperView chunk."
            )

    if not changed:
        return

    try:
        chunk_path.write_text(patched, encoding="utf-8")
    except OSError as exc:
        print(f"Frontend patch skipped: failed writing chunk ({exc})")


def _resolve_bind_host() -> tuple[str, str | None]:
    explicit_bind = os.environ.get("HYPERVIEW_BIND_HOST")
    if explicit_bind:
        return explicit_bind, None

    if SPACE_HOST in LOCAL_BIND_HOSTS:
        return SPACE_HOST, None

    return "0.0.0.0", f"SPACE_HOST='{SPACE_HOST}' is non-local; falling back to 0.0.0.0"


def _resolve_port() -> int:
    for key in ("SPACE_PORT", "PORT"):
        value = os.environ.get(key)
        if value:
            try:
                return int(value)
            except ValueError as exc:
                raise ValueError(f"Invalid integer value for {key}: {value}") from exc
    return 7860


def load_asset_manifest(path: Path) -> dict[str, Any]:
    if not path.exists():
        raise FileNotFoundError(
            f"Embedding asset manifest not found: {path}. "
            "Run scripts/build_hyperview_demo_assets.py first."
        )
    payload = json.loads(path.read_text(encoding="utf-8"))
    if "models" not in payload or not isinstance(payload["models"], list):
        raise ValueError(f"Invalid asset manifest format: {path}")
    return payload


def _load_hf_rows() -> HFDataset:
    repo_path = Path(HF_DATASET_REPO)
    if repo_path.exists():
        loaded = load_from_disk(str(repo_path))
        if isinstance(loaded, HFDatasetDict):
            if HF_DATASET_SPLIT in loaded:
                return loaded[HF_DATASET_SPLIT]
            if "train" in loaded:
                return loaded["train"]
            first_split = next(iter(loaded.keys()))
            return loaded[first_split]
        return loaded
    return load_dataset(HF_DATASET_REPO, name=HF_DATASET_CONFIG, split=HF_DATASET_SPLIT)


def ingest_hf_dataset_samples(dataset: hv.Dataset) -> None:
    rows = _load_hf_rows()
    media_root = Path(os.environ.get("HYPERVIEW_MEDIA_DIR", "./demo_data/media")) / DATASET_NAME
    media_root.mkdir(parents=True, exist_ok=True)

    added = 0
    for index, row in enumerate(rows):
        filename = str(row.get("filename", f"sample_{index:06d}.jpg"))
        sample_id = str(row.get("sample_id", filename))
        if dataset._storage.get_sample(sample_id) is not None:
            continue

        image_obj = row["image"]
        image_path = media_root / f"{Path(sample_id).stem}.jpg"
        if not image_path.exists():
            image_obj.convert("RGB").save(image_path, format="JPEG", quality=90, optimize=True)

        label = str(row.get("label", ""))
        metadata = {
            "filename": filename,
            "sample_id": sample_id,
            "split_tag": str(row.get("split_tag", "unknown")),
            "identity": label,
            "source_repo": HF_DATASET_REPO,
            "source_config": HF_DATASET_CONFIG,
            "source_split": HF_DATASET_SPLIT,
        }

        dataset.add_sample(
            Sample(
                id=sample_id,
                filepath=str(image_path),
                label=label,
                metadata=metadata,
            )
        )
        added += 1

    print(f"Ingested {added} HF samples into HyperView dataset '{DATASET_NAME}'.")


def ensure_embedding_spaces(dataset: hv.Dataset, asset_manifest: dict[str, Any], asset_dir: Path) -> None:
    known_sample_ids = {sample.id for sample in dataset.samples}

    for model in asset_manifest["models"]:
        model_key = str(model["model_key"])
        space_key = str(model["space_key"])
        embeddings_rel = model.get("embeddings_path")
        if not embeddings_rel:
            raise ValueError(f"Missing embeddings_path in asset manifest for model {model_key}")

        embeddings_path = asset_dir / str(embeddings_rel)
        if not embeddings_path.exists():
            raise FileNotFoundError(
                f"Missing embeddings file for model {model_key}: {embeddings_path}"
            )

        payload = np.load(embeddings_path, allow_pickle=False)
        ids = [str(x) for x in payload["ids"].tolist()]
        vectors = np.asarray(payload["vectors"], dtype=np.float32)

        if vectors.ndim != 2:
            raise ValueError(f"Embeddings for {model_key} must be 2D; got {vectors.shape}")
        if len(ids) != vectors.shape[0]:
            raise ValueError(
                f"Embeddings/ID mismatch for {model_key}: {len(ids)} ids vs {vectors.shape[0]} vectors"
            )

        missing_ids = sorted(set(ids) - known_sample_ids)
        if missing_ids:
            preview = ", ".join(missing_ids[:5])
            raise RuntimeError(
                f"Embedding IDs missing from loaded dataset for {model_key}. "
                f"First missing IDs: {preview}"
            )

        config = {
            "provider": "precomputed-checkpoint",
            "geometry": str(model["geometry"]),
            "comparison_key": model.get("comparison_key"),
            "family": model.get("family"),
            "checkpoint_path": model.get("checkpoint_path"),
        }

        dataset._storage.ensure_space(
            model_id=model_key,
            dim=int(vectors.shape[1]),
            space_key=space_key,
            config=config,
        )
        dataset._storage.add_embeddings(space_key, ids, vectors)

        print(f"Ensured space {space_key} ({vectors.shape[0]} x {vectors.shape[1]})")


def ensure_layouts(dataset: hv.Dataset, asset_manifest: dict[str, Any]) -> list[str]:
    layout_keys: list[str] = []
    for model in asset_manifest["models"]:
        space_key = str(model["space_key"])
        layout_spec = str(model.get("layout", "euclidean:2d"))
        layout_key = dataset.compute_visualization(
            space_key=space_key,
            layout=layout_spec,
            method="umap",
            force=False,
        )
        layout_keys.append(layout_key)
        print(f"Ensured layout {layout_key} for space={space_key}")
    return layout_keys


def build_dataset() -> hv.Dataset:
    asset_manifest = load_asset_manifest(ASSET_MANIFEST_PATH)

    dataset = hv.Dataset(DATASET_NAME)
    if len(dataset) == 0:
        print(
            f"Loading HF dataset rows from {HF_DATASET_REPO}[{HF_DATASET_CONFIG}] split={HF_DATASET_SPLIT}"
        )
        ingest_hf_dataset_samples(dataset)

    ensure_embedding_spaces(dataset, asset_manifest=asset_manifest, asset_dir=EMBEDDING_ASSET_DIR)
    layout_keys = ensure_layouts(dataset, asset_manifest=asset_manifest)

    print(f"Dataset '{DATASET_NAME}' has {len(dataset)} samples")
    print(f"Spaces: {[space.space_key for space in dataset.list_spaces()]}")
    print(f"Layouts: {layout_keys}")

    return dataset


def main() -> None:
    _patch_hyperview_default_panel()
    dataset = build_dataset()

    if os.environ.get("HYPERVIEW_DEMO_PREP_ONLY") == "1":
        print("Preparation-only mode enabled; skipping server launch.")
        return

    bind_host, bind_warning = _resolve_bind_host()
    bind_port = _resolve_port()

    if bind_warning:
        print(f"Bind host notice: {bind_warning}")
    print(
        "Starting HyperView with "
        f"bind_host={bind_host} bind_port={bind_port} "
        f"(SPACE_HOST={SPACE_HOST!r}, SPACE_PORT={os.environ.get('SPACE_PORT')!r}, "
        f"PORT={os.environ.get('PORT')!r})"
    )
    hv.launch(dataset, host=bind_host, port=bind_port, open_browser=False)


if __name__ == "__main__":
    main()