breakpilot-pwa/geo-service/services/tile_server.py

"""
Tile Server Service
Serves vector tiles from PMTiles format or generates on-demand from PostGIS
"""
import os
import gzip
from typing import Optional
from datetime import datetime
import structlog
from pmtiles.reader import Reader as PMTilesReader
from pmtiles.tile import TileType

from config import settings

logger = structlog.get_logger(__name__)


class MMapFileReader:
    """Memory-mapped file reader for PMTiles."""

    def __init__(self, path: str):
        self.path = path
        self._file = None
        self._size = 0

    def __enter__(self):
        self._file = open(self.path, "rb")
        self._file.seek(0, 2)  # Seek to end
        self._size = self._file.tell()
        self._file.seek(0)
        return self

    def __exit__(self, *args):
        if self._file:
            self._file.close()

    def read(self, offset: int, length: int) -> bytes:
        """Read bytes from file at offset."""
        self._file.seek(offset)
        return self._file.read(length)

    def size(self) -> int:
        """Get file size."""
        return self._size


class TileServerService:
    """
    Service for serving vector tiles from PMTiles format.

    PMTiles is a cloud-optimized format for tile archives that allows
    random access to individual tiles without extracting the entire archive.
    """

    def __init__(self):
        self.pmtiles_path = settings.pmtiles_path
        self.cache_dir = settings.tile_cache_dir
        self._reader = None
        self._metadata_cache = None

    def _get_reader(self) -> Optional[PMTilesReader]:
        """Get or create PMTiles reader."""
        if not os.path.exists(self.pmtiles_path):
            logger.warning("PMTiles file not found", path=self.pmtiles_path)
            return None

        if self._reader is None:
            try:
                file_reader = MMapFileReader(self.pmtiles_path)
                file_reader.__enter__()
                self._reader = PMTilesReader(file_reader)
                logger.info("PMTiles reader initialized", path=self.pmtiles_path)
            except Exception as e:
                logger.error("Failed to initialize PMTiles reader", error=str(e))
                return None

        return self._reader

    async def get_tile(self, z: int, x: int, y: int) -> Optional[bytes]:
        """
        Get a vector tile at the specified coordinates.

        Args:
            z: Zoom level
            x: Tile X coordinate
            y: Tile Y coordinate

        Returns:
            Tile data as gzipped protobuf, or None if not found
        """
        # Check cache first
        cache_path = os.path.join(self.cache_dir, str(z), str(x), f"{y}.pbf")
        if os.path.exists(cache_path):
            with open(cache_path, "rb") as f:
                return f.read()

        # Try to get from PMTiles
        reader = self._get_reader()
        if reader is None:
            raise FileNotFoundError("PMTiles file not available")

        try:
            tile_data = reader.get_tile(z, x, y)

            if tile_data is None:
                return None

            # Cache the tile
            await self._cache_tile(z, x, y, tile_data)

            return tile_data

        except Exception as e:
            logger.error("Error reading tile", z=z, x=x, y=y, error=str(e))
            return None

    async def _cache_tile(self, z: int, x: int, y: int, data: bytes):
        """Cache a tile to disk."""
        cache_path = os.path.join(self.cache_dir, str(z), str(x))
        os.makedirs(cache_path, exist_ok=True)

        tile_path = os.path.join(cache_path, f"{y}.pbf")
        with open(tile_path, "wb") as f:
            f.write(data)

    async def get_metadata(self) -> dict:
        """
        Get metadata about the tile archive.

        Returns:
            Dictionary with metadata including bounds, zoom levels, etc.
        """
        if self._metadata_cache is not None:
            return self._metadata_cache

        reader = self._get_reader()
        if reader is None:
            return {
                "data_available": False,
                "minzoom": 0,
                "maxzoom": 14,
                "bounds": [5.87, 47.27, 15.04, 55.06],
                "center": [10.45, 51.16, 6],
            }

        try:
            header = reader.header()
            metadata = reader.metadata()

            self._metadata_cache = {
                "data_available": True,
                "minzoom": header.get("minZoom", 0),
                "maxzoom": header.get("maxZoom", 14),
                "bounds": header.get("bounds", [5.87, 47.27, 15.04, 55.06]),
                "center": header.get("center", [10.45, 51.16, 6]),
                "tile_type": "mvt",  # Mapbox Vector Tiles
                "last_updated": datetime.fromtimestamp(
                    os.path.getmtime(self.pmtiles_path)
                ).isoformat() if os.path.exists(self.pmtiles_path) else None,
                **metadata,
            }

            return self._metadata_cache

        except Exception as e:
            logger.error("Error reading metadata", error=str(e))
            return {"data_available": False}

    def clear_cache(self):
        """Clear the tile cache."""
        import shutil

        if os.path.exists(self.cache_dir):
            shutil.rmtree(self.cache_dir)
            os.makedirs(self.cache_dir)
            logger.info("Tile cache cleared")

    def get_cache_size_mb(self) -> float:
        """Get the current cache size in MB."""
        total_size = 0
        for dirpath, dirnames, filenames in os.walk(self.cache_dir):
            for filename in filenames:
                filepath = os.path.join(dirpath, filename)
                total_size += os.path.getsize(filepath)

        return total_size / (1024 * 1024)