
from __future__ import annotations

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


def read_text_best_effort(path: Path) -> str:
    data = path.read_bytes()
    for enc in ("utf-8-sig", "utf-8", "cp1252", "latin-1"):
        try:
            return data.decode(enc)
        except UnicodeDecodeError:
            continue
    return data.decode("latin-1", errors="replace")


def normpath_omsi(p: str) -> str:
    p = p.strip().strip('"').strip()
    return p.replace("/", "\\")


def is_texture_folder(p: Path) -> bool:
    return p.is_dir() and p.name.lower() == "texture"


def nonempty_lines(lines: List[str]) -> List[str]:
    return [ln for ln in (ln.strip() for ln in lines) if ln]


TAG_RE = re.compile(r"^\[([^\]]+)\]\s*$")


def parse_sections(text: str) -> Dict[str, List[str]]:
    raw = text.splitlines()
    cleaned: List[str] = []
    for ln in raw:
        s = ln.rstrip("\n\r")
        s_strip = s.strip()
        if not s_strip:
            cleaned.append("") 
            continue
        if s_strip.startswith("//"):
            continue
        cleaned.append(s)

    sections: Dict[str, List[str]] = {}
    current_tag: Optional[str] = None
    buf: List[str] = []

    def flush():
        nonlocal current_tag, buf
        if current_tag is not None:
            sections[current_tag] = buf[:]
        buf = []

    for ln in cleaned:
        m = TAG_RE.match(ln.strip())
        if m:
            flush()
            current_tag = m.group(1).strip().lower()
            continue
        if current_tag is not None:
            buf.append(ln)

    flush()
    return sections


@dataclass
class BusVariant:
    bus_file: str
    rel_path: str

    friendly_name: Optional[str] = None
    friendly_name_lines: List[str] = None
    description_lines: List[str] = None

    model_cfg: Optional[str] = None
    scripts: List[str] = None
    sounds: List[str] = None
    coupled_bus: Optional[str] = None

    repaint_mode: str = "unknown"      
    repaint_ctc_key: Optional[str] = None

    model_slots: List[str] = None      
    repaints: Dict[str, Dict[str, str]] = None  
    repaint_sources: Dict[str, List[str]] = None 


def parse_bus(bus_path: Path) -> BusVariant:
    text = read_text_best_effort(bus_path)
    sections = parse_sections(text)

    friendly_lines = nonempty_lines(sections.get("friendlyname", []))
    desc_lines = nonempty_lines(sections.get("description", []))

    friendly_name = friendly_lines[0] if friendly_lines else None

    model_cfg = None
    if "model" in sections:
        model_lines = nonempty_lines(sections["model"])
        if model_lines:
            model_cfg = normpath_omsi(model_lines[0])

    scripts: List[str] = []
    if "script" in sections:
        for ln in nonempty_lines(sections["script"]):
            scripts.append(normpath_omsi(ln))

    sounds: List[str] = []
    if "sound" in sections:
        sound_lines = nonempty_lines(sections["sound"])
        if sound_lines:
            sounds.append(normpath_omsi(sound_lines[0]))
    if "sound_ai" in sections:
        sound_ai_lines = nonempty_lines(sections["sound_ai"])
        if sound_ai_lines:
            sounds.append(normpath_omsi(sound_ai_lines[0]))

    coupled_bus = None
    for t in ("coupling_back", "couple_back", "trailer", "trail", "coupled"):
        if t in sections:
            lines = nonempty_lines(sections[t])
            for ln in lines:
                if ".bus" in ln.lower():
                    coupled_bus = normpath_omsi(ln)
                    break
        if coupled_bus:
            break

    repaint_mode = "unknown"
    repaint_ctc_key = None
    if "ctc" in sections:
        ctc_lines = [ln.strip() for ln in sections["ctc"] if ln.strip()]
        for i in range(len(ctc_lines) - 1):
            if ctc_lines[i].strip().lower() == "repaint":
                repaint_mode = "ctc"
                repaint_ctc_key = ctc_lines[i + 1].strip()
                break

    return BusVariant(
        bus_file=bus_path.name,
        rel_path=str(bus_path.relative_to(bus_path.parents[1] if len(bus_path.parents) > 1 else bus_path.parent)),
        friendly_name=friendly_name,
        friendly_name_lines=friendly_lines,
        description_lines=desc_lines,
        model_cfg=model_cfg,
        scripts=scripts,
        sounds=sounds,
        coupled_bus=coupled_bus,
        repaint_mode=repaint_mode,
        repaint_ctc_key=repaint_ctc_key,
        model_slots=[],
        repaints={},
        repaint_sources={}
    )



FARBSHEMA_RE = re.compile(r"\b(farbschema_[A-Za-z0-9_]+)\b", re.IGNORECASE)

def parse_model_slots(model_path: Path) -> Set[str]:
    if not model_path.exists():
        return set()
    text = read_text_best_effort(model_path)
    slots = set(m.group(1).lower() for m in FARBSHEMA_RE.finditer(text))
    return slots



@dataclass
class CtiRepaintIndex:
    item_repaints: Dict[str, Dict[str, str]]
    ctc_repaints: Dict[str, Dict[str, Dict[str, str]]]

def looks_like_texture_path(s: str) -> bool:
    s = s.strip().strip('"')
    sl = s.lower()
    return any(sl.endswith(ext) for ext in (".dds", ".png", ".bmp", ".tga", ".jpg", ".jpeg"))

def parse_cti(cti_path: Path) -> CtiRepaintIndex:
    text = read_text_best_effort(cti_path)
    raw_lines = [ln.strip() for ln in text.splitlines()]
    lines = [ln for ln in raw_lines if ln and not ln.startswith("//")]

    item_repaints: Dict[str, Dict[str, str]] = {}
    ctc_repaints: Dict[str, Dict[str, Dict[str, str]]] = {}

    has_item = any(ln.lower() == "[item]" for ln in lines)
    has_ctc = any(ln.lower() == "[ctc]" for ln in lines) or any(ln.lower() == "[repaint]" for ln in lines)

    if has_item:
        i = 0
        while i < len(lines):
            if lines[i].lower() == "[item]":
                repaint_name = lines[i + 1].strip() if i + 1 < len(lines) else "UNKNOWN_REPAINT"
                i += 2
                slotmap: Dict[str, str] = {}

                while i < len(lines):
                    if lines[i].startswith("[") and lines[i].endswith("]"):
                        break

                    slot = lines[i].strip()
                    tex = lines[i + 1].strip() if i + 1 < len(lines) else ""
                    if not tex or (tex.startswith("[") and tex.endswith("]")):
                        i += 1
                        continue

                    slotmap[slot.lower()] = normpath_omsi(tex)
                    i += 2

                if slotmap:
                    item_repaints.setdefault(repaint_name, {}).update(slotmap)

                continue
            i += 1

    if has_ctc:
        i = 0
        current_ctc_key: Optional[str] = None
        while i < len(lines):
            ln = lines[i].lower()
            if ln == "[ctc]":
                candidate = lines[i + 1].strip() if i + 1 < len(lines) else None
                if candidate and not (candidate.startswith("[") and candidate.endswith("]")):
                    current_ctc_key = candidate
                    ctc_repaints.setdefault(current_ctc_key, {})
                i += 2
                continue

            if ln == "[repaint]":
                repaint_name = lines[i + 1].strip() if i + 1 < len(lines) else "UNKNOWN_REPAINT"
                i += 2
                slotmap: Dict[str, str] = {}
                while i < len(lines):
                    if lines[i].lower() in ("[repaint]", "[ctc]"):
                        break
                    slot = lines[i].strip()
                    tex = lines[i + 1].strip() if i + 1 < len(lines) else ""
                    if not tex or (tex.startswith("[") and tex.endswith("]")):
                        i += 1
                        continue
                    slotmap[slot.lower()] = normpath_omsi(tex)
                    i += 2

                if current_ctc_key is None:
                    current_ctc_key = "__NO_CTC_KEY__"
                    ctc_repaints.setdefault(current_ctc_key, {})
                if slotmap:
                    ctc_repaints[current_ctc_key][repaint_name] = slotmap
                continue

            i += 1

    return CtiRepaintIndex(item_repaints=item_repaints, ctc_repaints=ctc_repaints)



def find_bus_files(vehicle_dir: Path) -> List[Path]:
    return sorted(vehicle_dir.rglob("*.bus"))


def find_hof_files(vehicle_dir: Path) -> List[Path]:
    return sorted(vehicle_dir.rglob("*.hof"))


def find_cti_files(vehicle_dir: Path) -> List[Path]:
    cti_files: List[Path] = []
    for p in vehicle_dir.rglob("*"):
        if is_texture_folder(p):
            cti_files.extend(sorted(p.rglob("*.cti")))
    seen = set()
    unique = []
    for f in cti_files:
        rp = f.resolve()
        if rp not in seen:
            seen.add(rp)
            unique.append(f)
    return unique


def should_include_vehicle_folder(vehicle_folder: Path) -> bool:

    has_bus = any(vehicle_folder.rglob("*.bus"))
    has_hof = any(vehicle_folder.rglob("*.hof"))

    if not has_hof:
        return False

    has_cti = False
    for p in vehicle_folder.rglob("*"):
        if is_texture_folder(p):
            if any(p.rglob("*.cti")):
                has_cti = True
                break

    return has_bus or has_cti



def attach_repaints(
    variant: BusVariant,
    vehicle_dir: Path,
    all_item_repaints: Dict[str, Dict[str, str]],
    item_sources: Dict[str, List[str]],
    all_ctc_repaints: Dict[str, Dict[str, Dict[str, str]]],
    ctc_sources: Dict[Tuple[str, str], List[str]],
    parse_model: bool,
) -> None:
    model_slots: Set[str] = set()

    if parse_model and variant.model_cfg:
        model_path = vehicle_dir / variant.model_cfg
        model_slots = parse_model_slots(model_path)
        variant.model_slots = sorted(model_slots)

    if variant.repaint_mode == "unknown":
        if variant.repaint_ctc_key:
            variant.repaint_mode = "ctc"
        elif all_item_repaints:
            variant.repaint_mode = "item"

    repaints_out: Dict[str, Dict[str, str]] = {}
    sources_out: Dict[str, List[str]] = {}

    if variant.repaint_mode == "ctc":
        key = variant.repaint_ctc_key or "__NO_CTC_KEY__"
        pool = all_ctc_repaints.get(key, {})
        for rname, slotmap in pool.items():
            repaints_out[rname] = slotmap
            sources_out[rname] = sorted(set(ctc_sources.get((key, rname), [])))

    elif variant.repaint_mode == "item":
        if model_slots:
            for rname, slotmap in all_item_repaints.items():
                if any(slot in model_slots for slot in slotmap.keys()):
                    repaints_out[rname] = slotmap
                    sources_out[rname] = sorted(set(item_sources.get(rname, [])))
        else:
            for rname, slotmap in all_item_repaints.items():
                repaints_out[rname] = slotmap
                sources_out[rname] = sorted(set(item_sources.get(rname, [])))

    variant.repaints = repaints_out
    variant.repaint_sources = sources_out


def analyze_vehicles(vehicles_dir: Path, parse_model: bool = True) -> dict:
    out: Dict[str, dict] = {}

    for vehicle_folder in sorted([p for p in vehicles_dir.iterdir() if p.is_dir()]):
        if not should_include_vehicle_folder(vehicle_folder):
            continue

        bus_files = find_bus_files(vehicle_folder)
        cti_files = find_cti_files(vehicle_folder)
        hof_files = find_hof_files(vehicle_folder)

        all_item_repaints: Dict[str, Dict[str, str]] = {}
        item_sources: Dict[str, List[str]] = {}

        all_ctc_repaints: Dict[str, Dict[str, Dict[str, str]]] = {}
        ctc_sources: Dict[Tuple[str, str], List[str]] = {}

        for cti in cti_files:
            idx = parse_cti(cti)

            for rname, slotmap in idx.item_repaints.items():
                all_item_repaints.setdefault(rname, {}).update(slotmap)
                item_sources.setdefault(rname, []).append(str(cti.relative_to(vehicle_folder)))

            for ctc_key, repaints in idx.ctc_repaints.items():
                all_ctc_repaints.setdefault(ctc_key, {})
                for rname, slotmap in repaints.items():
                    all_ctc_repaints[ctc_key].setdefault(rname, {}).update(slotmap)
                    ctc_sources.setdefault((ctc_key, rname), []).append(str(cti.relative_to(vehicle_folder)))

        variants: List[BusVariant] = []
        for bus in bus_files:
            v = parse_bus(bus)
            attach_repaints(
                variant=v,
                vehicle_dir=vehicle_folder,
                all_item_repaints=all_item_repaints,
                item_sources=item_sources,
                all_ctc_repaints=all_ctc_repaints,
                ctc_sources=ctc_sources,
                parse_model=parse_model
            )
            variants.append(v)
        model_friendly = None
        if variants:
            for v in variants:
                if v.friendly_name:
                    model_friendly = v.friendly_name
                    break

        out[vehicle_folder.name] = {
            "model_key": vehicle_folder.name,
            "model_friendly_name": model_friendly,  # best-effort
            "vehicle_folder": str(vehicle_folder),
            "bus_count": len(variants),
            "cti_count": len(cti_files),
            "hof_count": len(hof_files),

            "repaint_pools": {
                "item_repaints_count": len(all_item_repaints),
                "ctc_keys": {k: len(v) for k, v in all_ctc_repaints.items()},
            },

            "variants": [asdict(v) for v in variants],
        }

    return out

def is_frozen() -> bool:
    return getattr(sys, "frozen", False)

def exe_dir() -> Path:
    return Path(sys.executable).resolve().parent if is_frozen() else Path(__file__).resolve().parent

def prompt_for_vehicles_dir() -> Optional[Path]:
    try:
        import tkinter as tk
        from tkinter import filedialog, messagebox
    except Exception:
        return None

    root = tk.Tk()
    root.withdraw()
    root.attributes("-topmost", True)

    folder = filedialog.askdirectory(
        title="Select your OMSI 2 folder (root or Vehicles)."
    )
    if not folder:
        return None

    p = Path(folder)

    if (p / "Vehicles").is_dir():
        vehicles = p / "Vehicles"
    elif p.name.lower() == "vehicles" and p.is_dir():
        vehicles = p
    else:
        messagebox.showerror(
            "Invalid folder",
            "Please select your OMSI 2 root folder (it must contain a 'Vehicles' folder),\n"
            "or select the 'Vehicles' folder itself."
        )
        return None

    return vehicles


def main():
    parser = argparse.ArgumentParser(
        description="Analyze OMSI 2 Vehicles folder and build nested JSON of variants/repaints."
    )
    parser.add_argument("--vehicles", help="Path to OMSI 2 'Vehicles' folder (optional; if omitted you will be prompted)")
    parser.add_argument("--out", help="Output JSON filename (optional; default writes next to the EXE)")
    parser.add_argument("--no-model-parse", action="store_true", help="Skip parsing model*.cfg for farbschema slots")
    args = parser.parse_args()

    vehicles_dir: Optional[Path] = None
    if args.vehicles:
        p = Path(args.vehicles).expanduser().resolve()
        if (p / "Vehicles").is_dir():
            vehicles_dir = p / "Vehicles"
        else:
            vehicles_dir = p
    else:
        vehicles_dir = prompt_for_vehicles_dir()
        if vehicles_dir is None:
            return 0

    if not vehicles_dir.exists() or not vehicles_dir.is_dir():
        raise SystemExit(f"Vehicles folder not found or not a directory: {vehicles_dir}")

    data = analyze_vehicles(vehicles_dir, parse_model=(not args.no_model_parse))

    if args.out:
        out_path = Path(args.out).expanduser().resolve()
    else:
        out_path = exe_dir() / "omsi_vehicles_index.json"

    out_path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"Wrote: {out_path}")

    if is_frozen():
        try:
            import tkinter as tk
            from tkinter import messagebox
            root = tk.Tk()
            root.withdraw()
            root.attributes("-topmost", True)
            messagebox.showinfo("Done", f"Created:\n{out_path}.\nNow open the 'omsi_index.html' file and load in the json file just created in this folder.")
        except Exception:
            pass

    return 0



if __name__ == "__main__":
    main()
