"""Retrieval and cache management tools.

These four tools own the local DuckDB cache lifecycle: download → inspect →
manage → refresh. They form the bridge between remote portal data and the
fast local SQL queries in querying.py.

Why a two-step download-then-query workflow instead of auto-caching?

An auto-caching query tool was considered (one call that downloads silently
if needed). Rejected because:
1. Downloads range from 2 seconds to 2 minutes and 50–500MB. The model
   should control when to pay that cost, not have it triggered implicitly.
2. SQL must reference the table name, which is only known after downloading.
   An auto-cache tool can't eliminate this — it just hides it, making the
   flow harder to debug when downloads fail or timeout.
3. Explicit download means the model sees row count, column list, and dtypes
   before writing SQL — reducing query errors from type assumptions.

The hint in download_resource always includes the ready-to-use SELECT
statement so the model can query immediately without a cache_info call.

Why separate cache tools instead of one "cache" tool with a mode parameter?

- ``download_resource`` — READONLY (annotated as such per MCP spec). Populates
  the cache but makes no remote mutations. Cannot be merged with destructive ops.
- ``cache_info`` — zero-argument listing of what's cached. Different access
  pattern from ``cache_manage`` (which requires action + optional resource_id).
- ``cache_manage`` — DESTRUCTIVE (annotated). Kept separate from READONLY tools
  so MCP clients that auto-approve read-only tools don't accidentally clear data.
- ``refresh_cache`` — DESTRUCTIVE. Re-download operation that overwrites cached
  data. Separate from cache_manage to make intent explicit (refresh ≠ delete).
"""
from __future__ import annotations

import asyncio
import io
import logging
from datetime import datetime, timezone
from typing import Any

import httpx
import pandas as pd
from fastmcp import Context

from ontario_data.ckan_client import CKANClient
from ontario_data.server import DESTRUCTIVE, READONLY, mcp
from ontario_data.staleness import compute_expires_at, get_staleness_info
from ontario_data.formatting import md_response
from ontario_data.utils import (
    get_lifespan_state,
    get_cache,
    get_deps,
    infer_portal_from_table,
    is_arcgis_portal,
    make_table_name,
    parse_portal_id,
    resolve_resource_portal,
)

logger = logging.getLogger("ontario_data.retrieval")


def _read_csv_bytes(content: bytes) -> pd.DataFrame:
    """Read CSV bytes, falling back to latin-1 if the file isn't valid UTF-8."""
    try:
        return pd.read_csv(io.BytesIO(content), encoding="utf-8")
    except UnicodeDecodeError:
        return pd.read_csv(io.BytesIO(content), encoding="latin-1")


async def _download_resource_data(
    ckan: CKANClient,
    resource_id: str,
    http_client: httpx.AsyncClient,
) -> tuple[pd.DataFrame, dict[str, Any], dict[str, Any]]:
    """Fetch resource data, preferring the CKAN datastore (structured API)
    and falling back to direct file download for CSV/XLSX/JSON/GeoJSON."""
    resource = await ckan.resource_show(resource_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", "")

    # Try datastore first (structured data)
    if resource.get("datastore_active"):
        result = await ckan.datastore_search_all(resource_id)
        df = pd.DataFrame(result["records"])
        internal_cols = [c for c in df.columns if c.startswith("_")]
        df = df.drop(columns=internal_cols, errors="ignore")
        return df, resource, dataset

    if not url:
        raise ValueError(f"Resource '{resource_id}' has no download URL and datastore is inactive")

    # Download file directly using shared client with extended timeout
    response = await http_client.get(url, timeout=120.0, follow_redirects=True)
    response.raise_for_status()
    content = response.content

    if fmt in ("CSV", "TXT"):
        df = _read_csv_bytes(content)
    elif fmt in ("XLS", "XLSX"):
        df = pd.read_excel(io.BytesIO(content))
    elif fmt == "JSON":
        df = pd.read_json(io.BytesIO(content))
    elif fmt == "GEOJSON":
        import geopandas as gpd
        gdf = gpd.read_file(io.BytesIO(content))
        df = pd.DataFrame(gdf)
        if "geometry" in df.columns:
            df["geometry_wkt"] = gdf["geometry"].apply(lambda g: g.wkt if g else None)
            df = df.drop(columns=["geometry"])
    else:
        raise ValueError(f"Unsupported format for tabular import: {fmt}. URL: {url}")

    return df, resource, dataset


async def _download_arcgis_resource_data(
    client,
    resource_id: str,
    http_client: httpx.AsyncClient,
) -> tuple[pd.DataFrame, dict[str, Any], dict[str, Any]]:
    """Fetch ArcGIS Hub resource data via Downloads API (bulk CSV)."""
    dataset = await client.package_show(resource_id)

    csv_url = await client.get_download_url(resource_id, fmt="csv")
    if csv_url:
        resp = await http_client.get(csv_url, timeout=120.0, follow_redirects=True)
        resp.raise_for_status()
        df = _read_csv_bytes(resp.content)
        resource_meta = {
            "id": resource_id,
            "package_id": resource_id,
            "format": "CSV",
            "url": csv_url,
            "datastore_active": False,
        }
        return df, resource_meta, dataset

    raise ValueError(
        f"No CSV download available for dataset '{resource_id}'. "
        f"Try a different dataset or check the portal directly at {client.base_url}."
    )


@mcp.tool(annotations=READONLY)
async def download_resource(
    resource_id: str,
    ctx: Context = None,
) -> str:
    """Download a dataset resource and cache it locally in DuckDB for fast querying.

    Supports CSV, XLSX, JSON, and datastore-active resources.
    If already cached, returns staleness info so you can decide whether to refresh.
    Numeric columns stored as text are automatically cast to DOUBLE.

    Workflow: search_datasets → get_dataset_info → download_resource → query_cached.

    Args:
        resource_id: Prefixed resource ID (e.g. "toronto:abc123") or bare ID
    """
    configs = get_lifespan_state(ctx)["portal_configs"]
    portal, bare_id = parse_portal_id(resource_id, set(configs.keys()))

    cache = get_cache(ctx)

    if cache.is_cached(bare_id):
        table_name = cache.get_table_name(bare_id)
        meta = cache.get_resource_meta(bare_id)
        staleness = get_staleness_info(cache, bare_id)
        return md_response(
            status="already_cached",
            table_name=table_name,
            row_count=meta["row_count"],
            downloaded_at=str(meta["downloaded_at"]),
            staleness=staleness,
            hint=f'Ready to query: SELECT * FROM "{table_name}" LIMIT 10. Use refresh_cache(resource_id=...) to re-download.',
        )

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

    ckan, _ = get_deps(ctx, portal)

    await ctx.report_progress(0, 100, "Downloading resource...")

    http_client = get_lifespan_state(ctx)["http_client"]
    if is_arcgis_portal(ctx, portal):
        df, resource, dataset = await _download_arcgis_resource_data(ckan, bare_id, http_client)
    else:
        df, resource, dataset = await _download_resource_data(ckan, bare_id, http_client)

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

    table_name = make_table_name(dataset.get("name", ""), bare_id, portal=portal)
    await asyncio.to_thread(
        cache.store_resource,
        resource_id=bare_id,
        dataset_id=dataset.get("id", ""),
        table_name=table_name,
        df=df,
        source_url=resource.get("url", ""),
    )
    await asyncio.to_thread(
        cache.store_dataset_metadata,
        dataset.get("id", ""),
        dataset,
    )

    # Set staleness expiry based on update frequency
    update_freq = dataset.get("update_frequency")
    expires_at = compute_expires_at(datetime.now(timezone.utc), update_freq)
    await asyncio.to_thread(cache.update_expires_at, bare_id, expires_at)

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

    return md_response(
        status="downloaded",
        table_name=table_name,
        row_count=len(df),
        columns=list(df.columns),
        dtypes={col: str(dtype) for col, dtype in df.dtypes.items()},
        hint=f'Use query_cached tool with SQL like: SELECT * FROM "{table_name}" LIMIT 10',
    )


@mcp.tool(annotations=READONLY)
async def cache_info(ctx: Context = None) -> str:
    """Get cache statistics and list all cached datasets.

    Returns size, table count, and details for every cached resource.
    """
    cache = get_cache(ctx)
    stats = cache.get_stats()
    cached = cache.list_cached()

    # Add staleness info for each cached resource
    items = []
    for c in cached:
        staleness = get_staleness_info(cache, c["resource_id"])
        size_bytes = c.get("size_bytes", 0) or 0
        items.append({
            "table_name": c["table_name"],
            "resource_id": c["resource_id"],
            "dataset_id": c["dataset_id"],
            "row_count": c["row_count"],
            "size_bytes": size_bytes,
            "size_mb": round(size_bytes / (1024 * 1024), 2),
            "downloaded_at": c["downloaded_at"],
            "is_stale": staleness["is_stale"] if staleness else None,
        })

    return md_response(
        **stats,
        total_size_mb=round(stats["total_size_bytes"] / (1024 * 1024), 2),
        datasets=items,
    )


@mcp.tool(annotations=DESTRUCTIVE)
async def cache_manage(
    action: str,
    resource_id: str | None = None,
    ctx: Context = None,
) -> str:
    """Manage the local DuckDB cache: remove or clear cached data.

    Args:
        action: One of "remove" (single resource) or "clear" (all)
        resource_id: Required for "remove" action. Prefixed or bare ID accepted.
    """
    cache = get_cache(ctx)

    if action == "remove":
        if not resource_id:
            raise ValueError("resource_id is required for 'remove' action")
        configs = get_lifespan_state(ctx)["portal_configs"]
        _, bare_id = parse_portal_id(resource_id, set(configs.keys()))
        cache.remove_resource(bare_id)
        return md_response(status="removed", resource_id=bare_id)

    elif action == "clear":
        count = len(cache.list_cached())
        cache.remove_all()
        return md_response(status="cleared", removed_count=count)

    else:
        raise ValueError(f"Invalid action '{action}'. Use 'remove' or 'clear'.")


@mcp.tool(annotations=DESTRUCTIVE)
async def refresh_cache(
    resource_id: str | None = None,
    ctx: Context = None,
) -> str:
    """Re-download cached resources to get the latest data.

    Args:
        resource_id: Specific resource to refresh (prefixed or bare ID), or omit to refresh all
    """
    cache = get_cache(ctx)
    cached = cache.list_cached()

    bare_id = None
    if resource_id:
        configs = get_lifespan_state(ctx)["portal_configs"]
        _, bare_id = parse_portal_id(resource_id, set(configs.keys()))
        cached = [c for c in cached if c["resource_id"] == bare_id]
        if not cached:
            raise ValueError(f"Resource {bare_id} not found in cache")

    http_client = get_lifespan_state(ctx)["http_client"]
    results = []
    for i, item in enumerate(cached):
        await ctx.report_progress(i, len(cached), f"Refreshing {item['table_name']}...")
        try:
            portal = infer_portal_from_table(item["table_name"])

            ckan, _ = get_deps(ctx, portal)
            if is_arcgis_portal(ctx, portal):
                df, resource, dataset = await _download_arcgis_resource_data(ckan, item["resource_id"], http_client)
            else:
                df, resource, dataset = await _download_resource_data(ckan, item["resource_id"], http_client)
            await asyncio.to_thread(
                cache.store_resource,
                resource_id=item["resource_id"],
                dataset_id=item["dataset_id"],
                table_name=item["table_name"],
                df=df,
                source_url=item["source_url"],
            )
            update_freq = dataset.get("update_frequency")
            expires_at = compute_expires_at(datetime.now(timezone.utc), update_freq)
            await asyncio.to_thread(cache.update_expires_at, item["resource_id"], expires_at)
            results.append({"resource_id": item["resource_id"], "status": "refreshed", "new_row_count": len(df)})
        except Exception as e:
            results.append({"resource_id": item["resource_id"], "status": "error", "error": str(e)})

    return md_response(refreshed=results)
