from datetime import datetime, timedelta
from pathlib import Path
import json
import math

from .product_definitions import PRODUCT_DEFINITIONS, get_product_definition

PROJECT_ROOT = Path(__file__).resolve().parents[2]
SOURCE_DIR = PROJECT_ROOT / "data" / "display" / "grib"
GRID_DIR = PROJECT_ROOT / "data" / "grid"
OUTPUT_DIR = PROJECT_ROOT / "data" / "display" / "products"
META_DIR = OUTPUT_DIR / "meta"
LEGEND_DIR = OUTPUT_DIR / "legend"

for path in (OUTPUT_DIR, META_DIR, LEGEND_DIR):
    path.mkdir(parents=True, exist_ok=True)


def parse_source_name(path):
    stem = path.name
    if stem.endswith(".grib2"):
        stem = stem[:-6]

    parts = stem.split("_")
    if len(parts) < 3:
        raise ValueError(f"unexpected display source filename: {path.name}")

    date = parts[0]
    cycle = parts[1]
    forecast = parts[2]
    product_parts = parts[3:]

    if not product_parts:
        product = "wind"
    elif product_parts == ["wave"]:
        product = "wave"
    elif product_parts == ["rain"]:
        product = "rain"
    elif product_parts == ["pressure"]:
        product = "pressure-isoline"
    else:
        raise ValueError(f"unexpected display source filename: {path.name}")

    return date, cycle, forecast, product


def frame_time_from_parts(date, cycle, forecast):
    cycle_time = datetime.strptime(f"{date}{cycle}", "%Y%m%d%H")
    forecast_hour = int(forecast.removeprefix("f"))
    return cycle_time + timedelta(hours=forecast_hour)


def iso_frame_key(date, cycle, forecast):
    return frame_time_from_parts(date, cycle, forecast).strftime("%Y%m%dT%H%MZ")


def find_grid_payload(date, cycle, forecast):
    grid_path = GRID_DIR / f"grid_{date}_{cycle}_{forecast}.json"
    if not grid_path.exists():
        return None
    with grid_path.open(encoding="utf-8") as file_handle:
        return json.load(file_handle)


def flatten_numbers(values):
    flat = []
    for row in values:
        for value in row:
            if value is None:
                continue
            if isinstance(value, float) and math.isnan(value):
                continue
            flat.append(float(value))
    return flat


def extract_values_for_product(product, grid):
    field_map = {
        "wind": "wind_speed",
        "wave": "wave_h",
        "rain": "rain",
        "pressure-isoline": "pressure",
    }
    field_name = field_map[product]
    values = grid["grid"].get(field_name)
    if values is None:
        return []
    return flatten_numbers(values)


def compute_data_range(product, date, cycle, forecast):
    grid_payload = find_grid_payload(date, cycle, forecast)
    if grid_payload is None:
        return {"min": None, "max": None}

    values = extract_values_for_product(product, grid_payload)
    if not values:
        return {"min": None, "max": None}
    return {"min": min(values), "max": max(values)}


def build_meta(product, time_key, date, cycle, forecast):
    definition = get_product_definition(product)
    data_range = compute_data_range(product, date, cycle, forecast)
    return {
        "product": product,
        "time": time_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"],
        "data_min": data_range["min"],
        "data_max": data_range["max"],
        "no_data": definition["no_data"],
        "supported_zoom": definition["supported_zoom"],
        "opacity_suggestion": definition["opacity"],
        "path": definition["path_template"].format(time=time_key, z="{z}", x="{x}", y="{y}"),
    }


def build_legend(product):
    definition = get_product_definition(product)
    legend = definition["legend"]
    payload = {
        "product": product,
        "title": definition["title"],
        "subtitle": definition["subtitle"],
        "unit": definition["unit"],
        "scale_type": legend["scale_type"],
        "legend_sections": legend["sections"],
        "color_stops": legend["color_stops"],
    }
    if "contour_levels" in legend:
        payload["contour_levels"] = legend["contour_levels"]
    return payload


def collect_frame_inventory():
    inventory = {product: set() for product in PRODUCT_DEFINITIONS}
    source_files = sorted(SOURCE_DIR.glob("*.grib2"))
    for path in source_files:
        date, cycle, forecast, product = parse_source_name(path)
        inventory[product].add(
            (
                iso_frame_key(date, cycle, forecast),
                date,
                cycle,
                forecast,
            )
        )
    return inventory


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 build_all_products():
    inventory = collect_frame_inventory()
    frame_index = {
        "available_frames": [],
        "frame_step_hours": 3,
        "earliest": None,
        "latest": None,
        "product_availability": {},
    }

    all_frames = set()
    for product, items in inventory.items():
        sorted_items = sorted(items)
        product_frames = [time_key for time_key, _, _, _ in sorted_items]
        frame_index["product_availability"][product] = product_frames
        all_frames.update(product_frames)

        write_json(LEGEND_DIR / f"{product}.json", build_legend(product))
        for time_key, date, cycle, forecast in sorted_items:
            meta_payload = build_meta(product, time_key, date, cycle, forecast)
            write_json(META_DIR / product / f"{time_key}.json", meta_payload)

    if all_frames:
        ordered_frames = sorted(all_frames)
        frame_index["available_frames"] = ordered_frames
        frame_index["earliest"] = ordered_frames[0]
        frame_index["latest"] = ordered_frames[-1]

    write_json(OUTPUT_DIR / "frame-index.json", frame_index)
    write_json(
        OUTPUT_DIR / "product-index.json",
        {
            "products": [
                {
                    "product": product,
                    "display_type": PRODUCT_DEFINITIONS[product]["display_type"],
                    "legend": f"/weather-display/legend/{product}",
                    "meta": f"/weather-display/meta/{product}" + "/{time}",
                }
                for product in PRODUCT_DEFINITIONS
            ]
        },
    )
    return frame_index


def main():
    build_all_products()


if __name__ == "__main__":
    main()
