#!/usr/bin/env python3
"""Merge OSM layers into a single master GeoPackage, enriched for rendering.

Reads data/raw/osm/{mode}.gpkg (lines + stops) and data/processed/area.gpkg,
and writes data/processed/master.gpkg with layers "lines" and "stops".

Enrichment:
  - style      : from config/modes.yaml (mode -> style key)
  - colour     : GTFS route_color (if present) -> OSM colour tag -> palette
  - bus_category : A / C / S / regular (regex from styling.yaml)
Lines and stops are clipped to the City Pass area polygon.
"""
import re
from pathlib import Path

import geopandas as gpd
import pandas as pd
import yaml

from _common import CONFIG, OSM_RAW, PROCESSED

MIXED_CRS_WARN = "mixed CRS"


def load_modes():
    return yaml.safe_load((CONFIG / "modes.yaml").read_text())["modes"]


def load_styling():
    return yaml.safe_load((CONFIG / "styling.yaml").read_text())


def collect_lines(modes):
    frames = []
    for mode, cfg in modes.items():
        p = OSM_RAW / f"{mode}.gpkg"
        if not p.exists():
            continue
        g = gpd.read_file(p, layer="lines")
        if g.empty:
            continue
        g["mode"] = mode
        g["style"] = cfg["style"]
        g["gtfs_route_type"] = cfg.get("gtfs_route_type")
        frames.append(g)
    if not frames:
        return gpd.GeoDataFrame(columns=["mode", "ref", "style", "geometry"], crs="EPSG:4326")
    return pd.concat(frames, ignore_index=True)


def collect_stops(modes):
    frames = []
    for mode, cfg in modes.items():
        p = OSM_RAW / f"{mode}.gpkg"
        if not p.exists():
            continue
        try:
            g = gpd.read_file(p, layer="stops")
        except Exception:
            continue
        if g.empty:
            continue
        g["mode"] = mode
        g["style"] = cfg["style"]
        frames.append(g)
    if not frames:
        return None
    return pd.concat(frames, ignore_index=True)


def gtfs_colour_map():
    """Return {(route_short_name, route_type): route_color} if GTFS data exists."""
    csv = PROCESSED / "route_colors.csv"
    if not csv.exists():
        return {}
    df = pd.read_csv(csv)
    out = {}
    for _, r in df.iterrows():
        name = str(r.get("route_short_name") or "").strip()
        if name:
            out[(name, r.get("route_type"))] = r.get("route_color")
    return out


def categorise_bus(ref, patterns):
    if ref is None:
        return "regular"
    s = str(ref)
    for cat, pat in patterns.items():
        if re.search(pat, s):
            return cat
    return "regular"


def resolve_colour(row, palette, gtfs):
    style = row["style"]
    ref = row.get("ref")
    colour = row.get("colour")
    # 1. GTFS (if available)
    if gtfs and isinstance(ref, str):
        c = gtfs.get((ref, row.get("gtfs_route_type")))
        if c:
            return str(c)
    # 2. OSM colour tag (only accept real hex/named strings, not NA/None)
    if isinstance(colour, str) and colour.strip():
        return colour
    # 3. palette
    p = palette.get(style, {})
    if style == "bus":
        cat = row.get("bus_category")
        if not isinstance(cat, str):
            cat = "regular"
        return p.get(cat) or p.get("default")
    key = ref if isinstance(ref, str) else None
    return p.get(key) or p.get("default")


def main():
    modes = load_modes()
    styling = load_styling()
    palette = styling["palette"]
    bus_patterns = styling["bus_filters"]["categories"]

    area = gpd.read_file(PROCESSED / "area.gpkg")
    area_geom = area.geometry.union_all()

    lines = collect_lines(modes)
    print(f"collected {len(lines)} raw lines across modes", flush=True)

    # bus categorisation (only meaningful for bus style)
    lines["bus_category"] = [
        categorise_bus(r, bus_patterns) if s == "bus" else None
        for r, s in zip(lines.get("ref"), lines["style"])
    ]

    gtfs = gtfs_colour_map()
    if gtfs:
        print(f"GTFS colours loaded: {len(gtfs)} routes", flush=True)
    else:
        print("no GTFS colour data; using OSM colour -> palette", flush=True)

    lines["colour_final"] = [
        resolve_colour(r, palette, gtfs) for _, r in lines.iterrows()
    ]

    # clip to area
    lines = lines[~lines.geometry.isna() & lines.geometry.is_valid]
    clipped = gpd.clip(lines, area_geom)
    print(f"clipped to area: {len(clipped)} line features remain", flush=True)

    out_cols = [
        "mode", "style", "ref", "name", "network", "operator", "route",
        "colour", "bus_category", "colour_final", "geometry",
    ]
    for c in out_cols:
        if c not in clipped.columns:
            clipped[c] = None
    clipped = clipped.set_geometry("geometry")
    clipped = gpd.GeoDataFrame(clipped[out_cols], crs="EPSG:4326")

    PROCESSED.mkdir(parents=True, exist_ok=True)
    # drop existing master.gpkg so layer overwrite is clean
    out = PROCESSED / "master.gpkg"
    if out.exists():
        out.unlink()
    clipped.to_file(out, driver="GPKG", layer="lines")
    print(f"wrote lines layer: {len(clipped)} features", flush=True)

    # stops (stop_positions along routes)
    stops = collect_stops(modes)
    if stops is not None and not stops.empty:
        stops = stops[~stops.geometry.isna() & stops.geometry.is_valid]
        stops_clipped = gpd.clip(stops, area_geom)
        keep = ["mode", "style", "name", "ref", "public_transport", "railway", "route_refs", "geometry"]
        for c in keep:
            if c not in stops_clipped.columns:
                stops_clipped[c] = None
        stops_clipped = gpd.GeoDataFrame(stops_clipped[keep], crs="EPSG:4326")
        stops_clipped.to_file(out, driver="GPKG", layer="stops")
        print(f"wrote stops layer: {len(stops_clipped)} features", flush=True)
    else:
        print("no stops to write", flush=True)

    # stations (named station nodes/areas) for labelling
    st_path = OSM_RAW / "stations.gpkg"
    if st_path.exists():
        st = gpd.read_file(st_path, layer="stations")
        if not st.empty:
            st = st[~st.geometry.isna() & st.geometry.is_valid]
            st_clipped = gpd.clip(st, area_geom)
            st_clipped = gpd.GeoDataFrame(st_clipped, crs="EPSG:4326")
            st_clipped.to_file(out, driver="GPKG", layer="stations")
            print(f"wrote stations layer: {len(st_clipped)} named stations", flush=True)
    else:
        print("no stations file; labelling will be limited", flush=True)

    # summary by style
    print("\nsummary by style:")
    for style, grp in clipped.groupby("style"):
        refs = grp["ref"].dropna().unique()
        print(f"  {style:10s}: {len(grp):4d} feats, {len(refs):3d} refs")


if __name__ == "__main__":
    main()