#!/usr/bin/env python3
"""Patch traffic lane connection probabilities toward highway routing."""

from __future__ import annotations

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


FLAGS = {
    "Road": 16,
    "Pavement": 8,
    "TrafficDisabled": 128,
    "CrossWalk": 256,
    "GPSOnly": 512,
    "Blockade": 2048,
    "Highway": 16384,
}


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def flag_value(flags: Any) -> int:
    return int(flags or 0)


def has(flags: int, name: str) -> bool:
    return bool(flags & FLAGS[name])


def category(flags: int) -> str:
    if has(flags, "Highway"):
        return "highway"
    if has(flags, "GPSOnly"):
        return "gpsonly"
    if has(flags, "Road"):
        return "road"
    if has(flags, "Pavement"):
        return "pavement"
    return "other"


def write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("traffic_persistent_json", type=Path)
    parser.add_argument("lane_connections_json", type=Path)
    parser.add_argument("out_json", type=Path)
    parser.add_argument("--high", type=int, default=255)
    parser.add_argument("--low", type=int, default=1)
    args = parser.parse_args()

    traffic = load_json(args.traffic_persistent_json)
    lanes = traffic["Data"]["RootChunk"]["data"]["lanes"]
    lane_categories = [category(flag_value(lane.get("flags", 0))) for lane in lanes]

    connections = load_json(args.lane_connections_json)
    rows = connections["Data"]["RootChunk"]["data"]

    changed = 0
    high_set = 0
    low_set = 0
    touched_rows = 0

    for row in rows:
        source_index = row["index"]
        source_category = lane_categories[source_index]
        outlanes = row["value"]["outlanes"]
        if not outlanes:
            continue

        target_categories = [
            lane_categories[out["laneIndex"]]
            for out in outlanes
            if 0 <= out["laneIndex"] < len(lane_categories)
        ]
        has_highway_target = "highway" in target_categories
        has_non_highway_target = any(cat != "highway" for cat in target_categories)
        row_changed = False

        for out in outlanes:
            target_category = lane_categories[out["laneIndex"]]
            old = out["exitProbabilityCompressed"]
            new = old

            if target_category == "highway" and source_category in {"road", "gpsonly", "highway"}:
                new = args.high
            elif (
                source_category == "highway"
                and has_highway_target
                and target_category != "highway"
            ):
                new = args.low
            elif (
                source_category in {"road", "gpsonly"}
                and has_highway_target
                and has_non_highway_target
                and target_category != "highway"
            ):
                new = args.low

            if new != old:
                out["exitProbabilityCompressed"] = new
                changed += 1
                row_changed = True
                if new == args.high:
                    high_set += 1
                elif new == args.low:
                    low_set += 1

        if row_changed:
            touched_rows += 1

    write_json(args.out_json, connections)
    print(f"rows seen: {len(rows)}")
    print(f"rows touched: {touched_rows}")
    print(f"edges changed: {changed}")
    print(f"edges set high: {high_set}")
    print(f"edges set low: {low_set}")
    return 0


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