"""Tag registry — loads parquet datasets and builds search indices."""

import json
import logging
import re
import threading
from collections import defaultdict
from pathlib import Path

logger = logging.getLogger("BooruGacha")

DATA_DIR = Path(__file__).parent / "data"

# Matches the outermost (tag:FLOAT) weight wrapper.
# Greedy .+ swallows inner parens so tags like kousaka_kirino_(cosplay)
# are captured intact: (kousaka_kirino_(cosplay):1.2) works correctly.
_WEIGHT_RE = re.compile(r"^\((.+):(\d+\.?\d*)\)$")


def parse_weighted_tag(text: str) -> tuple[str, float | None]:
    """Parse ``(tag:1.2)`` ComfyUI weight syntax.

    Returns *(raw_tag, weight)*.  Weight is ``None`` when the input
    has no weight wrapper.
    """
    text = text.strip()
    m = _WEIGHT_RE.match(text)
    if m:
        return m.group(1).strip(), float(m.group(2))
    return text, None


def format_weighted_tag(tag: str, weight: float | None) -> str:
    """Wrap *tag* in ``(tag:weight)`` syntax when weight is not 1.0."""
    if weight is not None and abs(weight - 1.0) > 1e-4:
        # Show minimal decimals: 1.2 not 1.20, but 1.05 keeps two
        wt = f"{weight:.2f}".rstrip("0").rstrip(".")
        if "." not in wt:
            wt += ".0"
        return f"({tag}:{wt})"
    return tag


def load_category_schema(data_dir: Path = DATA_DIR) -> dict:
    """Load categorization_schema.json and return the category tree.

    Returns a dict like::

        {
            "top_groups": ["artist", "character", ...],
            "subgroups": {
                "Hair": ["hair_color", "hair_length", ...],
                "Face": ["eye_color", ...],
                ...
            },
        }
    """
    schema_path = data_dir / "categorization_schema.json"
    with open(schema_path, "r", encoding="utf-8") as f:
        schema = json.load(f)

    top_groups: list[str] = []
    subgroups: dict[str, list[str]] = {}

    for group_name, group_info in schema.get("groups", {}).items():
        subs = group_info.get("subgroups")
        if subs is None:
            # Flat group (artist, character, etc.)
            top_groups.append(group_name)
        else:
            for sub_name, sub_info in subs.items():
                categories = list(sub_info.get("categories", {}).keys())
                if categories:
                    subgroups[sub_name] = categories

    return {"top_groups": top_groups, "subgroups": subgroups}


def tokenize(text: str) -> list[str]:
    """Split text into lowercase tokens on common separators."""
    return [
        t
        for t in text.lower()
        .replace("-", "_")
        .replace("/", "_")
        .replace(" ", "_")
        .split("_")
        if t
    ]


def flatten(text: str) -> str:
    """Remove all separators for whole-string comparison."""
    return (
        text.lower()
        .replace("_", "")
        .replace("-", "")
        .replace("/", "")
        .replace(" ", "")
        .replace("(", "")
        .replace(")", "")
    )


class TagRegistry:
    """Loads categorized Danbooru tags and builds fast lookup indices."""

    def __init__(self, data_dir: Path = DATA_DIR):
        self.data_dir = data_dir

        # Core storage — one entry per usable tag
        self.tags: list[dict] = []
        self.name_to_idx: dict[str, int] = {}
        self.danbooru_to_idx: dict[str, int] = {}

        # Search indices
        self.token_index: dict[str, set[int]] = defaultdict(set)
        self.flat_index: dict[str, list[int]] = defaultdict(list)

        # Category → tag indices (for gacha roll)
        self.by_category: dict[str, list[int]] = defaultdict(list)

        self._load()

    # ------------------------------------------------------------------

    def _load(self):
        import pyarrow.parquet as pq

        # --- categorized_tags (the curated ~55k tags) -----------------
        cat_table = pq.read_table(self.data_dir / "categorized_tags.parquet")
        cat_tags = cat_table.column("tag").to_pylist()
        cat_groups = cat_table.column("group").to_pylist()
        cat_subgroups = cat_table.column("subgroup").to_pylist()
        cat_categories = cat_table.column("category").to_pylist()
        cat_filtering = cat_table.column("filtering").to_pylist()

        cat_lookup: dict[str, tuple] = {}
        for i in range(len(cat_tags)):
            cat_lookup[cat_tags[i]] = (
                cat_groups[i],
                cat_subgroups[i],
                cat_categories[i],
                cat_filtering[i],
            )

        # --- unified_tags (danbooru name mapping) ---------------------
        uni_table = pq.read_table(
            self.data_dir / "unified_tags.parquet",
            columns=["name", "danbooru_name", "danbooru_post_count", "is_deprecated"],
        )
        names = uni_table.column("name").to_pylist()
        db_names = uni_table.column("danbooru_name").to_pylist()
        post_counts = uni_table.column("danbooru_post_count").to_pylist()
        deprecated = uni_table.column("is_deprecated").to_pylist()

        for i in range(len(names)):
            db_name = db_names[i]
            if db_name is None or deprecated[i]:
                continue
            name = names[i]
            if name not in cat_lookup:
                continue

            group, subgroup, category, filtering = cat_lookup[name]
            idx = len(self.tags)
            self.tags.append(
                {
                    "name": name,
                    "danbooru_name": db_name,
                    "group": group,
                    "subgroup": subgroup,
                    "category": category,
                    "filtering": filtering or "sfw",
                    "post_count": post_counts[i] or 0,
                }
            )
            self.name_to_idx[name] = idx
            if db_name not in self.danbooru_to_idx:
                self.danbooru_to_idx[db_name] = idx

            # Category index — fine-grained category + group
            if category:
                self.by_category[category].append(idx)
            self.by_category[group].append(idx)

            # Token index — index both normalized and danbooru forms
            for token in tokenize(name):
                self.token_index[token].add(idx)
            if db_name != name:
                for token in tokenize(db_name):
                    self.token_index[token].add(idx)

            # Flat index
            flat_name = flatten(name)
            self.flat_index[flat_name].append(idx)
            if db_name != name:
                flat_db = flatten(db_name)
                if flat_db != flat_name:
                    self.flat_index[flat_db].append(idx)

        logger.info(f"[BooruGacha] Loaded {len(self.tags)} tags")

    # ------------------------------------------------------------------

    def __len__(self) -> int:
        return len(self.tags)


class CooccurrenceIndex:
    """Lazy-loaded co-occurrence lookup from ``cooccurrence.parquet``."""

    def __init__(self, data_dir: Path = DATA_DIR):
        self._path = data_dir / "cooccurrence.parquet"
        self._table = None
        self._cache: dict[str, list[tuple]] = {}
        self._lock = threading.Lock()

    def _ensure_loaded(self):
        if self._table is not None:
            return
        with self._lock:
            if self._table is not None:
                return
            import pyarrow.parquet as pq

            self._table = pq.read_table(self._path)
            logger.info(
                f"[BooruGacha] Loaded co-occurrence index "
                f"({len(self._table)} pairs)"
            )

    def get_related(
        self, tag: str, limit: int = 100
    ) -> list[tuple[str, float, float, float]]:
        """Return co-occurring tags sorted by Jaccard descending.

        Each entry is ``(cooccur_tag, jaccard, pmi, lmi_score)``.
        """
        if tag in self._cache:
            return self._cache[tag][:limit]

        self._ensure_loaded()
        import pyarrow.compute as pc

        mask = pc.equal(self._table.column("tag"), tag)
        filt = self._table.filter(mask)

        rows: list[tuple[str, float, float, float]] = []
        if len(filt) > 0:
            ct = filt.column("cooccur_tag").to_pylist()
            jc = filt.column("jaccard").to_pylist()
            pm = filt.column("pmi").to_pylist()
            lm = filt.column("lmi_score").to_pylist()
            for i in range(len(ct)):
                rows.append((ct[i], jc[i], pm[i], lm[i]))
            rows.sort(key=lambda r: -r[1])  # Jaccard desc

        self._cache[tag] = rows
        return rows[:limit]
