from __future__ import annotations

from collections import defaultdict
from pathlib import Path
import json

import numpy as np
from PIL import Image

from src import gfs_downloader
from src.grid_builder_v2 import OUTPUT_DIR as GRID_OUTPUT_DIR
from src.grid_builder_v2 import INPUT_DIR as GRIB_INPUT_DIR
from src.grid_builder_v2 import process_file

from .build_products import build_all_products
from .downloader import download_cycle, get_cycle
from .mask_compositor import apply_mask_policy
from .pressure_isoline_generator import (
    PRODUCT_DEFINITIONS,
    encode_tile,
    get_zoom_levels as get_vector_zoom_levels,
    load_grid as load_pressure_grid,
    marching_squares_segments,
    mercantile,
    tiles_for_segment,
)
from .raster_generator import (
    GRID_DIR,
    ensure_directories,
    generate_tiles_for_grid,
    get_zoom_levels as get_raster_zoom_levels,
    load_grid,
)

DISPLAY_ROOT = Path("/home/wwwroot/weather/display/raster")
MASK_ROOT = Path("/home/wwwroot/weather/geo-mask/land-sea")


def refresh_display_source_data(date: str, cycle: str) -> None:
    print(f"refreshing display source data for {date} {cycle}")
    download_cycle()


def refresh_regional_grib(date: str, cycle: str) -> None:
    print(f"refreshing regional grib for {date} {cycle}")
    session = gfs_downloader.requests.Session()
    session.headers["User-Agent"] = "weather-refresh/1.0"
    try:
        for forecast_hour in gfs_downloader.FORECAST_HOURS:
            forecast_hour_str = f"{forecast_hour:03d}"
            gfs_downloader.download_forecast(session, date, cycle, forecast_hour_str)
    finally:
        session.close()


def build_cycle_grids(date: str, cycle: str) -> list[Path]:
    print(f"building cycle grids for {date} {cycle}")
    output_paths: list[Path] = []
    grib_paths = sorted(GRIB_INPUT_DIR.glob(f"{date}_{cycle}_f*.grib2"))
    for path in grib_paths:
        if path.stem.endswith("_wave"):
            continue
        print(f"processing grid {path.name}")
        grid = process_file(path)
        output_path = GRID_OUTPUT_DIR / f"grid_{path.stem}.json"
        with output_path.open("w", encoding="utf-8") as file_handle:
            json.dump({"time": path.stem, "grid": grid}, file_handle)
        output_paths.append(output_path)
    return output_paths


def render_cycle_raster(grid_paths: list[Path]) -> None:
    print("rendering cycle raster")
    ensure_directories()
    zooms = get_raster_zoom_levels()
    for path in grid_paths:
        grid_info = load_grid(path)
        print(f"rendering raster {path.name}")
        generate_tiles_for_grid(grid_info, zooms, None, {}, {})


def apply_cycle_mask(date: str, cycle: str) -> None:
    print("applying geo mask to cycle raster")
    cycle_prefix = f"{date}T"
    products = ["wind", "wave", "rain", "pressure"]
    masked_count = 0

    for product in products:
        product_root = DISPLAY_ROOT / product
        if not product_root.exists():
            continue
        for frame_dir in sorted(product_root.iterdir()):
            if not frame_dir.is_dir() or not frame_dir.name.startswith(cycle_prefix):
                continue
            for tile_path in frame_dir.glob("*/*/*.png"):
                z = tile_path.parts[-3]
                x = tile_path.parts[-2]
                y = tile_path.name
                mask_path = MASK_ROOT / z / x / y
                if not mask_path.exists():
                    continue

                rgba = np.array(Image.open(tile_path).convert("RGBA"), dtype=np.uint8)
                mask_rgba = np.array(Image.open(mask_path).convert("RGBA"), dtype=np.uint8)
                land_mask = mask_rgba[:, :, 0] >= 128
                coast_band = mask_rgba[:, :, 1].astype(np.float32) / 255.0
                masked = apply_mask_policy(product, rgba, land_mask, coast_band)
                Image.fromarray(masked, mode="RGBA").save(tile_path)
                masked_count += 1

    print(f"applied mask to {masked_count} raster tiles")


def render_cycle_pressure_isolines(grid_paths: list[Path]) -> None:
    print("rendering cycle pressure isolines")
    contour_levels = PRODUCT_DEFINITIONS["pressure-isoline"]["legend"]["contour_levels"]
    zooms = get_vector_zoom_levels()

    for path in grid_paths:
        grid_info = load_pressure_grid(path)
        print(f"rendering pressure isolines {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,
                )


def main() -> None:
    date, cycle = get_cycle()
    print(f"scheduled refresh cycle: {date} {cycle}")
    refresh_display_source_data(date, cycle)
    refresh_regional_grib(date, cycle)
    grid_paths = build_cycle_grids(date, cycle)
    build_all_products()
    render_cycle_raster(grid_paths)
    apply_cycle_mask(date, cycle)
    render_cycle_pressure_isolines(grid_paths)
    print("scheduled refresh done")


if __name__ == "__main__":
    main()
