"""
Loads the data-pipeline's JSON output and builds fast lookup indices.
Accepts either the plain .json files or the gzip-compressed .json.gz ones.
"""
import gzip
import json
import os


def _load_json(path):
    gz_path = path + ".gz"
    if os.path.exists(path):
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    if os.path.exists(gz_path):
        with gzip.open(gz_path, "rt", encoding="utf-8") as f:
            return json.load(f)
    raise FileNotFoundError(f"Neither {path} nor {gz_path} exists")


class Database:
    def __init__(self, data_dir):
        self.data_dir = data_dir

        self.enhancements = _load_json(os.path.join(data_dir, "enhancements.json"))
        self.enhancement_sets = _load_json(os.path.join(data_dir, "enhancement_sets.json"))
        self.archetypes = _load_json(os.path.join(data_dir, "archetypes.json"))
        self.powersets = _load_json(os.path.join(data_dir, "powersets.json"))
        self.powers = _load_json(os.path.join(data_dir, "powers.json"))
        self.entities = _load_json(os.path.join(data_dir, "entities.json"))
        self.maths = _load_json(os.path.join(data_dir, "maths_ed_schedules.json"))
        self.enhancement_classes = _load_json(os.path.join(data_dir, "enhancement_classes.json"))

        # --- indices ---
        self.enhancement_by_uid = {e["uid"]: e for e in self.enhancements if e.get("uid")}
        self.enhancement_by_static_index = {
            e["static_index"]: e for e in self.enhancements if e.get("static_index") is not None
        }
        self.enhancement_set_by_uid = {s["uid"]: s for s in self.enhancement_sets if s.get("uid")}
        self.power_by_full_name = {p["full_name"]: p for p in self.powers if p.get("full_name")}
        self.archetype_by_class_name = {a["class_name"]: a for a in self.archetypes if a.get("class_name")}
        self.powerset_by_full_name = {p["full_name"]: p for p in self.powersets if p.get("full_name")}

        # enhancement set lookup by enhancement uid (an enhancement points at its set via
        # uid_set on the Enhancement record itself)
        self.set_for_enhancement_uid = {}
        for e in self.enhancements:
            if e.get("uid_set"):
                self.set_for_enhancement_uid[e["uid"]] = e["uid_set"]

        # Power.GetValidEnhancements' generic (non-set) path matches an enhancement's
        # 0-based `class_id` indices against `enhancement_classes[idx]["id"]`, then checks
        # membership in the power's `enhancements_allowed_raw` list of those same "id" values.
        # This list IS enhancement_classes sorted by array position already (id is 1-based,
        # position is 0-based) - store it as-is and expose the by-position accessor.
        self._enhancement_class_id_by_index = [c["id"] for c in self.enhancement_classes]

        # powers grouped by their powerset's full_name ("<group_name>.<set_name>"), in the
        # order they appear in the source database (which is level/pick order for real
        # powersets) - powersets.json itself doesn't list its member powers.
        self.powers_by_powerset_full_name = {}
        for p in self.powers:
            if not p.get("group_name") or not p.get("set_name"):
                continue
            key = f"{p['group_name']}.{p['set_name']}"
            self.powers_by_powerset_full_name.setdefault(key, []).append(p)

    def enhancement_class_id(self, zero_based_index):
        """Mirrors `DatabaseAPI.Database.EnhancementClasses[classId].ID` - classId is a
        0-based index (as stored in an enhancement's `class_id` list); returns the 1-based
        "id" that a power's `enhancements_allowed_raw` list actually contains."""
        if 0 <= zero_based_index < len(self._enhancement_class_id_by_index):
            return self._enhancement_class_id_by_index[zero_based_index]
        return None

    def powers_in_powerset(self, powerset_full_name):
        return self.powers_by_powerset_full_name.get(powerset_full_name, [])

    def enhancement(self, uid):
        return self.enhancement_by_uid.get(uid)

    def enhancement_set(self, uid):
        return self.enhancement_set_by_uid.get(uid)

    def power(self, full_name):
        return self.power_by_full_name.get(full_name)

    def archetype(self, class_name):
        return self.archetype_by_class_name.get(class_name)


_default_db = None


def get_default_db():
    global _default_db
    if _default_db is None:
        here = os.path.dirname(os.path.abspath(__file__))
        candidates = [
            os.path.join(here, "data"),
            "/home/claude/coh-data-pipeline/out",
        ]
        for c in candidates:
            if os.path.isdir(c):
                _default_db = Database(c)
                return _default_db
        raise FileNotFoundError(f"Could not find a data directory among: {candidates}")
    return _default_db


if __name__ == "__main__":
    db = get_default_db()
    print(f"Loaded {len(db.powers)} powers, {len(db.enhancements)} enhancements, "
          f"{len(db.enhancement_sets)} enhancement sets, {len(db.archetypes)} archetypes")