#!/usr/bin/env python3
"""
Build map/mode info JSON for tank missions with planes + capture zones.

Output shape:
{
  "<map_name>": {
    "<mode>": [
      {
        "info": { ... }
      }
    ]
  }
}
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any


SCRIPT_DIR = Path(__file__).resolve().parent
BOTS_DIR = SCRIPT_DIR.parent.parent

DATAMINE_LEVELS_DIR = BOTS_DIR / "War-Thunder-Datamine" / "aces.vromfs.bin_u" / "levels"
DATAMINE_MISSIONS_DIR = BOTS_DIR / "War-Thunder-Datamine" / "mis.vromfs.bin_u"
MISSION_ROOT = DATAMINE_MISSIONS_DIR / "gamedata" / "missions" / "cta" / "tanks"

SHARED_LEVELS_DIR = BOTS_DIR / "SHARED" / "MAPS" / "LEVELS"
SHARED_MISSIONS_DIR = SHARED_LEVELS_DIR / "MISSIONS"

SKIP_PARTS = {"mainareas", "battleareas", "spawnareas"}


def _clean_map_key(raw: str) -> str:
    key = Path(raw).name.strip().rstrip("*")
    if key.endswith(".png"):
        key = key[:-4]
    return key


def _normalize_rel(path: str) -> str:
    rel = str(path or "").strip().replace("\\", "/").lstrip("/")
    if rel[:9].lower() == "gamedata/":
        rel = "gamedata/" + rel[9:]
    return rel.lower()


def _load_json(path: Path, cache: dict[Path, Any]) -> Any:
    if path in cache:
        return cache[path]
    try:
        cache[path] = json.loads(path.read_text())
    except Exception:
        cache[path] = None
    return cache[path]


def _repo_rel(path: Path) -> str:
    try:
        return str(path.resolve().relative_to(BOTS_DIR.resolve())).replace("\\", "/")
    except Exception:
        return str(path)


def _datamine_level_path(level_path: str) -> str:
    rel = str(level_path or "").replace("\\", "/").lstrip("/")
    if rel.lower().startswith("levels/"):
        return f"War-Thunder-Datamine/aces.vromfs.bin_u/{rel}"
    return f"War-Thunder-Datamine/{rel}" if rel else rel


def _resolve_mission_file(rel: str) -> Path | None:
    norm = _normalize_rel(rel)
    if not norm:
        return None
    for root in (SHARED_MISSIONS_DIR, DATAMINE_MISSIONS_DIR):
        raw = root / norm
        if raw.suffix.lower() == ".blk":
            cands = [raw.with_suffix(".blkx"), raw]
        elif raw.suffix.lower() == ".blkx":
            cands = [raw, raw.with_suffix(".blk")]
        else:
            cands = [raw, raw.with_suffix(".blkx"), raw.with_suffix(".blk")]
        for p in cands:
            if p.exists():
                return p
    return None


def _import_paths(mission_def: dict) -> list[str]:
    imports = mission_def.get("imports")
    if not isinstance(imports, dict):
        return []
    import_record = imports.get("import_record")
    records = import_record if isinstance(import_record, list) else [import_record]
    out: list[str] = []
    for rec in records:
        if not isinstance(rec, dict):
            continue
        f = rec.get("file")
        if isinstance(f, str) and f.strip():
            out.append(f)
    return out


def _collect_defs(root_file: Path, cache: dict[Path, Any]) -> tuple[dict[str, dict], list[str], list[str]]:
    root_def = _load_json(root_file, cache)
    if not isinstance(root_def, dict):
        return {}, [], []

    areas: dict[str, dict] = {}
    files_used: list[str] = []
    imports_used: list[str] = []
    queue: list[Path] = [root_file]
    seen: set[Path] = set()

    while queue:
        cur_file = queue.pop(0).resolve()
        if cur_file in seen:
            continue
        seen.add(cur_file)
        files_used.append(str(cur_file))

        cur_def = _load_json(cur_file, cache)
        if not isinstance(cur_def, dict):
            continue

        cur_areas = cur_def.get("areas")
        if isinstance(cur_areas, dict):
            for name, area in cur_areas.items():
                if isinstance(name, str) and isinstance(area, dict) and name not in areas:
                    areas[name] = area

        for rel in _import_paths(cur_def):
            imports_used.append(_normalize_rel(rel))
            p = _resolve_mission_file(rel)
            if p is not None:
                queue.append(p)

    return areas, files_used, sorted(set(imports_used))


def _radius_from_tm(tm: list) -> float:
    try:
        ax = float(tm[0][0])
        az = float(tm[0][2])
        bx = float(tm[2][0])
        bz = float(tm[2][2])
        ra = (ax * ax + az * az) ** 0.5
        rb = (bx * bx + bz * bz) ** 0.5
        return max(ra, rb)
    except Exception:
        return 0.0


def _capture_prefixes(mission_type: str, mode_key: str) -> list[str]:
    m = (mission_type or "").lower()
    k = (mode_key or "").lower()
    if "conq" in m or "conq" in k:
        return ["conq_capture_area_"]
    if "dom" in m or "dom" in k:
        return ["dom_capture_area_"]
    if "bttl" in m or "bttl" in k:
        return ["bttl_t1_capture_area_", "bttl_t2_capture_area_"]
    return ["capture_area_"]


def _capture_sort(name: str) -> tuple[int, str]:
    digits = "".join(ch for ch in name if ch.isdigit())
    return (int(digits) if digits else 999, name)


def _pick_capture_areas(areas: dict[str, dict], mission_type: str, mode_key: str) -> list[dict]:
    prefixes = _capture_prefixes(mission_type, mode_key)
    names: list[str] = []
    for name, area in areas.items():
        if not isinstance(area, dict):
            continue
        lname = name.lower()
        if "capture_area" not in lname or lname.startswith("briefing_"):
            continue
        if not any(lname.startswith(p) for p in prefixes):
            continue
        names.append(name)

    if not names:
        return []

    suffix_prio = ["_arcade", "_realistic", "_hardcore", "_simulator", ""]
    groups: dict[str, dict[str, str]] = {}
    for name in names:
        lname = name.lower()
        suffix = ""
        base = lname
        for sfx in suffix_prio[:-1]:
            if lname.endswith(sfx):
                suffix = sfx
                base = lname[:-len(sfx)]
                break
        groups.setdefault(base, {})[suffix] = name

    chosen: list[str] = []
    for base in sorted(groups.keys(), key=_capture_sort):
        variants = groups[base]
        for sfx in suffix_prio:
            if sfx in variants:
                chosen.append(variants[sfx])
                break

    out: list[dict] = []
    for name in chosen:
        area = areas.get(name)
        if not isinstance(area, dict):
            continue
        tm = area.get("tm")
        if not isinstance(tm, list) or len(tm) < 4 or not isinstance(tm[3], list):
            continue
        try:
            cx = float(tm[3][0])
            cz = float(tm[3][2])
        except Exception:
            continue
        out.append({
            "name": name,
            "type": str(area.get("type", "")),
            "x": cx,
            "z": cz,
            "radius": _radius_from_tm(tm),
        })
    return out


def _battle_areas(areas: dict[str, dict]) -> list[dict]:
    out: list[dict] = []
    for name, area in areas.items():
        if not isinstance(area, dict):
            continue
        lname = name.lower()
        if "battle_area" not in lname or lname.startswith("briefing_"):
            continue
        tm = area.get("tm")
        if not isinstance(tm, list) or len(tm) < 4 or not isinstance(tm[3], list):
            continue
        try:
            cx = float(tm[3][0])
            cz = float(tm[3][2])
        except Exception:
            continue
        out.append({
            "name": name,
            "type": str(area.get("type", "")),
            "x": cx,
            "z": cz,
            "radius": _radius_from_tm(tm),
        })
    out.sort(key=lambda x: x["name"])
    return out


def _level_info(level_stem: str, cache: dict[Path, Any]) -> dict:
    cands = [
        DATAMINE_LEVELS_DIR / f"{level_stem}.blkx",
        SHARED_LEVELS_DIR / f"{level_stem}.blkx",
        SHARED_LEVELS_DIR / "DATAMINE" / f"{level_stem}.blkx",
    ]
    for p in cands:
        if not p.exists():
            continue
        data = _load_json(p, cache)
        if not isinstance(data, dict):
            continue
        return {
            "file": _repo_rel(p),
            "tankMapCoord0": data.get("tankMapCoord0"),
            "tankMapCoord1": data.get("tankMapCoord1"),
            "aiTanksMapCoord0": data.get("aiTanksMapCoord0"),
            "aiTanksMapCoord1": data.get("aiTanksMapCoord1"),
            "customLevelTankMap": data.get("customLevelTankMap"),
            "customLevelMap": data.get("customLevelMap"),
        }
    return {}


def _has_planes(allowed: dict | None, imports_used: list[str]) -> bool:
    if isinstance(allowed, dict):
        val = allowed.get("isAirplanesAllowed")
        if isinstance(val, bool):
            return val
    return any("air_spawns" in p for p in imports_used)


def _has_tanks(allowed: dict | None) -> bool:
    if isinstance(allowed, dict):
        val = allowed.get("isTanksAllowed")
        if isinstance(val, bool):
            return val
    return True


def _mode_key(mission_cfg: dict, mission_file: Path) -> str:
    postfix = mission_cfg.get("postfix")
    if isinstance(postfix, str) and postfix.strip():
        return postfix.lstrip("_")
    mtype = mission_cfg.get("type")
    if isinstance(mtype, str) and mtype.strip():
        return mtype
    return mission_file.stem


def should_skip_mission_file(path: Path) -> bool:
    if path.name.startswith("template_"):
        return True
    parts = {p.lower() for p in path.parts}
    if SKIP_PARTS & parts:
        return True
    return False


def build_index() -> dict[str, dict[str, list[dict[str, Any]]]]:
    cache: dict[Path, Any] = {}
    output: dict[str, dict[str, list[dict[str, Any]]]] = {}

    for mission_file in sorted(MISSION_ROOT.rglob("*.blkx")):
        if should_skip_mission_file(mission_file):
            continue

        mission_def = _load_json(mission_file, cache)
        if not isinstance(mission_def, dict):
            continue
        mission_settings = mission_def.get("mission_settings")
        if not isinstance(mission_settings, dict):
            continue
        mission_cfg = mission_settings.get("mission")
        if not isinstance(mission_cfg, dict):
            continue

        level_path = mission_cfg.get("level", "")
        if not isinstance(level_path, str) or not level_path.strip():
            continue
        level_stem = Path(level_path).stem
        mode = _mode_key(mission_cfg, mission_file)
        mission_type = str(mission_cfg.get("type", ""))
        allowed = mission_cfg.get("allowedUnitTypes")

        areas, files_used, imports_used = _collect_defs(mission_file, cache)
        captures = _pick_capture_areas(areas, mission_type, mode)
        if not captures:
            continue

        has_tanks = _has_tanks(allowed if isinstance(allowed, dict) else None)
        has_planes = _has_planes(allowed if isinstance(allowed, dict) else None, imports_used)
        if not (has_tanks and has_planes):
            continue

        battle_areas = _battle_areas(areas)
        info = {
            "mission_file": _repo_rel(mission_file),
            "level_path": _datamine_level_path(level_path),
            "mission_type": mission_type,
            "mode_postfix": mission_cfg.get("postfix", ""),
            "allowed_unit_types": allowed if isinstance(allowed, dict) else {},
            "useAlternativeMapCoord": mission_cfg.get("useAlternativeMapCoord"),
            "capture_areas": captures,
            "battle_areas": battle_areas,
            "level_def": _level_info(level_stem, cache),
            "imports_count": len(imports_used),
            "files_used_count": len(files_used),
        }

        output.setdefault(level_stem, {}).setdefault(mode, []).append({"info": info})

    return output


def main() -> None:
    parser = argparse.ArgumentParser(description="Build map_name -> mode -> [info] JSON")
    parser.add_argument(
        "--out",
        type=Path,
        default=SCRIPT_DIR.parent / "map_mode_info.json",
        help="Output JSON path",
    )
    args = parser.parse_args()

    data = build_index()
    args.out.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")

    maps = len(data)
    modes = sum(len(v) for v in data.values())
    entries = sum(len(items) for v in data.values() for items in v.values())
    print(f"Wrote {args.out} (maps={maps}, modes={modes}, entries={entries})")


if __name__ == "__main__":
    main()