#!/usr/bin/env python3
"""Patch WolvenKit JSON exports of Cyberpunk traffic lane resources.

The game GPS planner appears to be native and data-driven over traffic lanes.
This tool changes lane metadata that is likely to feed edge-cost selection:
`maxSpeed` and, optionally, non-highway speed caps.
"""

from __future__ import annotations

import argparse
import copy
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any


FLAGS = {
    "FromRoadSpline": 1,
    "Bidirectional": 2,
    "PatrolRoute": 4,
    "Pavement": 8,
    "Road": 16,
    "Intersection": 32,
    "NeverDeadEnd": 64,
    "TrafficDisabled": 128,
    "CrossWalk": 256,
    "GPSOnly": 512,
    "ShowDebug": 1024,
    "Blockade": 2048,
    "Yield": 4096,
    "NoAIDriving": 8192,
    "Highway": 16384,
    "NoAutodrive": 32768,
}


@dataclass
class PatchStats:
    files_seen: int = 0
    files_changed: int = 0
    lane_sets_seen: int = 0
    lanes_seen: int = 0
    lanes_changed: int = 0
    highway_lanes_changed: int = 0
    road_lanes_changed: int = 0
    capped_lanes_changed: int = 0


def unwrap(value: Any) -> Any:
    if isinstance(value, dict):
        for key in ("value", "Value", "$value", "_value"):
            if key in value and len(value) <= 3:
                return value[key]
    return value


def set_wrapped(container: dict[str, Any], key: str, value: Any) -> None:
    current = container[key]
    if isinstance(current, dict):
        for value_key in ("value", "Value", "$value", "_value"):
            if value_key in current:
                current[value_key] = value
                return
    container[key] = value


def as_int(value: Any) -> int | None:
    value = unwrap(value)
    if isinstance(value, bool):
        return None
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        return int(value)
    if isinstance(value, str):
        try:
            return int(value, 0)
        except ValueError:
            return None
    return None


def looks_like_lane(value: Any) -> bool:
    if not isinstance(value, dict):
        return False
    return "flags" in value and "maxSpeed" in value and (
        "length" in value or "playerGPSInfo" in value or "laneNumber" in value
    )


def has_flag(flags: int, flag: str) -> bool:
    return (flags & FLAGS[flag]) == FLAGS[flag]


def patch_lane(
    lane: dict[str, Any],
    *,
    highway_speed: int,
    road_speed_floor: int,
    non_highway_cap: int | None,
) -> tuple[bool, str | None]:
    flags = as_int(lane.get("flags"))
    speed = as_int(lane.get("maxSpeed"))
    if flags is None or speed is None:
        return False, None

    if has_flag(flags, "TrafficDisabled") or has_flag(flags, "Blockade"):
        return False, None
    if has_flag(flags, "Pavement") or has_flag(flags, "CrossWalk"):
        return False, None

    is_highway = has_flag(flags, "Highway")
    is_road = has_flag(flags, "Road") or has_flag(flags, "GPSOnly")

    target = speed
    reason: str | None = None

    if is_highway and speed < highway_speed:
        target = highway_speed
        reason = "highway"
    elif is_road and speed < road_speed_floor:
        target = road_speed_floor
        reason = "road"

    if non_highway_cap is not None and not is_highway and target > non_highway_cap:
        target = non_highway_cap
        reason = "cap"

    if target == speed:
        return False, None

    set_wrapped(lane, "maxSpeed", target)
    return True, reason


def walk_and_patch(value: Any, stats: PatchStats, args: argparse.Namespace) -> None:
    if isinstance(value, dict):
        lanes = value.get("lanes")
        if isinstance(lanes, list) and any(looks_like_lane(lane) for lane in lanes):
            stats.lane_sets_seen += 1
            for lane in lanes:
                if not looks_like_lane(lane):
                    continue
                stats.lanes_seen += 1
                changed, reason = patch_lane(
                    lane,
                    highway_speed=args.highway_speed,
                    road_speed_floor=args.road_speed_floor,
                    non_highway_cap=args.non_highway_cap,
                )
                if changed:
                    stats.lanes_changed += 1
                    if reason == "highway":
                        stats.highway_lanes_changed += 1
                    elif reason == "road":
                        stats.road_lanes_changed += 1
                    elif reason == "cap":
                        stats.capped_lanes_changed += 1

        for child in value.values():
            walk_and_patch(child, stats, args)
    elif isinstance(value, list):
        for child in value:
            walk_and_patch(child, stats, args)


def iter_json_files(path: Path) -> list[Path]:
    if path.is_file():
        return [path]
    return sorted(path.rglob("*.json"))


def output_path(input_file: Path, input_root: Path, output_root: Path) -> Path:
    if input_root.is_file():
        return output_root
    return output_root / input_file.relative_to(input_root)


def patch_file(input_file: Path, input_root: Path, output_root: Path, args: argparse.Namespace, stats: PatchStats) -> None:
    stats.files_seen += 1
    original = input_file.read_text(encoding="utf-8")
    data = json.loads(original)
    patched = copy.deepcopy(data)

    before_changed = stats.lanes_changed
    walk_and_patch(patched, stats, args)
    changed = stats.lanes_changed != before_changed
    if changed:
        stats.files_changed += 1

    destination = output_path(input_file, input_root, output_root)
    if args.dry_run:
        return

    destination.parent.mkdir(parents=True, exist_ok=True)
    if changed or args.copy_unchanged:
        destination.write_text(
            json.dumps(patched, indent=2, ensure_ascii=False) + "\n",
            encoding="utf-8",
        )
    elif input_root.is_file():
        destination.write_text(original, encoding="utf-8")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("input", type=Path, help="WolvenKit JSON file or export directory")
    parser.add_argument("output", type=Path, help="Patched JSON file or output directory")
    parser.add_argument("--highway-speed", type=int, default=25)
    parser.add_argument("--road-speed-floor", type=int, default=15)
    parser.add_argument("--non-highway-cap", type=int, default=None)
    parser.add_argument("--copy-unchanged", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if not args.input.exists():
        raise SystemExit(f"Input does not exist: {args.input}")

    stats = PatchStats()
    for input_file in iter_json_files(args.input):
        patch_file(input_file, args.input, args.output, args, stats)

    print(f"files seen: {stats.files_seen}")
    print(f"files changed: {stats.files_changed}")
    print(f"lane sets seen: {stats.lane_sets_seen}")
    print(f"lanes seen: {stats.lanes_seen}")
    print(f"lanes changed: {stats.lanes_changed}")
    print(f"highway boosts: {stats.highway_lanes_changed}")
    print(f"road boosts: {stats.road_lanes_changed}")
    print(f"non-highway caps: {stats.capped_lanes_changed}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
