"""Geospatial tools: load and query geographic datasets.

These three tools handle the spatial data workflow separately from the tabular
workflow in retrieval.py and querying.py, for structural reasons.

Why separate from the main retrieval/query tools?

- ``load_geodata`` — downloads SHP, KML, and GeoJSON resources and stores them
  with spatial column handling (geometry_wkt, spatial index). Kept separate from
  ``download_resource`` because: (1) it requires geopandas which may not be
  installed, (2) it uses a different table prefix (geo_ vs ds_) so spatial tables
  are distinguishable in cache listings, (3) it has different format handling
  (ZIP extraction, coordinate validation). Merging into download_resource would
  make that tool conditionally depend on geopandas and add spatial branching to
  a tool that non-spatial users never need.
- ``spatial_query`` — runs point-in-polygon, bounding-box, and radius queries
  against geo-cached tables using DuckDB spatial extension. Kept separate from
  ``query_cached`` because: (1) it constructs the spatial SQL for the model so
  the model doesn't need to know DuckDB spatial syntax; (2) it validates that
  the spatial extension is available and gives a clear error if not; (3) it
  accepts geographic parameters (lat/lng, radius_km, bbox) rather than raw SQL.
- ``list_geo_datasets`` — searches specifically for datasets with spatial resource
  formats (SHP, KML, GeoJSON, WMS, WFS). Kept separate from ``search_datasets``
  because spatial format filtering requires awareness of which format strings
  indicate geographic data — the model would need to know to pass
  resource_format="SHP" OR "KML" OR "GEOJSON" etc. to search_datasets, and
  would still miss WMS/WFS service endpoints.

Consolidation considered and rejected: merging load_geodata into spatial_query
(auto-load on first query) would hide the geopandas import error until query
time, and would lose the explicit "did this load correctly?" confirmation step
that is useful for large shapefiles that may fail mid-parse.
"""
from __future__ import annotations

import asyncio
import io
from fastmcp import Context

from ontario_data.server import READONLY, mcp
from ontario_data.formatting import format_records, md_response
from ontario_data.utils import (
    SpatialExtensionError,
    get_lifespan_state,
    arcgis_guard,
    fan_out,
    get_cache,
    get_deps,
    is_arcgis_portal,
    make_geo_table_name,
    parse_portal_id,
    require_cached,
    resolve_resource_portal,
)


@mcp.tool(annotations=READONLY)
async def load_geodata(
    resource_id: str,
    force_refresh: bool = False,
    ctx: Context = None,
) -> str:
    """Download and cache a geospatial resource (SHP, KML, GeoJSON) into DuckDB with spatial support.

    Args:
        resource_id: Prefixed resource ID (e.g. "toronto:abc123") or bare ID
        force_refresh: Re-download even if cached
    """
    import geopandas as gpd
    import pandas as pd

    configs = get_lifespan_state(ctx)["portal_configs"]
    portal, bare_id = parse_portal_id(resource_id, set(configs.keys()))

    if portal and is_arcgis_portal(ctx, portal):
        return arcgis_guard(
            resource_id,
            alternative=f"download_resource(resource_id='{resource_id}') for CSV data, or access the portal directly",
        )

    cache = get_cache(ctx)

    if cache.is_cached(bare_id) and not force_refresh:
        table_name = cache.get_table_name(bare_id)
        return md_response(status="already_cached", table_name=table_name)

    if not portal:
        portal, bare_id = await resolve_resource_portal(ctx, resource_id)

    ckan, _ = get_deps(ctx, portal)
    resource = await ckan.resource_show(bare_id)
    dataset_id = resource.get("package_id", "")
    dataset = await ckan.package_show(dataset_id) if dataset_id else {}
    fmt = (resource.get("format") or "").upper()
    url = resource.get("url", "")

    if not url:
        raise ValueError(f"Resource '{bare_id}' has no download URL")

    await ctx.report_progress(0, 100, "Downloading geospatial data...")

    http_client = get_lifespan_state(ctx)["http_client"]
    response = await http_client.get(url, timeout=120.0, follow_redirects=True)
    response.raise_for_status()
    content = response.content

    await ctx.report_progress(50, 100, "Parsing geospatial data...")

    if fmt == "GEOJSON":
        gdf = gpd.read_file(io.BytesIO(content), driver="GeoJSON")
    elif fmt == "KML":
        gdf = gpd.read_file(io.BytesIO(content), driver="KML")
    elif fmt == "ZIP" or (fmt == "SHP" and content[:4] == b"PK\x03\x04"):
        import os
        import tempfile
        import zipfile
        with tempfile.TemporaryDirectory() as tmpdir:
            with zipfile.ZipFile(io.BytesIO(content)) as zf:
                for member in zf.namelist():
                    dest = os.path.realpath(os.path.join(tmpdir, member))
                    if not dest.startswith(os.path.realpath(tmpdir) + os.sep) and dest != os.path.realpath(tmpdir):
                        raise ValueError(f"Zip entry '{member}' would extract outside target directory")
                zf.extractall(tmpdir)
            gdf = gpd.read_file(tmpdir)
    else:
        raise ValueError(f"Unsupported geospatial format: {fmt}")

    await ctx.report_progress(80, 100, "Storing in DuckDB...")

    # Convert geometry to WKT for DuckDB storage
    df = gdf.copy()
    if "geometry" in df.columns:
        df["geometry_wkt"] = df["geometry"].apply(lambda g: g.wkt if g else None)
        df["geometry_type"] = df["geometry"].apply(lambda g: g.geom_type if g else None)
        if hasattr(gdf, "crs") and gdf.crs:
            df["crs"] = str(gdf.crs)
        bounds = gdf.total_bounds  # [minx, miny, maxx, maxy]
        df = df.drop(columns=["geometry"])
    else:
        bounds = None

    table_name = make_geo_table_name(dataset.get("name", ""), bare_id, portal=portal)

    await asyncio.to_thread(
        cache.store_resource,
        resource_id=bare_id,
        dataset_id=dataset_id,
        table_name=table_name,
        df=pd.DataFrame(df),
        source_url=url,
    )

    await ctx.report_progress(100, 100, "Done")

    return md_response(
        status="loaded",
        table_name=table_name,
        row_count=len(df),
        columns=list(df.columns),
        geometry_types=df["geometry_type"].unique().tolist() if "geometry_type" in df.columns else [],
        bounds={"minx": bounds[0], "miny": bounds[1], "maxx": bounds[2], "maxy": bounds[3]} if bounds is not None else None,
        crs=str(gdf.crs) if hasattr(gdf, "crs") and gdf.crs else None,
        hint=f'Query with: SELECT * FROM "{table_name}" LIMIT 10',
    )


@mcp.tool(annotations=READONLY)
async def spatial_query(
    resource_id: str,
    operation: str,
    latitude: float | None = None,
    longitude: float | None = None,
    radius_km: float | None = None,
    bbox: list[float] | None = None,
    limit: int = 100,
    ctx: Context = None,
) -> str:
    """Run spatial queries against cached geospatial data.

    Args:
        resource_id: Resource ID (must be cached via load_geodata)
        operation: "contains_point", "within_bbox", or "within_radius"
        latitude: Latitude for point queries (-90 to 90)
        longitude: Longitude for point queries (-180 to 180)
        radius_km: Radius in kilometers (for within_radius, must be > 0)
        bbox: Bounding box as [min_lng, min_lat, max_lng, max_lat] (for within_bbox)
        limit: Max results
    """
    # Clamp limit to a safe range — negative values (e.g. -1) would bypass
    # the intended cap in DuckDB, and very large values could exhaust memory.
    limit = max(1, min(limit, 10_000))

    cache = get_cache(ctx)
    table_name = require_cached(cache, resource_id)

    if not cache.has_spatial_extension:
        raise SpatialExtensionError(
            "DuckDB spatial extension is not available. "
            "Spatial queries require the spatial extension to be installed."
        )

    # Validate coordinates
    if latitude is not None and not (-90 <= latitude <= 90):
        raise ValueError(f"Latitude {latitude} out of range. Must be between -90 and 90.")
    if longitude is not None and not (-180 <= longitude <= 180):
        raise ValueError(f"Longitude {longitude} out of range. Must be between -180 and 180.")
    if radius_km is not None and radius_km <= 0:
        raise ValueError(f"Radius must be positive, got {radius_km}.")
    if bbox is not None:
        if len(bbox) != 4:
            raise ValueError(f"Bounding box must have 4 values [min_lng, min_lat, max_lng, max_lat], got {len(bbox)}.")
        if not (-180 <= bbox[0] <= 180 and -180 <= bbox[2] <= 180):
            raise ValueError(f"Bounding box longitudes out of range (-180 to 180).")
        if not (-90 <= bbox[1] <= 90 and -90 <= bbox[3] <= 90):
            raise ValueError(f"Bounding box latitudes out of range (-90 to 90).")

    if operation == "contains_point" and latitude is not None and longitude is not None:
        sql = f"""
            SELECT *, ST_Distance(
                ST_GeomFromText(geometry_wkt),
                ST_Point(?, ?)
            ) as distance
            FROM "{table_name}"
            WHERE geometry_wkt IS NOT NULL
            AND ST_Contains(ST_GeomFromText(geometry_wkt), ST_Point(?, ?))
            LIMIT {limit}
        """
        params = [longitude, latitude, longitude, latitude]
    elif operation == "within_radius" and latitude is not None and longitude is not None and radius_km is not None:
        import math
        # Haversine-based radius query (correct at all latitudes).
        # Use a degree bounding box as a cheap pre-filter, with cosine
        # correction so the longitude band isn't ~40% too narrow at Ontario latitudes.
        lat_delta = radius_km / 111.0
        cos_lat = math.cos(math.radians(latitude))
        lon_delta = radius_km / (111.0 * cos_lat) if cos_lat > 0.01 else 360.0
        sql = f"""
            WITH centroid AS (
                SELECT ST_Centroid(ST_GeomFromText(geometry_wkt)) AS c, *
                FROM "{table_name}"
                WHERE geometry_wkt IS NOT NULL
                  AND ST_Y(ST_Centroid(ST_GeomFromText(geometry_wkt)))
                      BETWEEN ? - ? AND ? + ?
                  AND ST_X(ST_Centroid(ST_GeomFromText(geometry_wkt)))
                      BETWEEN ? - ? AND ? + ?
            )
            SELECT * EXCLUDE (c),
                   2 * 6371 * ASIN(SQRT(
                       POWER(SIN(RADIANS(ST_Y(c) - ?) / 2), 2) +
                       COS(RADIANS(?)) * COS(RADIANS(ST_Y(c))) *
                       POWER(SIN(RADIANS(ST_X(c) - ?) / 2), 2)
                   )) AS distance_km
            FROM centroid
            WHERE distance_km <= ?
            ORDER BY distance_km
            LIMIT {limit}
        """
        params = [
            latitude, lat_delta, latitude, lat_delta,   # lat bbox
            longitude, lon_delta, longitude, lon_delta,  # lon bbox
            latitude, latitude, longitude,               # haversine
            radius_km,                                   # final filter
        ]
    elif operation == "within_bbox" and bbox and len(bbox) == 4:
        sql = f"""
            SELECT *
            FROM "{table_name}"
            WHERE geometry_wkt IS NOT NULL
            AND ST_Intersects(
                ST_GeomFromText(geometry_wkt),
                ST_MakeEnvelope(?, ?, ?, ?)
            )
            LIMIT {limit}
        """
        params = list(bbox)
    else:
        raise ValueError(
            f"Invalid operation '{operation}' or missing parameters. "
            f"Valid: contains_point (lat, lng), within_radius (lat, lng, radius_km), within_bbox (bbox)"
        )

    records = cache.execute_sql_dict(sql, params=params)

    if not records:
        # Provide context when no results found
        parts = [format_records(records, row_count=0)]
        try:
            total = cache.execute_sql(f'SELECT COUNT(*) FROM "{table_name}" WHERE geometry_wkt IS NOT NULL')[0][0]
            if total > 0:
                context = f"No features found"
                if operation == "within_radius":
                    context += f" within {radius_km}km of ({latitude}, {longitude})"
                elif operation == "contains_point":
                    context += f" containing point ({latitude}, {longitude})"
                context += f". The table has {total:,} features with geometry."
                parts.append(context)
        except Exception:
            pass
        return "\n".join(parts)

    return format_records(records, row_count=len(records))


@mcp.tool(annotations=READONLY)
async def list_geo_datasets(
    format_filter: str | None = None,
    limit: int = 50,
    portal: str | None = None,
    ctx: Context = None,
) -> str:
    """Find all datasets that contain geospatial resources (SHP, KML, GeoJSON).

    Args:
        format_filter: Filter to specific format: "SHP", "KML", "GEOJSON", or None for all
        limit: Max results per portal
        portal: Narrow to one portal. Default: all portals.
    """
    geo_formats = [format_filter.upper()] if format_filter else ["SHP", "KML", "GEOJSON"]

    async def _list_geo(portal_key: str) -> list[dict]:
        ckan, _ = get_deps(ctx, portal_key)
        datasets = []
        seen_ids = set()
        for fmt in geo_formats:
            result = await ckan.package_search(filters={"res_format": fmt}, rows=min(limit, 50))
            for ds in result["results"]:
                if ds["id"] not in seen_ids:
                    seen_ids.add(ds["id"])
                    geo_resources = [
                        {"id": r["id"], "name": r.get("name"), "format": r.get("format"), "size": r.get("size")}
                        for r in ds.get("resources", [])
                        if (r.get("format") or "").upper() in ("SHP", "KML", "GEOJSON", "ZIP")
                    ]
                    datasets.append({
                        "id": f"{portal_key}:{ds['id']}",
                        "title": ds.get("title"),
                        "organization": ds.get("organization", {}).get("title"),
                        "geo_resources": geo_resources,
                    })
        return datasets

    raw = await fan_out(ctx, portal, _list_geo)
    all_datasets = []
    for _, result, error in raw:
        if result and not error:
            all_datasets.extend(result)

    return md_response(total=len(all_datasets), datasets=all_datasets[:limit])
