#!/usr/bin/env python3
from __future__ import annotations

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

from patch_traffic_lanes import FLAGS, as_int, has_flag, looks_like_lane


def visit(value: Any):
    if isinstance(value, dict):
        lanes = value.get("lanes")
        if isinstance(lanes, list) and any(looks_like_lane(lane) for lane in lanes):
            for lane in lanes:
                if looks_like_lane(lane):
                    yield lane
        for child in value.values():
            yield from visit(child)
    elif isinstance(value, list):
        for child in value:
            yield from visit(child)


def bucket(flags: int) -> str:
    if has_flag(flags, "Highway"):
        return "highway"
    if has_flag(flags, "Pavement"):
        return "pavement"
    if has_flag(flags, "CrossWalk"):
        return "crosswalk"
    if has_flag(flags, "Road"):
        return "road"
    return "other"


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("json_file", type=Path)
    args = parser.parse_args()

    data = json.loads(args.json_file.read_text(encoding="utf-8"))
    speeds: dict[str, Counter[int]] = {
        "all": Counter(),
        "highway": Counter(),
        "road": Counter(),
        "pavement": Counter(),
        "crosswalk": Counter(),
        "other": Counter(),
    }
    flags_counter: Counter[int] = Counter()

    for lane in visit(data):
        flags = as_int(lane.get("flags"))
        speed = as_int(lane.get("maxSpeed"))
        if flags is None or speed is None:
            continue
        speeds["all"][speed] += 1
        speeds[bucket(flags)][speed] += 1
        flags_counter[flags] += 1

    for name, counter in speeds.items():
        print(f"{name}: {sum(counter.values())} lanes")
        print("  " + ", ".join(f"{speed}:{count}" for speed, count in counter.most_common(20)))

    print("top flags:")
    for flags, count in flags_counter.most_common(20):
        names = [name for name, bit in FLAGS.items() if flags & bit]
        print(f"  {flags}: {count} ({', '.join(names)})")

    return 0


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