from __future__ import annotations

import math
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import rasterio
from cachetools import LRUCache
from rasterio.io import DatasetReader
from rasterio.transform import rowcol


@dataclass(frozen=True)
class TileInfo:
    path: Path
    bounds: tuple[float, float, float, float]  # left, bottom, right, top
    nodata: float | None


class RasterService:
  """Base GeoTIFF reader with tile index and LRU cache."""

  def __init__(self, data_dir: str | Path, cache_size: int = 16) -> None:
    self.data_dir = Path(data_dir)
    self._tiles: list[TileInfo] = []
    self._cache: LRUCache[str, DatasetReader] = LRUCache(maxsize=cache_size)
    self._scan_tiles()

  def _scan_tiles(self) -> None:
    if not self.data_dir.exists():
      return

    for path in sorted(self.data_dir.glob("*.tif")) + sorted(self.data_dir.glob("*.tiff")):
      try:
        with rasterio.open(path) as dataset:
          bounds = dataset.bounds
          nodata = dataset.nodata
      except rasterio.errors.RasterioIOError:
        continue

      self._tiles.append(
        TileInfo(
          path=path,
          bounds=(bounds.left, bounds.bottom, bounds.right, bounds.top),
          nodata=nodata,
        )
      )

  def is_available(self) -> bool:
    return len(self._tiles) > 0

  def _find_tile(self, lng: float, lat: float) -> TileInfo | None:
    for tile in self._tiles:
      left, bottom, right, top = tile.bounds
      if left <= lng <= right and bottom <= lat <= top:
        return tile
    return None

  def _open_dataset(self, tile: TileInfo) -> DatasetReader:
    key = str(tile.path)
    if key not in self._cache:
      self._cache[key] = rasterio.open(tile.path)
    return self._cache[key]

  def close(self) -> None:
    for dataset in self._cache.values():
      dataset.close()
    self._cache.clear()

  def _normalize_value(self, value: float, nodata: float | None, default: float) -> float:
    if not math.isfinite(value):
      return default
    if nodata is not None and value == nodata:
      return default
    return float(value)

  def get_value(self, lat: float, lng: float, default: float = 0.0) -> float:
    tile = self._find_tile(lng, lat)
    if tile is None:
      return default

    dataset = self._open_dataset(tile)
    row, col = rowcol(dataset.transform, lng, lat)
    if row < 0 or col < 0 or row >= dataset.height or col >= dataset.width:
      return default

    value = float(dataset.read(1, window=((row, row + 1), (col, col + 1)))[0, 0])
    return self._normalize_value(value, tile.nodata, default)

  def get_values(self, coordinates: Iterable[tuple[float, float]], default: float = 0.0) -> list[float]:
    coords = list(coordinates)
    if not coords:
      return []

    results = [default] * len(coords)
    by_tile: dict[str, list[tuple[int, float, float]]] = {}

    for index, (lat, lng) in enumerate(coords):
      tile = self._find_tile(lng, lat)
      if tile is None:
        continue
      by_tile.setdefault(str(tile.path), []).append((index, lat, lng))

    for tile_path, items in by_tile.items():
      tile = next(t for t in self._tiles if str(t.path) == tile_path)
      dataset = self._open_dataset(tile)
      lats = [lat for _, lat, _ in items]
      lngs = [lng for _, _, lng in items]
      rows, cols = rowcol(dataset.transform, lngs, lats)

      for (index, lat, lng), row, col in zip(items, rows, cols, strict=True):
        if row < 0 or col < 0 or row >= dataset.height or col >= dataset.width:
          continue
        value = float(dataset.read(1, window=((row, row + 1), (col, col + 1)))[0, 0])
        results[index] = self._normalize_value(value, tile.nodata, default)

    return results

  def get_value_at_distance(self, lat: float, lng: float) -> float:
    return self.get_value(lat, lng, default=0.0)