from __future__ import annotations

"""Local DEM inspection and sampling for the desktop calculation client."""

from dataclasses import dataclass
import hashlib
import importlib
import json
from math import hypot
from pathlib import Path
import sys
from typing import Any, Iterable
from urllib import request
import zipfile

# The desktop launcher executes the unpacked client without the developer's
# Python environment. Keep the raster reader with the feature package.
_VENDOR_DIR = Path(__file__).resolve().parent / "vendor"
if _VENDOR_DIR.is_dir() and str(_VENDOR_DIR) not in sys.path:
    sys.path.insert(0, str(_VENDOR_DIR))

try:
    import rasterio as _rasterio_lib  # type: ignore
except ImportError:
    _rasterio_lib = None

_APP_ROOT = Path(__file__).resolve().parent.parent
_RUNTIME_STATE_DIR = _APP_ROOT / ".dem_runtime"


SUPPORTED_SUFFIXES = {".tif", ".tiff", ".asc", ".grd"}


class DemImportError(ValueError):
    pass


@dataclass(frozen=True)
class DemMetadata:
    path: str
    format: str
    width: int
    height: int
    crs: str
    x_resolution: float
    y_resolution: float
    bounds: tuple[float, float, float, float]
    nodata: float | None
    sample_min: float | None
    sample_max: float | None

    def as_dict(self) -> dict[str, Any]:
        left, bottom, right, top = self.bounds
        return {
            "ok": True,
            "path": self.path,
            "format": self.format,
            "width": self.width,
            "height": self.height,
            "crs": self.crs or "未写入坐标系",
            "resolution": {"x": self.x_resolution, "y": self.y_resolution},
            "bounds": {"left": left, "bottom": bottom, "right": right, "top": top},
            "nodata": self.nodata,
            "sample_min": self.sample_min,
            "sample_max": self.sample_max,
            "message": "DEM 已读取。高程统计为抽样结果；工程使用前请核对坐标系、高程基准和测绘成果。",
        }


def _path(value: str | Path) -> Path:
    path = Path(str(value or "")).expanduser()
    if not path.is_file():
        raise DemImportError("DEM 文件不存在，请重新选择文件。")
    if path.suffix.lower() not in SUPPORTED_SUFFIXES:
        raise DemImportError("仅支持 GeoTIFF（.tif/.tiff）和 ESRI ASCII Grid（.asc/.grd）。")
    return path


def _float(value: Any, name: str) -> float:
    try:
        return float(value)
    except (TypeError, ValueError) as exc:
        raise DemImportError(f"{name}必须是数值。") from exc


def _asc_header(path: Path) -> tuple[dict[str, float], list[list[float]]]:
    lines = path.read_text(encoding="utf-8-sig", errors="replace").splitlines()
    header: dict[str, float] = {}
    start = 0
    for index, line in enumerate(lines[:12]):
        parts = line.split()
        if len(parts) < 2:
            continue
        key = parts[0].lower()
        if key in {"ncols", "nrows", "xllcorner", "yllcorner", "xllcenter", "yllcenter", "cellsize", "nodata_value"}:
            header[key] = _float(parts[1], key)
            start = index + 1
    for required in ("ncols", "nrows", "cellsize"):
        if required not in header:
            raise DemImportError("ASCII DEM 缺少 ncols、nrows 或 cellsize 头信息。")
    if not ({"xllcorner", "xllcenter"} & set(header)) or not ({"yllcorner", "yllcenter"} & set(header)):
        raise DemImportError("ASCII DEM 缺少左下角坐标信息。")
    rows: list[list[float]] = []
    for line in lines[start:]:
        if line.strip():
            rows.append([_float(item, "DEM 高程") for item in line.split()])
    ncols, nrows = int(header["ncols"]), int(header["nrows"])
    if len(rows) != nrows or any(len(row) != ncols for row in rows):
        raise DemImportError("ASCII DEM 的栅格行列数与头信息不一致。")
    return header, rows


def _asc_metadata(path: Path) -> DemMetadata:
    header, rows = _asc_header(path)
    cell = header["cellsize"]
    width, height = int(header["ncols"]), int(header["nrows"])
    left = header.get("xllcorner", header.get("xllcenter", 0.0) - cell / 2)
    bottom = header.get("yllcorner", header.get("yllcenter", 0.0) - cell / 2)
    nodata = header.get("nodata_value")
    values = [value for row in rows for value in row if nodata is None or value != nodata]
    return DemMetadata(str(path), "ESRI ASCII Grid", width, height, "未写入坐标系", cell, cell,
                       (left, bottom, left + width * cell, bottom + height * cell), nodata,
                       min(values) if values else None, max(values) if values else None)


def _rasterio() -> Any:
    if _rasterio_lib is None:
        raise DemImportError("当前软件缺少 GeoTIFF 解析组件，请安装含 DEM 功能的最新版。")
    return _rasterio_lib


def runtime_status() -> dict[str, Any]:
    return {
        "ok": True,
        "installed": _rasterio_lib is not None,
        "version": str(getattr(_rasterio_lib, "__version__", "")),
        "message": "DEM GeoTIFF 组件已就绪。" if _rasterio_lib is not None else "首次使用 GeoTIFF 时需要下载 DEM 栅格组件。",
    }


def _runtime_config() -> dict[str, str]:
    try:
        data = json.loads((_APP_ROOT / "version.json").read_text(encoding="utf-8-sig"))
    except (OSError, json.JSONDecodeError) as exc:
        raise DemImportError("无法读取 DEM 组件配置。") from exc
    config = data.get("dem_runtime") if isinstance(data, dict) else None
    if not isinstance(config, dict):
        raise DemImportError("当前版本未配置 DEM 栅格组件下载地址。")
    url, digest = str(config.get("url") or ""), str(config.get("sha256") or "").lower()
    if not url.startswith("https://") or len(digest) != 64:
        raise DemImportError("DEM 栅格组件配置不完整。")
    return {"url": url, "sha256": digest}


def _safe_extract(archive: zipfile.ZipFile, target: Path) -> None:
    root = target.resolve()
    for member in archive.infolist():
        destination = (target / member.filename).resolve()
        if destination != root and root not in destination.parents:
            raise DemImportError("DEM 组件包含不安全的文件路径。")
    archive.extractall(target)


def install_runtime() -> dict[str, Any]:
    global _rasterio_lib
    if _rasterio_lib is not None:
        return runtime_status()
    config = _runtime_config()
    _RUNTIME_STATE_DIR.mkdir(parents=True, exist_ok=True)
    archive_path = _RUNTIME_STATE_DIR / "dem_runtime.zip.part"
    try:
        with request.urlopen(config["url"], timeout=180) as response, archive_path.open("wb") as target:
            while chunk := response.read(1024 * 1024):
                target.write(chunk)
        digest = hashlib.sha256(archive_path.read_bytes()).hexdigest()
        if digest != config["sha256"]:
            raise DemImportError("DEM 组件完整性校验失败，请重新下载。")
        with zipfile.ZipFile(archive_path) as archive:
            _safe_extract(archive, _APP_ROOT)
    except (OSError, zipfile.BadZipFile) as exc:
        raise DemImportError(f"下载或安装 DEM 组件失败：{exc}") from exc
    finally:
        archive_path.unlink(missing_ok=True)
    importlib.invalidate_caches()
    if str(_VENDOR_DIR) not in sys.path:
        sys.path.insert(0, str(_VENDOR_DIR))
    try:
        _rasterio_lib = importlib.import_module("rasterio")
    except ImportError as exc:
        raise DemImportError("DEM 组件安装后仍无法加载，请重新安装软件。") from exc
    return runtime_status()


def _tiff_metadata(path: Path) -> DemMetadata:
    rasterio = _rasterio()
    with rasterio.open(path) as dataset:
        # Sample at most 40,000 cells so very large DEMs remain responsive.
        band = dataset.read(1, masked=True, out_shape=(min(dataset.height, 200), min(dataset.width, 200)))
        values = band.compressed()
        bounds = dataset.bounds
        return DemMetadata(str(path), "GeoTIFF", dataset.width, dataset.height, str(dataset.crs or ""),
                           abs(dataset.transform.a), abs(dataset.transform.e),
                           (bounds.left, bounds.bottom, bounds.right, bounds.top), dataset.nodata,
                           float(values.min()) if values.size else None, float(values.max()) if values.size else None)


def read_metadata(path: str | Path) -> dict[str, Any]:
    source = _path(path)
    return (_asc_metadata(source) if source.suffix.lower() in {".asc", ".grd"} else _tiff_metadata(source)).as_dict()


def _asc_sample(path: Path, x: float, y: float) -> float | None:
    header, rows = _asc_header(path)
    cell = header["cellsize"]
    left = header.get("xllcorner", header.get("xllcenter", 0.0) - cell / 2)
    bottom = header.get("yllcorner", header.get("yllcenter", 0.0) - cell / 2)
    col, row = int((x - left) // cell), int((bottom + int(header["nrows"]) * cell - y) // cell)
    if row < 0 or col < 0 or row >= len(rows) or col >= len(rows[0]):
        return None
    value = rows[row][col]
    return None if header.get("nodata_value") == value else value


def _tiff_samples(path: Path, points: list[tuple[float, float]], input_crs: str = "") -> list[float | None]:
    rasterio = _rasterio()
    with rasterio.open(path) as dataset:
        coordinates = points
        if input_crs and dataset.crs and input_crs != str(dataset.crs):
            from rasterio.warp import transform
            xs, ys = transform(input_crs, dataset.crs, [item[0] for item in points], [item[1] for item in points])
            coordinates = list(zip(xs, ys))
        samples = dataset.sample(coordinates, indexes=1, masked=True)
        result: list[float | None] = []
        for value in samples:
            result.append(None if bool(getattr(value, "mask", False)) else float(value[0]))
        return result


def sample_points(path: str | Path, points: Iterable[dict[str, Any]], input_crs: str = "") -> dict[str, Any]:
    source = _path(path)
    normalized = [(_float(item.get("x"), "X 坐标"), _float(item.get("y"), "Y 坐标")) for item in points]
    if not normalized:
        raise DemImportError("请至少提供一个采样点。")
    if len(normalized) > 1000:
        raise DemImportError("单次最多采样 1000 个点。")
    values = ([_asc_sample(source, x, y) for x, y in normalized] if source.suffix.lower() in {".asc", ".grd"}
              else _tiff_samples(source, normalized, input_crs))
    rows = [{"index": index + 1, "x": x, "y": y, "elevation_m": value, "status": "ok" if value is not None else "outside_or_nodata"}
            for index, ((x, y), value) in enumerate(zip(normalized, values))]
    return {"ok": True, "points": rows, "message": "已完成 DEM 高程采样；空值表示点位超出栅格范围或落在 NoData 区。"}


def sample_profile(path: str | Path, start_x: Any, start_y: Any, end_x: Any, end_y: Any, count: Any = 51, input_crs: str = "") -> dict[str, Any]:
    sx, sy, ex, ey = (_float(start_x, "起点 X"), _float(start_y, "起点 Y"), _float(end_x, "终点 X"), _float(end_y, "终点 Y"))
    try:
        count_value = int(count)
    except (TypeError, ValueError) as exc:
        raise DemImportError("采样点数必须是整数。") from exc
    if not 2 <= count_value <= 1000:
        raise DemImportError("采样点数必须在 2 至 1000 之间。")
    distance = hypot(ex - sx, ey - sy)
    if distance == 0:
        raise DemImportError("起点与终点不能相同。")
    points = [{"x": sx + (ex - sx) * index / (count_value - 1), "y": sy + (ey - sy) * index / (count_value - 1)} for index in range(count_value)]
    result = sample_points(path, points, input_crs)
    rows = result["points"]
    valid = [row for row in rows if row["elevation_m"] is not None]
    for index, row in enumerate(rows):
        row["distance"] = distance * index / (count_value - 1)
    slopes = [(valid[index + 1]["elevation_m"] - valid[index]["elevation_m"]) / (valid[index + 1]["distance"] - valid[index]["distance"])
              for index in range(len(valid) - 1) if valid[index + 1]["distance"] > valid[index]["distance"]]
    result.update({
        "distance": distance,
        "profile": rows,
        "summary": {
            "valid_count": len(valid),
            "min_elevation_m": min((row["elevation_m"] for row in valid), default=None),
            "max_elevation_m": max((row["elevation_m"] for row in valid), default=None),
            "end_to_end_slope_percent": ((valid[-1]["elevation_m"] - valid[0]["elevation_m"]) / distance * 100) if len(valid) >= 2 else None,
            "max_segment_slope_percent": max((abs(value) * 100 for value in slopes), default=None),
        },
        "message": "已提取 DEM 纵断面。距离单位与 DEM 坐标单位一致；工程使用前必须确认其为米及高程基准一致。",
    })
    return result
