import mercantile
import numpy as np
from shapely import contains_xy
from shapely.geometry import box
from shapely.ops import unary_union
from shapely.strtree import STRtree

from src.geo_mask.foundation import SOURCE_ASSET
from src.geo_mask.coast_transition import (
    DEFAULT_TRANSITION_RADIUS_BY_ZOOM,
    build_coast_transition_band,
)
from src.geo_mask.raster_tile_generator import load_land_geometries

from .mask_policy import DISPLAY_MASK_POLICIES


def build_land_mask_context():
    if not SOURCE_ASSET.exists():
        return None

    geometries = load_land_geometries(SOURCE_ASSET)
    return {
        "geometries": geometries,
        "tree": STRtree(geometries),
    }


def candidate_union(context, tile):
    if context is None:
        return None

    bounds = mercantile.bounds(tile)
    tile_box = box(bounds.west, bounds.south, bounds.east, bounds.north)
    candidates = context["tree"].query(tile_box)
    if len(candidates) == 0:
        return None

    geometries = [context["geometries"][int(index)] for index in candidates]
    tile_geometries = [geometry for geometry in geometries if geometry.intersects(tile_box)]
    if not tile_geometries:
        return None

    return unary_union(tile_geometries)


def sample_land_mask(context, tile, lon_grid, lat_grid):
    if context is None:
        return np.zeros(lon_grid.shape, dtype=bool)

    geom = candidate_union(context, tile)
    if geom is None:
        return np.zeros(lon_grid.shape, dtype=bool)

    return contains_xy(geom, lon_grid, lat_grid)


def build_coast_mask_cache(land_mask_cache):
    coast_cache = {}
    for tile_key, land_mask in land_mask_cache.items():
        zoom = tile_key[0]
        radius = DEFAULT_TRANSITION_RADIUS_BY_ZOOM.get(zoom, 6.0)
        coast_cache[tile_key] = build_coast_transition_band(land_mask, radius)
    return coast_cache


def apply_mask_policy(product, rgba, land_mask, coast_band):
    policy = DISPLAY_MASK_POLICIES.get(product, {"sea": "normal", "land": "normal", "coast": "normal"})
    if policy["land"] == "normal":
        if policy.get("coast") == "normal":
            return rgba

    active_land = land_mask & (rgba[:, :, 3] > 0)
    active_sea = (~land_mask) & (rgba[:, :, 3] > 0)
    output = rgba.copy()

    if policy["land"] == "mask":
        output[active_land] = np.array([0, 0, 0, 0], dtype=np.uint8)
    elif policy["land"] == "attenuate" and np.any(active_land):
        rgb = output[:, :, :3].astype(np.float32)
        inland_blend = 0.65 - 0.25 * coast_band[active_land]
        rgb[active_land] = np.round(
            rgb[active_land] * (1.0 - inland_blend[:, None]) + 255.0 * inland_blend[:, None]
        )
        output[:, :, :3] = np.clip(rgb, 0, 255).astype(np.uint8)
        output[:, :, 3][active_land] = np.minimum(
            output[:, :, 3][active_land],
            np.round(90 + 70 * coast_band[active_land]).astype(np.uint8),
        )

    if policy.get("coast") == "soft-mask" and np.any(active_sea):
        fade = 1.0 - 0.7 * coast_band[active_sea]
        output[:, :, 3][active_sea] = np.round(output[:, :, 3][active_sea] * fade).astype(np.uint8)
    elif policy.get("coast") == "soft-attenuate" and np.any(active_sea):
        fade = 1.0 - 0.25 * coast_band[active_sea]
        output[:, :, 3][active_sea] = np.round(output[:, :, 3][active_sea] * fade).astype(np.uint8)

    return output
