from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path
import json
import os

import mercantile
import numpy as np

try:
    import mapbox_vector_tile
except ModuleNotFoundError:  # pragma: no cover
    mapbox_vector_tile = None

from .product_definitions import PRODUCT_DEFINITIONS, get_product_definition

PROJECT_ROOT = Path(__file__).resolve().parents[2]
GRID_DIR = PROJECT_ROOT / "data" / "grid"
OUTPUT_ROOT = Path("/home/wwwroot/weather/display")
VECTOR_ROOT = OUTPUT_ROOT / "vector" / "pressure-isoline"
META_ROOT = OUTPUT_ROOT / "meta" / "pressure-isoline"
LEGEND_ROOT = OUTPUT_ROOT / "legend"
FRAME_INDEX_PATH = OUTPUT_ROOT / "frame-index.json"
PRODUCT_INDEX_PATH = OUTPUT_ROOT / "product-index.json"

DEFAULT_ZOOMS = [2, 4, 6, 8]
JAPAN_DISPLAY_BOUNDS = (120.0, 20.0, 150.0, 50.0)
FIELD_NAME = "pressure"

# Marching squares lookup where each edge is between two cell corners:
# 0 bottom, 1 right, 2 top, 3 left
CASE_TO_EDGES = {
    0: [],
    1: [(3, 0)],
    2: [(0, 1)],
    3: [(3, 1)],
    4: [(1, 2)],
    5: [(3, 2), (0, 1)],
    6: [(0, 2)],
    7: [(3, 2)],
    8: [(2, 3)],
    9: [(0, 2)],
    10: [(0, 1), (2, 3)],
    11: [(1, 2)],
    12: [(3, 1)],
    13: [(0, 1)],
    14: [(3, 0)],
    15: [],
}


def get_zoom_levels():
    raw_value = os.environ.get("DISPLAY_VECTOR_ZOOMS", "").strip()
    if not raw_value:
        return DEFAULT_ZOOMS

    zooms = []
    for chunk in raw_value.split(","):
        chunk = chunk.strip()
        if not chunk:
            continue
        zooms.append(int(chunk))
    return sorted(set(zooms))


def zoom_display_bounds(zoom):
    if zoom <= 2:
        return None
    return JAPAN_DISPLAY_BOUNDS


def grid_time_to_frame_key(grid_time):
    date, cycle, forecast = grid_time.split("_")
    cycle_time = datetime.strptime(f"{date}{cycle}", "%Y%m%d%H")
    forecast_hour = int(forecast.removeprefix("f"))
    frame_time = cycle_time + timedelta(hours=forecast_hour)
    return frame_time.strftime("%Y%m%dT%H%MZ")


def load_grid(path):
    with path.open(encoding="utf-8") as file_handle:
        payload = json.loads(file_handle.read().replace("NaN", "null"))

    grid = payload["grid"]
    return {
        "grid_time": payload["time"],
        "frame_key": grid_time_to_frame_key(payload["time"]),
        "lat": np.asarray(grid["lat"], dtype=np.float32),
        "lon": np.asarray(grid["lon"], dtype=np.float32),
        "pressure": np.asarray(grid[FIELD_NAME], dtype=np.float32),
    }


def interpolate_point(edge_id, lon0, lat0, lon1, lat1, v0, v1, level):
    if v1 == v0:
        ratio = 0.5
    else:
        ratio = float((level - v0) / (v1 - v0))
    ratio = min(max(ratio, 0.0), 1.0)

    if edge_id == 0:  # bottom
        return (lon0 + (lon1 - lon0) * ratio, lat1)
    if edge_id == 1:  # right
        return (lon1, lat1 + (lat0 - lat1) * ratio)
    if edge_id == 2:  # top
        return (lon0 + (lon1 - lon0) * ratio, lat0)
    return (lon0, lat1 + (lat0 - lat1) * ratio)  # left


def cell_edge_point(edge_id, lon0, lat0, lon1, lat1, v00, v10, v11, v01, level):
    if edge_id == 0:
        return interpolate_point(edge_id, lon0, lat0, lon1, lat1, v01, v11, level)
    if edge_id == 1:
        return interpolate_point(edge_id, lon0, lat0, lon1, lat1, v11, v10, level)
    if edge_id == 2:
        return interpolate_point(edge_id, lon0, lat0, lon1, lat1, v00, v10, level)
    return interpolate_point(edge_id, lon0, lat0, lon1, lat1, v01, v00, level)


def marching_squares_segments(grid_info, level):
    latitudes = grid_info["lat"]
    longitudes = grid_info["lon"]
    values = grid_info["pressure"]
    segments = []

    for row in range(len(latitudes) - 1):
        lat0 = float(latitudes[row])
        lat1 = float(latitudes[row + 1])
        for col in range(len(longitudes) - 1):
            lon0 = float(longitudes[col])
            lon1 = float(longitudes[col + 1])

            v00 = values[row, col]
            v10 = values[row, col + 1]
            v01 = values[row + 1, col]
            v11 = values[row + 1, col + 1]

            if np.isnan(v00) or np.isnan(v10) or np.isnan(v01) or np.isnan(v11):
                continue

            case_id = 0
            if v01 >= level:
                case_id |= 1
            if v11 >= level:
                case_id |= 2
            if v10 >= level:
                case_id |= 4
            if v00 >= level:
                case_id |= 8

            for edge_a, edge_b in CASE_TO_EDGES[case_id]:
                point_a = cell_edge_point(edge_a, lon0, lat0, lon1, lat1, v00, v10, v11, v01, level)
                point_b = cell_edge_point(edge_b, lon0, lat0, lon1, lat1, v00, v10, v11, v01, level)
                if point_a != point_b:
                    segments.append((point_a, point_b))

    return segments


def tiles_for_segment(point_a, point_b, zoom):
    lon_values = [point_a[0], point_b[0]]
    lat_values = [point_a[1], point_b[1]]
    west = min(lon_values)
    east = max(lon_values)
    south = min(lat_values)
    north = max(lat_values)
    coverage = zoom_display_bounds(zoom)
    if coverage is not None:
        coverage_west, coverage_south, coverage_east, coverage_north = coverage
        west = max(west, coverage_west)
        south = max(south, coverage_south)
        east = min(east, coverage_east)
        north = min(north, coverage_north)
        if west >= east or south >= north:
            return []
    return mercantile.tiles(west, south, east, north, zooms=[zoom])


def encode_tile(frame_key, tile, tile_features):
    bounds = mercantile.bounds(tile)
    tile_dir = VECTOR_ROOT / frame_key / str(tile.z) / str(tile.x)
    tile_dir.mkdir(parents=True, exist_ok=True)
    tile_path = tile_dir / f"{tile.y}.pbf"

    features = []
    for level, point_a, point_b in tile_features:
        features.append(
            {
                "geometry": {
                    "type": "LineString",
                    "coordinates": [point_a, point_b],
                },
                "properties": {
                    "level": int(level),
                    "unit": "Pa",
                },
            }
        )

    tile_data = mapbox_vector_tile.encode(
        {"name": "pressure_isoline", "features": features},
        default_options={
            "quantize_bounds": (bounds.west, bounds.south, bounds.east, bounds.north),
        },
    )
    tile_path.write_bytes(tile_data)


def write_json(path, payload):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as file_handle:
        json.dump(payload, file_handle, ensure_ascii=False, indent=2)


def update_supporting_metadata(frame_keys):
    definition = get_product_definition("pressure-isoline")
    legend = definition["legend"]
    write_json(
        LEGEND_ROOT / "pressure-isoline.json",
        {
            "product": "pressure-isoline",
            "title": definition["title"],
            "subtitle": definition["subtitle"],
            "unit": definition["unit"],
            "scale_type": legend["scale_type"],
            "legend_sections": legend["sections"],
            "color_stops": legend["color_stops"],
            "contour_levels": legend["contour_levels"],
        },
    )

    for frame_key in frame_keys:
        write_json(
            META_ROOT / f"{frame_key}.json",
            {
                "product": "pressure-isoline",
                "time": frame_key,
                "unit": definition["unit"],
                "display_type": definition["display_type"],
                "palette_id": definition["palette_id"],
                "recommended_min": definition["recommended_range"]["min"],
                "recommended_max": definition["recommended_range"]["max"],
                "no_data": definition["no_data"],
                "supported_zoom": definition["supported_zoom"],
                "opacity_suggestion": definition["opacity"],
                "path": definition["path_template"].format(time=frame_key, z="{z}", x="{x}", y="{y}"),
            },
        )

    if PRODUCT_INDEX_PATH.exists():
        with PRODUCT_INDEX_PATH.open(encoding="utf-8") as file_handle:
            product_index = json.load(file_handle)
    else:
        product_index = {"products": []}

    products = {item["product"]: item for item in product_index.get("products", [])}
    products["pressure-isoline"] = {
        "product": "pressure-isoline",
        "display_type": "vector",
        "legend": "/weather-display/legend/pressure-isoline",
        "meta": "/weather-display/meta/pressure-isoline/{time}",
    }
    product_index["products"] = list(products.values())
    write_json(PRODUCT_INDEX_PATH, product_index)

    if FRAME_INDEX_PATH.exists():
        with FRAME_INDEX_PATH.open(encoding="utf-8") as file_handle:
            frame_index = json.load(file_handle)
    else:
        frame_index = {"available_frames": [], "product_availability": {}}

    frame_index.setdefault("product_availability", {})
    frame_index["product_availability"]["pressure-isoline"] = list(frame_keys)
    write_json(FRAME_INDEX_PATH, frame_index)


def generate_all():
    if mapbox_vector_tile is None:
        raise RuntimeError("mapbox-vector-tile is required to generate pressure isoline vector tiles")

    zooms = get_zoom_levels()
    contour_levels = PRODUCT_DEFINITIONS["pressure-isoline"]["legend"]["contour_levels"]
    frame_keys = []

    for path in sorted(GRID_DIR.glob("grid_*.json")):
        grid_info = load_grid(path)
        frame_keys.append(grid_info["frame_key"])
        print("processing", path.name)

        segments = []
        for level in contour_levels:
            for point_a, point_b in marching_squares_segments(grid_info, level):
                segments.append((level, point_a, point_b))

        for zoom in zooms:
            buckets = defaultdict(list)
            for level, point_a, point_b in segments:
                for tile in tiles_for_segment(point_a, point_b, zoom):
                    buckets[(tile.x, tile.y)].append((level, point_a, point_b))

            for (tile_x, tile_y), tile_features in buckets.items():
                encode_tile(grid_info["frame_key"], mercantile.Tile(x=tile_x, y=tile_y, z=zoom), tile_features)

    update_supporting_metadata(sorted(set(frame_keys)))


def main():
    generate_all()


if __name__ == "__main__":
    main()
