"""The dictionary: which words exist as ready-made assets, and where they live.

Tiers, in fallback order (rejecting an item moves it to the next tier):
    simready -> objaverse -> primitive -> generation
A word is *available* only when a simready or objaverse asset exists for it;
primitive and generation are construction, so they never make a word available.
A missing word can still be used if the caller explicitly maps it to a primitive.
"""
import bisect
import gzip
import json
import os
import re
import threading

import requests
from huggingface_hub import hf_hub_download

CACHE_DIR = os.environ.get("VLI_CACHE", "cache")
TIERS = ["simready", "objaverse", "primitive", "generation"]
ASSET_TIERS = ["simready", "objaverse"]
PRIMITIVE_SHAPES = ["box", "sphere", "cylinder", "cone", "plane"]

SIMREADY_TREE = "https://api.github.com/repos/NVIDIA-Omniverse/simready-central/git/trees/main?recursive=1"
SIMREADY_RAW = "https://raw.githubusercontent.com/NVIDIA-Omniverse/simready-central/main/"
SIMREADY_FALLBACK = ["submissions/nvidia/simready-assets"]
PACKAGE_URL = re.compile(
    r"https://huggingface\.co/datasets/([^/\s]+)/([^/\s]+)/resolve/([0-9a-f]{40})/([^/\s]+)/"
    r"com\.nvidia\.simready\.packaging\.json"
)
VARIANT = re.compile(r"[a-z]?\d+")  # a01, b02, 01


def _split_part(part):
    """'handDryer' -> ['hand', 'dryer']"""
    return re.sub(r"([a-z])([A-Z])", r"\1 \2", part).lower().split()


def simready_phrases(name):
    """Words an asset name like 'interiorPart_sun_visor_a01' can be looked up by.

    Names read category_object..., where the object can span several parts
    ('sun_visor'). Only whole phrases and the final part are indexed, never a
    lone modifier: 'sun' must not find a sun visor.
    """
    parts = [p for p in name.split("_") if p and not VARIANT.fullmatch(p)]
    body = parts[1:] or parts  # drop the category prefix
    out = []
    phrase = " ".join(w for p in body for w in _split_part(p))
    if phrase:
        out.append((phrase, 0))  # 'sun visor', 'hand dryer automatic'
    final = _split_part(body[-1]) if body else []
    if final:
        out.append((" ".join(final), 0))  # 'visor', 'bucky table'
        if len(final) > 1:
            out.append((final[-1], 1))  # 'table' from 'buckyTable': weaker
    return list(dict.fromkeys(out))


def objaverse_word(category):
    """LVIS category 'bow_(weapon)' -> 'bow'"""
    return re.sub(r"_?\(.*?\)", "", category).replace("_", " ").strip().lower()


class Lexicon:
    def __init__(self):
        self._lock = threading.Lock()
        self._index = None  # word -> {tier: [candidate, ...]}
        self._words = []
        self._object_paths = None

    # ---------- building ----------
    def _cached_json(self, name, build):
        path = os.path.join(CACHE_DIR, name)
        if os.path.exists(path):
            with open(path, encoding="utf-8") as f:
                return json.load(f)
        data = build()
        os.makedirs(CACHE_DIR, exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f)
        return data

    def _simready_packages(self):
        try:
            tree = requests.get(SIMREADY_TREE, timeout=30).json()["tree"]
            files = [t["path"] for t in tree if t["type"] == "blob" and t["path"].startswith("submissions/")]
        except Exception:
            files = SIMREADY_FALLBACK
        packages = []
        for f in files:
            try:
                text = requests.get(SIMREADY_RAW + f, timeout=30).text
            except Exception:
                continue
            for ns, ds, rev, name in PACKAGE_URL.findall(text):
                packages.append({"source": "simready", "repo": f"{ns}/{ds}", "revision": rev, "name": name})
        return packages

    def _objaverse_lvis(self):
        path = hf_hub_download("allenai/objaverse", "lvis-annotations.json.gz", repo_type="dataset")
        with gzip.open(path, "rt", encoding="utf-8") as f:
            return json.load(f)

    def _build(self):
        index = {}

        def add(word, tier, cand, rank=0):
            index.setdefault(word, {}).setdefault(tier, []).append((rank, cand))

        for pkg in self._cached_json("simready_packages.json", self._simready_packages):
            for phrase, rank in simready_phrases(pkg["name"]):
                add(phrase, "simready", pkg, rank)
        for category, uids in self._cached_json("objaverse_lvis.json", self._objaverse_lvis).items():
            for uid in uids:
                add(objaverse_word(category), "objaverse", {"source": "objaverse", "uid": uid, "category": category})

        for tiers in index.values():
            for tier, ranked in tiers.items():
                ranked.sort(key=lambda rc: (rc[0], len(rc[1].get("name", ""))))
                tiers[tier] = [c for _, c in ranked]
        self._index = index
        self._words = sorted(index)

    def ensure(self):
        with self._lock:
            if self._index is None:
                self._build()

    # ---------- queries ----------
    def lookup(self, word):
        """{tier: [candidates]} for a word (plural 's' stripped as a fallback)."""
        self.ensure()
        if word in self._index:
            return self._index[word]
        if word.endswith("s") and word[:-1] in self._index:
            return self._index[word[:-1]]
        return {}

    def available(self, word):
        return any(self.lookup(word).get(t) for t in ASSET_TIERS)

    def counts(self, word):
        found = self.lookup(word)
        return {t: len(found.get(t, [])) for t in ASSET_TIERS}

    def complete(self, prefix, limit=12):
        self.ensure()
        i = bisect.bisect_left(self._words, prefix)
        out = []
        while i < len(self._words) and self._words[i].startswith(prefix) and len(out) < limit:
            out.append(self._words[i])
            i += 1
        return out

    def max_phrase_words(self):
        self.ensure()
        return max((len(w.split()) for w in self._words), default=1)

    def object_path(self, uid):
        """Objaverse uid -> 'glbs/000-023/<uid>.glb'"""
        with self._lock:
            if self._object_paths is None:
                path = hf_hub_download("allenai/objaverse", "object-paths.json.gz", repo_type="dataset")
                with gzip.open(path, "rt", encoding="utf-8") as f:
                    self._object_paths = json.load(f)
        return self._object_paths[uid]


LEXICON = Lexicon()