from __future__ import annotations

import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import rasterio
from pyproj import Geod, Transformer
from rasterio.crs import CRS
from rasterio.features import shapes
from rasterio.merge import merge
from rasterio.warp import (
    Resampling,
    calculate_default_transform,
    reproject,
    transform,
    transform_bounds,
)
from shapely.geometry import mapping, shape
from shapely.ops import transform as shapely_transform
from shapely.ops import unary_union

from app.core.dem import _dem_files, _sample_dataset
from app.core.geo import GeoPoint
from app.models.viewshed import ViewshedRequest


@dataclass(frozen=True)
class RasterResult:
    uri: str
    metadata: dict[str, Any]


_GEOD = Geod(ellps="WGS84")


def _utm_epsg(lon: float, lat: float) -> int:
    zone = int((lon + 180) // 6) + 1
    return 32600 + zone if lat >= 0 else 32700 + zone


def _intersects(
    a: tuple[float, float, float, float],
    b: tuple[float, float, float, float],
) -> bool:
    return not (a[2] < b[0] or a[0] > b[2] or a[3] < b[1] or a[1] > b[3])


def _radius_bbox(observer: GeoPoint, radius_m: float) -> tuple[float, float, float, float]:
    coords = [
        _GEOD.fwd(observer.lon, observer.lat, azimuth, radius_m)[:2]
        for azimuth in range(0, 360, 45)
    ]
    lons = [observer.lon, *(coord[0] for coord in coords)]
    lats = [observer.lat, *(coord[1] for coord in coords)]
    return min(lons), min(lats), max(lons), max(lats)


def _find_dem_tile(lat: float, lon: float, dem_path: Path) -> Path | None:
    for path in _dem_files(dem_path):
        with rasterio.open(path) as dataset:
            if _sample_dataset(dataset, lat, lon) is not None:
                return path
    return None


def _find_dem_tiles(
    dem_path: Path,
    bbox_wgs84: tuple[float, float, float, float],
) -> list[Path]:
    matches: list[Path] = []
    for path in _dem_files(dem_path):
        with rasterio.open(path) as dataset:
            bounds = dataset.bounds
            tile_bbox = (
                bounds.left,
                bounds.bottom,
                bounds.right,
                bounds.top,
            )
            if dataset.crs is not None and dataset.crs != CRS.from_epsg(4326):
                tile_bbox = transform_bounds(dataset.crs, CRS.from_epsg(4326), *tile_bbox)
            if _intersects(tile_bbox, bbox_wgs84):
                matches.append(path)
    return matches


def _prepare_metric_dem(
    source_paths: list[Path],
    observer: GeoPoint,
    radius_m: float,
    output_path: Path,
) -> tuple[Path, float, float, int]:
    if not source_paths:
        raise FileNotFoundError("No DEM tiles found for viewshed radius")

    dst_crs = CRS.from_epsg(_utm_epsg(observer.lon, observer.lat))
    bbox_wgs84 = _radius_bbox(observer, radius_m)
    datasets = [rasterio.open(path) for path in source_paths]
    try:
        src_crs = datasets[0].crs or CRS.from_epsg(4326)
        merge_bounds = bbox_wgs84
        if src_crs != CRS.from_epsg(4326):
            merge_bounds = transform_bounds(CRS.from_epsg(4326), src_crs, *bbox_wgs84)
        mosaic, src_transform = merge(datasets, bounds=merge_bounds)
        source_data = mosaic[0].astype(np.float32)
        height, width = source_data.shape
        src_bounds = rasterio.transform.array_bounds(height, width, src_transform)
        observer_x, observer_y = transform(
            CRS.from_epsg(4326),
            dst_crs,
            [observer.lon],
            [observer.lat],
        )
        transform_affine, width, height = calculate_default_transform(
            src_crs,
            dst_crs,
            width,
            height,
            *src_bounds,
        )
        data = np.empty((height, width), dtype=np.float32)
        reproject(
            source=source_data,
            destination=data,
            src_transform=src_transform,
            src_crs=src_crs,
            dst_transform=transform_affine,
            dst_crs=dst_crs,
            resampling=Resampling.bilinear,
        )
        profile = {
            "driver": "GTiff",
            "crs": dst_crs,
            "transform": transform_affine,
            "width": width,
            "height": height,
            "dtype": "float32",
            "count": 1,
            "nodata": -9999,
        }
        with rasterio.open(output_path, "w", **profile) as dst:
            dst.write(data, 1)
    finally:
        for dataset in datasets:
            dataset.close()
    return output_path, observer_x[0], observer_y[0], len(source_paths)


def _run_gdal_viewshed(
    dem_path: Path,
    observer_x: float,
    observer_y: float,
    observer_height: float,
    target_height: float,
    max_distance: float,
    output_path: Path,
    gdal_bin: str,
) -> None:
    command = [
        gdal_bin,
        "-ox",
        str(observer_x),
        "-oy",
        str(observer_y),
        "-oz",
        str(observer_height),
        "-tz",
        str(target_height),
        "-md",
        str(max_distance),
        "-vv",
        "255",
        "-iv",
        "0",
        "-ov",
        "0",
        str(dem_path),
        str(output_path),
    ]
    subprocess.run(command, check=True, capture_output=True, text=True)


def _to_wgs84(geom, src_crs: CRS | None):
    if src_crs is None:
        return geom
    crs_text = src_crs.to_string()
    if crs_text in {"EPSG:4326", "OGC:CRS84", "WGS84"}:
        return geom
    transformer = Transformer.from_crs(src_crs, "EPSG:4326", always_xy=True)
    return shapely_transform(transformer.transform, geom)


def _viewshed_geojson(viewshed_path: Path, observer: GeoPoint) -> dict[str, Any]:
    features: list[dict[str, Any]] = []
    with rasterio.open(viewshed_path) as dataset:
        data = dataset.read(1, masked=True)
        mask = (data == 255).astype(np.uint8)
        for geom, value in shapes(mask, mask=mask.astype(bool), transform=dataset.transform):
            if int(value) != 1:
                continue
            polygon = _to_wgs84(shape(geom), dataset.crs)
            if polygon.is_empty:
                continue
            features.append(
                {
                    "type": "Feature",
                    "properties": {"visible": True},
                    "geometry": mapping(polygon),
                }
            )
    if not features:
        return {
            "type": "FeatureCollection",
            "features": [],
            "properties": {"observer": {"lat": observer.lat, "lon": observer.lon}},
        }
    merged = unary_union([shape(feature["geometry"]) for feature in features])
    return {
        "type": "FeatureCollection",
        "features": [
            {
                "type": "Feature",
                "properties": {"visible": True},
                "geometry": mapping(merged),
            }
        ],
        "properties": {"observer": {"lat": observer.lat, "lon": observer.lon}},
    }


def compute_viewshed(
    request: ViewshedRequest,
    *,
    dem_path: str | Path,
    output_dir: str | Path,
    gdal_viewshed_bin: str = "gdal_viewshed",
) -> RasterResult:
    dem_root = Path(dem_path)
    output_root = Path(output_dir)
    output_root.mkdir(parents=True, exist_ok=True)

    observer = GeoPoint(lat=request.observer.lat, lon=request.observer.lon)
    source_tile = _find_dem_tile(observer.lat, observer.lon, dem_root)
    calc_radius_m = request.calc_radius_m or request.radius_m
    bbox_wgs84 = _radius_bbox(observer, calc_radius_m)
    source_tiles = _find_dem_tiles(dem_root, bbox_wgs84)
    if source_tile is not None and source_tile not in source_tiles:
        source_tiles.append(source_tile)
    if not source_tiles:
        raise FileNotFoundError(
            f"No DEM tile covers viewshed radius at lat={observer.lat}, lon={observer.lon}"
        )

    with tempfile.TemporaryDirectory(prefix="viewshed-") as tmp_dir:
        tmp = Path(tmp_dir)
        metric_dem = tmp / "metric_dem.tif"
        viewshed_raster = tmp / "viewshed.tif"
        metric_path, observer_x, observer_y, dem_tiles_count = _prepare_metric_dem(
            source_tiles,
            observer,
            calc_radius_m,
            metric_dem,
        )
        _run_gdal_viewshed(
            metric_path,
            observer_x,
            observer_y,
            request.observer.height_agl,
            request.target_height_agl,
            calc_radius_m,
            viewshed_raster,
            gdal_viewshed_bin,
        )

        if request.format == "geojson":
            result = _viewshed_geojson(viewshed_raster, observer)
            output_file = output_root / f"viewshed_{observer.lat:.5f}_{observer.lon:.5f}.json"
            output_file.write_text(
                __import__("json").dumps(result),
                encoding="utf-8",
            )
            return RasterResult(
                uri=str(output_file),
                metadata={
                    "format": "geojson",
                    "surface": request.surface,
                    "radius_m": request.radius_m,
                    "calc_radius_m": calc_radius_m,
                    "dem_tiles_count": dem_tiles_count,
                    "feature_count": len(result.get("features", [])),
                },
            )

        output_file = output_root / f"viewshed_{observer.lat:.5f}_{observer.lon:.5f}.tif"
        output_file.write_bytes(viewshed_raster.read_bytes())
        return RasterResult(
            uri=str(output_file),
            metadata={
                "format": request.format,
                "surface": request.surface,
                "radius_m": request.radius_m,
                "calc_radius_m": calc_radius_m,
                "dem_tiles_count": dem_tiles_count,
            },
        )