"""
mlos_sort.py
Python port of MLOS_sorting.lua (Mod Load Order Sorter, by REfRigERatoR).

Faithful to the Lua algorithm:
  - preorder: ModManager, ModManagerServer, modoptions
  - category buckets: coreRequirement -> tweaks -> resource -> map -> vehicle ->
                      code -> clothes -> ui -> other -> translation -> undefined
  - loadFirst / loadLast: on (0) | category (1) | off (2)
  - topological sort by `require` + `loadAfter` with cycle detection
  - sorting_rules.txt overrides supported (loadAfter/loadBefore/incompatibleMods/
    loadFirst/loadLast/category)

Limitations vs in-game Lua:
  - mod.info-only input. We do NOT walk /media/* folders for category detection.
    We rely on mod.info `category=` if present, then `frameworkKeys` name
    heuristic, then default "other" (or "undefined" if uncategorizable).
  - `loadBefore` is converted into corresponding `loadAfter` edges on other mods,
    matching the Lua mod's behavior.
"""

import argparse
import json
import os
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# -----------------------------------------------------------------------------
# Constants (mirrors MLOS_sorting.lua)
# -----------------------------------------------------------------------------

PREORDER: Dict[str, int] = {"ModManager": 1, "ModManagerServer": 2, "modoptions": 3}

RAW_CATEGORY_ORDER: List[str] = [
    "coreRequirement",
    "tweaks",        # libraries / frameworks / APIs (matches existing 'lib' pill)
    "tile",          # tile asset packs
    "debug",         # error logger, cheat menus
    "resource",      # other generic resources
    "map",
    "qol",           # QOL changes
    "moodle",        # moodles / moodlets
    "tweak_minor",   # tiny tweaks (working aircon, ear protection, …)
    "music",         # music + music addons (load before vehicles per user spec)
    "wearable",      # clothing, hair, tattoos (NOT armor)
    "profession",    # profession mods
    "movement",      # drop-and-roll, crawl, jump, ladders
    "building",      # building menus, barricades, light switches
    "farming",
    "zombie",        # zombie behaviour mods (OccultZed, HordeNight, ReactiveZombies)
    "zone",          # hazardous zones, spore zones
    "armor",         # armor mods (separate from wearables)
    "food",
    "health",        # first aid, medical
    "weapon",        # weapons (load before vehicles per user spec)
    "crafting",
    "container",     # backpacks, boxes, tubs
    "vehicle",
    "vehicle_spawn", # vehicle spawn zones
    "loot",          # loot tables
    "code",          # generic gameplay code (legacy fallback)
    "ui",            # interface
    "sound",         # audio (non-music)
    "texture",
    "translation",
    "multiplayer",   # MP-specific utilities
    "server_only",   # admin tools, server logs
    "fix",           # bug-fix overlays (distinct from 'patch' loadLast tier)
    "other",
    "undefined",
    # Spec G-patch: "patch" is a category like any other, but `_initial_sort_key`
    # routes patches above all sub-axes via a leading is_patch tuple element so
    # they sort strictly LAST (after every loadLast=on map mod).
    "patch",
]
CATEGORY_ORDER: Dict[str, int] = {c: i for i, c in enumerate(RAW_CATEGORY_ORDER)}

LOAD_CATEGORIES: Dict[str, int] = {"on": 0, "category": 1, "off": 2}

# from MLOS_sorting.lua: frameworkKeys for name-based tweak detection.
# Lua uses string.find on lowercased name (substring match, no regex anchors).
FRAMEWORK_KEYS: List[str] = [
    "framework",
    " api",
    "_api",
    "tweak",
    "interface",
    "utilit",  # matches utility, utilities
    "bugfix",
    "librar",  # matches library, libraries — covers damnlib/tsarslib/StarlitLibrary/etc.
               # `derive_category` checks mod.maps before FRAMEWORK_KEYS, so a map
               # mod whose name contains "library" still classifies as `map` first.
]

# Multi-key list fields in mod.info (lowercased keys)
LIST_KEYS_MAP = {
    "require": "requirements",
    "loadafter": "loadAfter",
    "loadbefore": "loadBefore",
    "incompatiblemods": "incompatibleMods",
    "tags": "tags",
}


# -----------------------------------------------------------------------------
# Dataclasses
# -----------------------------------------------------------------------------


@dataclass
class ModInfo:
    id: str
    name: str = ""
    workshop_id: Optional[str] = None
    category: str = "undefined"
    requirements: List[str] = field(default_factory=list)
    loadAfter: List[str] = field(default_factory=list)
    loadBefore: List[str] = field(default_factory=list)
    incompatibleMods: List[str] = field(default_factory=list)
    loadFirst: str = "off"
    loadLast: str = "off"
    tags: List[str] = field(default_factory=list)
    maps: List[str] = field(default_factory=list)  # map folder names from media/maps/
    flags: List[str] = field(default_factory=list)
    is_addon: bool = False  # Spec A addon: default-off in multi-mod wsids
    # Steam Workshop's controlled-vocab tags (workshop_meta.tags). Canonical
    # signal for build / multiplayer / category detection. Distinct from
    # `tags` which is mod.info-side (freeform).
    workshop_tags: List[str] = field(default_factory=list)
    warnings: Dict[str, List[str]] = field(default_factory=dict)


@dataclass
class SortingRule:
    loadAfter: List[str] = field(default_factory=list)
    loadBefore: List[str] = field(default_factory=list)
    incompatibleMods: List[str] = field(default_factory=list)
    loadFirst: str = "off"
    loadLast: str = "off"
    category: Optional[str] = None


# -----------------------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------------------


def _split_csv(value: str) -> List[str]:
    """
    Mirrors Refr_Utils:splitStringBySeparator - split on commas, trim, drop empties.
    Defensively strips `word=` prefixes that some malformed mod.info lines include
    (the Lua does this too). Also strips a leading backslash on each entry: B42
    mod.info files write deps as `require=\\StarlitLibrary` (path-style); we want
    plain modIds so dep names match for the topo sort and missing-dep warnings.
    Finally strips a leading "<digits>/" wsid-path prefix some authors put in
    front of the modId (e.g. require=2256623447/firearmmod).
    """
    if value is None:
        return []
    cleaned = re.sub(r"\w+\s*=", "", value)
    out: List[str] = []
    for p in cleaned.split(","):
        s = p.strip().lstrip("\\")
        s = re.sub(r"^\d+/", "", s)
        if s:
            out.append(s)
    return out


def _convert_load_category(value) -> str:
    """Mirrors convertToLoadCategoryString: normalize to 'on' | 'category' | 'off'."""
    if value in (True, "true", 0, "0"):
        return "on"
    if value in (None, False, "false", 2, "2", ""):
        return "off"
    if value not in LOAD_CATEGORIES:
        return "off"
    return str(value)


def _str_contains_any(haystack: str, needles: List[str]) -> bool:
    if not haystack:
        return False
    h = haystack.lower()
    return any(n and n in h for n in needles)


# -----------------------------------------------------------------------------
# Parsers
# -----------------------------------------------------------------------------


def parse_mod_info(text: str, workshop_id: Optional[str] = None) -> Optional[ModInfo]:
    """
    Parse a mod.info file body. Returns None if no `id=` line found.
    Lines are `key=value`; keys lowercased; list-fields comma-separated.
    """
    fields: Dict[str, object] = {}
    for raw in text.splitlines():
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        m = re.match(r"^\s*(.+?)\s*=\s*(.*?)\s*$", line)
        if not m:
            continue
        key = m.group(1).strip().lower()
        value = m.group(2).strip()
        if key in LIST_KEYS_MAP:
            fields[LIST_KEYS_MAP[key]] = _split_csv(value)
        elif key == "loadfirst":
            fields["loadFirst"] = _convert_load_category(value)
        elif key == "loadlast":
            fields["loadLast"] = _convert_load_category(value)
        elif key == "category":
            fields["category"] = value if value in CATEGORY_ORDER else "undefined"
        elif key == "name":
            fields["name"] = value
        elif key == "id":
            # Some authors prefix the wsid into the id (e.g. id=2256623447/firearmmod).
            # Strip a leading "<digits>/" so the canonical mod_id is the clean form.
            fields["id"] = re.sub(r"^\d+/", "", value)

    if "id" not in fields:
        return None

    return ModInfo(
        id=fields["id"],
        name=fields.get("name", ""),
        workshop_id=workshop_id,
        category=fields.get("category", "undefined"),
        requirements=fields.get("requirements", []),
        loadAfter=fields.get("loadAfter", []),
        loadBefore=fields.get("loadBefore", []),
        incompatibleMods=fields.get("incompatibleMods", []),
        loadFirst=fields.get("loadFirst", "off"),
        loadLast=fields.get("loadLast", "off"),
        tags=fields.get("tags", []),
    )


def parse_sorting_rules(text: str) -> Dict[str, SortingRule]:
    """
    Parse a sorting_rules.txt file. Format:
        [modId]
        loadAfter=mod1,mod2
        loadBefore=mod3
        incompatibleMods=mod4
        loadFirst=on
        loadLast=off
        category=tweaks
    """
    rules: Dict[str, SortingRule] = {}
    current: Optional[str] = None
    for raw in text.splitlines():
        line = raw.strip()
        if not line:
            continue
        m = re.match(r"^\s*\[\s*(.+?)\s*\]\s*$", line)
        if m:
            current = m.group(1)
            rules.setdefault(current, SortingRule())
            continue
        if current is None:
            continue
        kv = re.match(r"^\s*(.+?)\s*=\s*(.*?)\s*$", line)
        if not kv:
            continue
        key, value = kv.group(1).lower(), kv.group(2)
        rule = rules[current]
        if key in ("loadafter", "loadmodafter"):
            rule.loadAfter = list(dict.fromkeys(rule.loadAfter + _split_csv(value)))
        elif key in ("loadbefore", "loadmodbefore"):
            rule.loadBefore = list(dict.fromkeys(rule.loadBefore + _split_csv(value)))
        elif key in ("incompatiblemods", "incompatible"):
            rule.incompatibleMods = list(dict.fromkeys(rule.incompatibleMods + _split_csv(value)))
        elif key == "loadfirst":
            rule.loadFirst = _convert_load_category(value)
        elif key == "loadlast":
            rule.loadLast = _convert_load_category(value)
        elif key == "category":
            rule.category = value if value in CATEGORY_ORDER else None
    return rules


# -----------------------------------------------------------------------------
# Filesystem ingestion (DepotDownloader output layout)
# -----------------------------------------------------------------------------


def load_mods_from_dir(root: Path) -> List[ModInfo]:
    """
    Walk `<root>/<workshop_id>/mods/<mod_id>/mod.info` (DepotDownloader output).
    Also: `<root>/<workshop_id>/mods/<mod_id>/media/maps/<map_folder>/map.info`
    populates `maps` for that mod.
    """
    mods: List[ModInfo] = []
    if not root.exists():
        raise FileNotFoundError(f"Mods root does not exist: {root}")

    for workshop_dir in sorted(root.iterdir()):
        if not workshop_dir.is_dir():
            continue
        workshop_id = workshop_dir.name if workshop_dir.name.isdigit() else None
        mods_root = workshop_dir / "mods"
        if not mods_root.exists():
            # also support: some layouts put mods/ at root level directly
            mods_root = workshop_dir
        if not mods_root.exists() or not mods_root.is_dir():
            continue

        for mod_dir in sorted(mods_root.iterdir()):
            if not mod_dir.is_dir():
                continue
            mod_info_path = mod_dir / "mod.info"
            if not mod_info_path.exists():
                continue
            try:
                text = mod_info_path.read_text(encoding="utf-8", errors="replace")
            except OSError as e:
                print(f"WARN: cannot read {mod_info_path}: {e}", file=sys.stderr)
                continue
            mod = parse_mod_info(text, workshop_id=workshop_id)
            if mod is None:
                print(f"WARN: no `id=` in {mod_info_path}, skipping", file=sys.stderr)
                continue

            # Collect map folders
            maps_dir = mod_dir / "media" / "maps"
            if maps_dir.exists() and maps_dir.is_dir():
                for map_folder in sorted(maps_dir.iterdir()):
                    if map_folder.is_dir() and (map_folder / "map.info").exists():
                        mod.maps.append(map_folder.name)

            mods.append(mod)

    return mods


# -----------------------------------------------------------------------------
# Category derivation (degraded vs in-game; no folder walk)
# -----------------------------------------------------------------------------


_PATCH_NAME_RE = re.compile(r"\b(patch|compat|compatibility)\b", re.IGNORECASE)


# Substring-based category hints (kept in sync with api/mlos_sort.py)
_LIB_NAME_HINTS  = ["library", "libraries", "framework"]
_LIB_NAME_RE = re.compile(
    r'(?<![A-Za-z])(?:lib|api|core)(?![A-Za-z])'
    r'|(?<=[a-z])(?:Lib|API|Core)(?![A-Za-z])',
    re.IGNORECASE,
)
_MUSIC_NAME_HINTS    = ["music", "moozic", "jukebox"]
_MOODLE_NAME_HINTS   = ["moodle", "moodlet"]
_PROFESSION_HINTS    = ["profession"]
_MOVEMENT_HINTS      = [
    "true action", "trueaction", "true_action",
    "drop and roll", "dropandroll", "drop_and_roll",
    "crawl", "ladder",
]
_ARMOR_NAME_HINTS    = ["armor", "armour"]
_HEALTH_NAME_HINTS   = ["first aid", "firstaid", "medical", "injur", "disease", "sickness"]
_CRAFTING_HINTS      = ["craft"]
_CONTAINER_HINTS     = ["backpack", "container", "storage"]
_LOOT_NAME_HINTS     = ["loot"]
_TILE_NAME_HINTS     = ["tiles", "tileset", "tilepack"]
_DEBUG_NAME_HINTS    = ["debug menu", "cheat menu", "error log", "errormagnifier"]
_ZONE_NAME_HINTS     = ["hazard zone", "spore zone", "spore zones"]
_ZOMBIE_NAME_HINTS   = ["zombie", "horde", "undead"]
_FIX_NAME_HINTS      = [" fix", "_fix", "bugfix", "hotfix"]


def _name_has(name: str, hints: List[str]) -> bool:
    if not name:
        return False
    n = name.lower()
    return any(h in n for h in hints)


def derive_category(mod: ModInfo) -> str:
    """Best-effort category from mod.info + workshop_meta.tags + name.
    Mirrors api/mlos_sort.py; keep both copies in sync.
    """
    if mod.category in CATEGORY_ORDER and mod.category != "undefined":
        return mod.category

    name = mod.name or ""
    if name and _PATCH_NAME_RE.search(name):
        return "patch"
    if _name_has(name, _LIB_NAME_HINTS) or (name and _LIB_NAME_RE.search(name)):
        return "tweaks"

    if mod.maps:
        return "map"

    if _name_has(name, _MUSIC_NAME_HINTS):
        return "music"
    if _name_has(name, _MOVEMENT_HINTS):
        return "movement"
    if _name_has(name, _MOODLE_NAME_HINTS):
        return "moodle"
    if _name_has(name, _DEBUG_NAME_HINTS):
        return "debug"
    if _name_has(name, _TILE_NAME_HINTS):
        return "tile"

    ws_tags = set(mod.workshop_tags or [])
    has_audio = "Audio" in ws_tags

    if "Weapons" in ws_tags:
        return "weapon"
    if "Vehicles" in ws_tags:
        if name and "spawn zone" in name.lower():
            return "vehicle_spawn"
        return "vehicle"
    if "Clothing/Armor" in ws_tags:
        if _name_has(name, _ARMOR_NAME_HINTS):
            return "armor"
        return "wearable"
    if "Food" in ws_tags:
        return "food"
    if "building" in {t.lower() for t in ws_tags}:
        return "building"
    if "Farming" in ws_tags:
        return "farming"
    if "Skills" in ws_tags:
        if _name_has(name, _PROFESSION_HINTS):
            return "profession"
        return "code"
    if "Interface" in ws_tags:
        return "ui"
    if "Textures" in ws_tags:
        return "texture"
    if "Language/Translation" in ws_tags:
        return "translation"
    if "QOL" in ws_tags:
        return "qol"
    if "Map" in ws_tags:
        return "map"

    if _name_has(name, _TILE_NAME_HINTS):
        return "tile"
    if _name_has(name, _ZOMBIE_NAME_HINTS):
        return "zombie"
    if _name_has(name, _HEALTH_NAME_HINTS):
        return "health"
    if _name_has(name, _CRAFTING_HINTS):
        return "crafting"
    if _name_has(name, _CONTAINER_HINTS):
        return "container"
    if _name_has(name, _LOOT_NAME_HINTS):
        return "loot"
    if _name_has(name, _FIX_NAME_HINTS):
        return "fix"
    if _name_has(name, _ZONE_NAME_HINTS):
        return "zone"

    if has_audio:
        return "sound"

    tags_lc = [t.lower() for t in mod.tags]
    if any("translation" in t for t in tags_lc):
        return "translation"
    if any("vehicle" in t for t in tags_lc):
        return "vehicle"
    if any("interface" in t or "ui" in t for t in tags_lc):
        return "ui"
    if any("clothing" in t or "skin" in t for t in tags_lc):
        return "wearable"
    if any("armor" in t for t in tags_lc):
        return "armor"
    if any("map" in t for t in tags_lc):
        return "map"

    if _str_contains_any(name, FRAMEWORK_KEYS):
        return "tweaks"

    if "Multiplayer" in ws_tags:
        return "multiplayer"

    return "other"


# -----------------------------------------------------------------------------
# Sort
# -----------------------------------------------------------------------------


def _apply_overrides(mods: List[ModInfo], rules: Dict[str, SortingRule]) -> None:
    """In-place: merge rules into mods, then propagate loadBefore -> reverse loadAfter."""
    by_id = {m.id: m for m in mods}

    for mod in mods:
        rule = rules.get(mod.id)
        if rule:
            mod.loadAfter = list(dict.fromkeys(mod.loadAfter + rule.loadAfter))
            mod.loadBefore = list(dict.fromkeys(mod.loadBefore + rule.loadBefore))
            mod.incompatibleMods = list(dict.fromkeys(mod.incompatibleMods + rule.incompatibleMods))
            mod.loadFirst = _convert_load_category(rule.loadFirst if rule.loadFirst != "off" else mod.loadFirst)
            mod.loadLast = _convert_load_category(rule.loadLast if rule.loadLast != "off" else mod.loadLast)
            if rule.category:
                mod.category = rule.category

        # Derive category if still undefined
        if mod.category not in CATEGORY_ORDER or mod.category == "undefined":
            mod.category = derive_category(mod)

    # Translate loadBefore into reverse loadAfter on the target mod
    # (mirrors updateSortingRulesLoadAfter)
    for mod in mods:
        for target in mod.loadBefore:
            target_mod = by_id.get(target)
            if target_mod and mod.id not in target_mod.loadAfter:
                target_mod.loadAfter.append(mod.id)


def _initial_sort_key(mod: ModInfo):
    """Mirrors initialSortMods comparator. Returns sortable tuple.

    Spec G-patch: index 0 is `is_patch` so patches sort strictly last - they
    have to override loadLast=on map mods at runtime. Within the patch tier
    the existing sub-axes still apply (PREORDER, alpha, etc.).
    """
    is_patch = 1 if mod.category == "patch" else 0
    pre = PREORDER.get(mod.id, 10000)
    return (
        is_patch,
        pre,
        LOAD_CATEGORIES[mod.loadFirst],          # global loadFirst (on first)
        -LOAD_CATEGORIES[mod.loadLast] + 100,    # global loadLast (on last) -- keep parity by
                                                  # sorting "on" last; we invert so smaller=earlier
        CATEGORY_ORDER.get(mod.category, CATEGORY_ORDER["undefined"]),
        LOAD_CATEGORIES[mod.loadFirst],          # in-category loadFirst
        -LOAD_CATEGORIES[mod.loadLast] + 100,    # in-category loadLast
        mod.id.lower(),
    )


def _topological_sort(mods: List[ModInfo]) -> Tuple[List[str], List[List[str]]]:
    """DFS topo sort on (requirements + loadAfter). Returns (order, cycles)."""
    by_id = {m.id: m for m in mods}
    visited: Dict[str, bool] = {}
    visiting: Dict[str, bool] = {}
    order: List[str] = []
    cycles: List[List[str]] = []

    def visit(mod: ModInfo, path: List[str]):
        if visiting.get(mod.id):
            cycles.append(path + [mod.id])
            return
        if visited.get(mod.id):
            return
        visiting[mod.id] = True
        for dep in mod.requirements:
            target = by_id.get(dep)
            if target:
                visit(target, path + [mod.id])
        for dep in mod.loadAfter:
            target = by_id.get(dep)
            if target:
                visit(target, path + [mod.id])
        visiting[mod.id] = False
        visited[mod.id] = True
        order.append(mod.id)

    for mod in mods:
        visit(mod, [])
    return order, cycles


def sort_mods(
    mods: List[ModInfo],
    rules: Optional[Dict[str, SortingRule]] = None,
) -> Dict[str, object]:
    """
    Top-level entry: returns dict with ordered IDs + warnings.
    """
    rules = rules or {}
    _apply_overrides(mods, rules)

    # Initial deterministic sort (preorder, loadFirst, category, loadLast, alpha)
    mods.sort(key=_initial_sort_key)

    order, cycles = _topological_sort(mods)

    by_id = {m.id: m for m in mods}
    enabled = set(by_id.keys())

    missing: Dict[str, List[str]] = {}
    incompat: Dict[str, List[str]] = {}

    for mod in mods:
        miss = [r for r in mod.requirements if r not in enabled]
        if miss:
            missing[mod.id] = miss
        inc = [r for r in mod.incompatibleMods if r in enabled]
        if inc:
            incompat[mod.id] = inc

    # Output blocks for the server's .ini file
    mods_line = order  # already mod IDs in load order
    workshop_seen: List[str] = []
    workshop_set = set()
    for mod_id in order:
        wid = by_id[mod_id].workshop_id
        if wid and wid not in workshop_set:
            workshop_seen.append(wid)
            workshop_set.add(wid)
    # MAP_LINE convention: dependencies first (leftmost), dependents next.
    # Vanilla Muldraugh, KY is ALWAYS appended at the very end by
    # adapters.build_response. `order` is already topo-sorted by mod-level
    # deps (require= / loadAfter= / loadBefore=), so dependencies appear
    # before their dependents — walk it forward.
    map_folders: List[str] = []
    for mod_id in order:
        for mf in by_id[mod_id].maps:
            if mf not in map_folders:
                map_folders.append(mf)

    return {
        "Mods": mods_line,
        "WorkshopItems": workshop_seen,
        "Map": map_folders,
        "warnings": {
            "cycles": cycles,
            "missing_requirements": missing,
            "incompatible_enabled": incompat,
        },
    }


# -----------------------------------------------------------------------------
# CLI
# -----------------------------------------------------------------------------


def main():
    ap = argparse.ArgumentParser(
        description="Sort PZ mods by load order. Reads DepotDownloader output layout.",
    )
    ap.add_argument("mods_root", help="Path containing <workshop_id>/mods/<mod_id>/mod.info trees")
    ap.add_argument("--rules", help="Optional sorting_rules.txt path")
    ap.add_argument("--json", action="store_true", help="Output JSON instead of ini blocks")
    args = ap.parse_args()

    root = Path(args.mods_root).resolve()
    mods = load_mods_from_dir(root)
    if not mods:
        print("ERROR: no mods found", file=sys.stderr)
        sys.exit(2)

    rules: Dict[str, SortingRule] = {}
    if args.rules:
        rules = parse_sorting_rules(Path(args.rules).read_text(encoding="utf-8"))

    result = sort_mods(mods, rules)

    if args.json:
        print(json.dumps(result, indent=2))
        return

    print("WorkshopItems=" + ";".join(result["WorkshopItems"]))
    print("Mods=" + ";".join(result["Mods"]))
    if result["Map"]:
        print("Map=" + ";".join(result["Map"]))
    w = result["warnings"]
    if w["cycles"] or w["missing_requirements"] or w["incompatible_enabled"]:
        print("\n# Warnings")
        if w["cycles"]:
            print("# cycles:", w["cycles"])
        if w["missing_requirements"]:
            print("# missing_requirements:", w["missing_requirements"])
        if w["incompatible_enabled"]:
            print("# incompatible_enabled:", w["incompatible_enabled"])


if __name__ == "__main__":
    main()