"""Querying tools: run SQL against cached or remote data.

These four tools expose two distinct backends — local DuckDB and the remote
CKAN datastore API — as separate tools rather than a unified interface. The
split is intentional and each tool exists for a specific reason.

Why not one merged query tool that auto-selects the backend?

- ``query_cached`` — the primary query tool. Fast, composable, supports full
  DuckDB SQL (CTEs, window functions, JOINs across multiple cached datasets).
  Should be the default after any download_resource call.
- ``query_resource`` — CKAN Datastore structured API (filters as dict, no SQL).
  Exists for quick filtered fetches on datastore-active resources *without
  downloading* the full dataset first. Useful when the model only needs a
  filtered slice, not the whole table. Cannot be merged with query_cached
  because it requires no prior download.
- ``sql_query`` — CKAN Datastore raw SQL endpoint (remote). Exists for
  one-off exploratory queries on datastore resources when downloading isn't
  warranted. Rate-limited (429-prone); the docstring warns to prefer
  query_cached for repeated use. ArcGIS portals have no datastore, so this
  tool is CKAN-only.
- ``preview_data`` — fetches the first N rows via CKAN datastore *without
  downloading the full resource*. Distinct from query_resource because it
  has no filter/field parameters — it's purely "show me what this looks like"
  before deciding whether to download. Replacing this with LIMIT 10 in
  query_cached would require a full download of potentially 500MB to see 10 rows.

Consolidation considered: query_resource and preview_data could theoretically
merge (preview_data is query_resource with limit and no filters). Kept separate
because preview_data's zero-argument simplicity signals clearly to the model
that it's a "just show me" call, while query_resource's filter/field/sort params
imply purposeful data extraction. Different mental models, different use cases.
"""
from __future__ import annotations

import logging
import re
from datetime import datetime, timezone
from typing import Any

from fastmcp import Context
from fastmcp.tools.tool import ToolResult

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

from ontario_data.cache import InvalidQueryError
from ontario_data.formatting import format_records
from ontario_data.server import READONLY, mcp
from ontario_data.utils import (
    get_lifespan_state,
    arcgis_guard,
    fan_out,
    get_cache,
    get_deps,
    is_arcgis_portal,
    parse_portal_id,
    strip_internal_fields,
    unwrap_first_match,
)

MAX_QUERY_ROWS = 2000

_QUERY_CACHED_OUTPUT_SCHEMA = {
    "type": "object",
    "properties": {
        "sql": {"type": "string", "description": "The SQL query that was executed"},
        "rows": {
            "type": "array",
            "items": {"type": "object"},
            "description": "Result rows as JSON objects",
        },
        "columns": {
            "type": "array",
            "items": {"type": "string"},
            "description": "Column names in result order",
        },
        "total_rows": {
            "type": "integer",
            "description": "Total rows matched (may exceed rows if truncated)",
        },
        "truncated": {
            "type": "boolean",
            "description": "True if results were truncated to MAX_QUERY_ROWS",
        },
        "warnings": {
            "type": "array",
            "items": {"type": "string"},
            "description": "Heuristic warnings (e.g. COUNT(*) vs SUM suggestion)",
        },
    },
    "required": ["sql", "rows", "columns", "truncated", "warnings"],
}

_COUNT_STAR_RE = re.compile(r"\bCOUNT\s*\(\s*\*\s*\)", re.IGNORECASE)
_GROUP_BY_RE = re.compile(r"\bGROUP\s+BY\b", re.IGNORECASE)
_TABLE_RE = re.compile(r'FROM\s+"([^"]+)"', re.IGNORECASE)
_QUANTITY_PATTERNS = re.compile(
    r"(count|quantity|number|no_of|total|amount|exceedances|num_)",
    re.IGNORECASE,
)



def _generate_query_warnings(
    sql: str,
    results: list[dict[str, Any]],
    fields: list[dict[str, str]],
    cache,
) -> list[str]:
    """Post-query heuristic warnings to catch common accuracy traps."""
    warnings: list[str] = []

    # Extract table name(s) from SQL
    table_matches = _TABLE_RE.findall(sql)

    # 1. COUNT(*) when a quantity column exists
    if _COUNT_STAR_RE.search(sql) and table_matches:
        table = table_matches[0]
        try:
            columns = cache.execute_sql(f'DESCRIBE "{table}"')
            col_names = [c[0] for c in columns]
            quantity_cols = [c for c in col_names if _QUANTITY_PATTERNS.search(c)]
            if quantity_cols:
                warnings.append(
                    f"This table has columns that may contain per-row counts: "
                    f"{quantity_cols}. Consider SUM(\"{quantity_cols[0]}\") instead of COUNT(*)."
                )
        except Exception:
            logger.debug("Failed to check quantity columns for table %s", table, exc_info=True)

    # 2. 0 rows from non-empty table
    if not results and table_matches:
        table = table_matches[0]
        try:
            total = cache.execute_sql(f'SELECT COUNT(*) FROM "{table}"')[0][0]
            if total > 0:
                warnings.append(
                    f"0 rows returned but table has {total:,} rows. "
                    f"Check your WHERE/JOIN conditions."
                )
        except Exception:
            logger.debug("Failed to check row count for table %s", table, exc_info=True)

    # 3. Very few rows from GROUP BY on large table
    if 1 <= len(results) <= 3 and _GROUP_BY_RE.search(sql) and table_matches:
        table = table_matches[0]
        try:
            total = cache.execute_sql(f'SELECT COUNT(*) FROM "{table}"')[0][0]
            if total > 1000:
                warnings.append(
                    f"Only {len(results)} groups from {total:,} rows. "
                    f"Verify your GROUP BY columns are correct."
                )
        except Exception:
            logger.debug("Failed to check GROUP BY for table %s", table, exc_info=True)

    return warnings


@mcp.tool(annotations=READONLY)
async def query_resource(
    resource_id: str,
    filters: dict[str, Any] | None = None,
    fields: list[str] | None = None,
    sort: str | None = None,
    limit: int = 100,
    offset: int = 0,
    ctx: Context = None,
) -> str:
    """Query a resource via the CKAN Datastore API (remote, no download needed).

    Only works for resources with datastore_active=True.

    Args:
        resource_id: Prefixed resource ID (e.g. "toronto:abc123") or bare ID
        filters: Column filters as {column: value} pairs
        fields: List of columns to return (default: all)
        sort: Sort string (e.g. "date desc", "name asc")
        limit: Max rows (1-1000)
        offset: Row offset for pagination
    """
    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)

    async def _query(pk: str):
        ckan, _ = get_deps(ctx, pk)
        return await ckan.datastore_search(
            resource_id=bare_id,
            filters=filters,
            fields=fields,
            sort=sort,
            limit=min(limit, 1000),
            offset=offset,
        )

    if portal:
        result = await _query(portal)
    else:
        results = await fan_out(ctx, None, _query, first_match=True)
        _, result = unwrap_first_match(results, bare_id, "Resource")

    field_info = [{"name": f["id"], "type": f.get("type")} for f in result.get("fields", []) if not f["id"].startswith("_")]
    clean_records = strip_internal_fields(result.get("records", []))

    return format_records(clean_records, row_count=len(clean_records), total=result.get("total", 0), fields=field_info)


@mcp.tool(annotations=READONLY)
async def sql_query(
    sql: str,
    portal: str = "ontario",
    ctx: Context = None,
) -> str:
    """Run a SQL query against the CKAN Datastore (remote).

    NOTE: Prefer download_resource + query_cached for repeated queries —
    the remote API has rate limits (429 errors). Use this tool only for
    quick one-off queries on datastore-active resources.

    Use resource IDs as table names in double quotes.
    Example: SELECT "Column Name" FROM "resource-id-here" WHERE "Year" > 2020 LIMIT 10

    Args:
        sql: SQL query string (read-only, SELECT only)
        portal: Portal to query (default: "ontario"). Required because SQL embeds resource IDs directly.
    """
    ckan, _ = get_deps(ctx, portal)
    if is_arcgis_portal(ctx, portal):
        return arcgis_guard("", alternative="download_resource + query_cached")

    result = await ckan.datastore_sql(sql)
    field_info = [{"name": f["id"], "type": f.get("type")} for f in result.get("fields", []) if not f["id"].startswith("_")]
    clean_records = strip_internal_fields(result.get("records", []))

    return format_records(clean_records, row_count=len(clean_records), fields=field_info)


@mcp.tool(annotations=READONLY, output_schema=_QUERY_CACHED_OUTPUT_SCHEMA)
async def query_cached(
    sql: str,
    ctx: Context = None,
) -> ToolResult:
    """Run a SQL query against locally cached data in DuckDB.

    Use table names from download_resource or cache_info.
    Supports full DuckDB SQL: aggregations, window functions, CTEs, JOINs across tables.

    Use SUM(quantity_col) not COUNT(*) when rows contain per-row counts (e.g.
    a "count" or "number_of" column). COUNT(*) counts rows, not quantities.
    Column names vary across resources in the same dataset — always DESCRIBE first.
    Use SELECT * RENAME ("old col" AS new_col) to normalize column names when joining across datasets.
    Values containing semicolons should be matched with LIKE patterns, not = equality.
    Quote table names with double quotes in SQL.

    Args:
        sql: SQL query (e.g. SELECT * FROM "ds_my_table_abc12345" LIMIT 10)
    """
    cache = get_cache(ctx)
    try:
        results, fields = cache.query_with_meta(sql, max_rows=MAX_QUERY_ROWS)

        # --- Truncation (Item 2) ---
        truncated_total = None
        if len(results) > MAX_QUERY_ROWS:
            try:
                count_row = cache.execute_sql(f"SELECT COUNT(*) FROM ({sql})")
                truncated_total = count_row[0][0]
            except Exception:
                truncated_total = len(results)
            results = results[:MAX_QUERY_ROWS]

        # --- Post-query heuristic warnings (Item 5) ---
        warnings = _generate_query_warnings(sql, results, fields, cache)

        # --- Build response ---
        parts: list[str] = []

        # Echo SQL (Item 3)
        parts.append(f"**Query:** `{sql}`")

        # Format the records table with column types
        table_output = format_records(
            results,
            row_count=len(results),
            total=truncated_total,
            fields=fields,
        )
        parts.append(table_output)

        # Truncation warning
        if truncated_total is not None and truncated_total > MAX_QUERY_ROWS:
            parts.append(
                f"\n**Warning:** Results truncated to {MAX_QUERY_ROWS:,} of "
                f"{truncated_total:,} rows. Add a LIMIT, WHERE, or aggregation "
                f"to your query for accurate analysis."
            )

        # Append warnings
        if warnings:
            parts.append("")
            for w in warnings:
                parts.append(f"⚠ {w}")

        # --- Data provenance (Item 10) ---
        table_names_in_sql = _TABLE_RE.findall(sql)
        if table_names_in_sql:
            try:
                table_metas = cache.get_tables_metadata(table_names_in_sql)
                for tm in table_metas:
                    downloaded = str(tm["downloaded_at"]).split(".")[0] if tm["downloaded_at"] else "unknown"
                    expires = tm.get("expires_at")
                    if expires:
                        now = datetime.now(timezone.utc)
                        is_stale = now > expires if hasattr(expires, '__gt__') else False
                        status = "**stale**" if is_stale else "fresh"
                    else:
                        status = "unknown"
                    parts.append(
                        f"\n**Source:** `{tm['table_name']}` | "
                        f"Resource: `{tm['resource_id']}` | "
                        f"Downloaded: {downloaded} | "
                        f"Status: {status}"
                    )
            except Exception:
                logger.debug("Failed to build data provenance for tables %s", table_names_in_sql, exc_info=True)

        markdown = "\n".join(parts)
        col_names = [f["name"] for f in fields] if fields else []
        return ToolResult(
            content=markdown,
            structured_content={
                "sql": sql,
                "rows": results,
                "columns": col_names,
                "total_rows": truncated_total if truncated_total is not None else len(results),
                "truncated": truncated_total is not None,
                "warnings": warnings,
            },
        )
    except Exception as e:
        cached = cache.list_cached()
        table_names = [c["table_name"] for c in cached]
        hints = ["Quote table names with double quotes."]
        augmented = f"{e}\n\nAvailable tables: {table_names}\nHints: {' '.join(hints)}"
        raise InvalidQueryError(augmented) from e


@mcp.tool(annotations=READONLY)
async def preview_data(
    resource_id: str,
    rows: int = 10,
    ctx: Context = None,
) -> str:
    """Quick preview of the first N rows of a resource (fetched remotely).

    Args:
        resource_id: Prefixed resource ID (e.g. "toronto:abc123") or bare ID
        rows: Number of rows to preview (1-100)
    """
    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)

    async def _preview(pk: str):
        ckan, _ = get_deps(ctx, pk)
        return await ckan.datastore_search(bare_id, limit=min(rows, 100))

    if portal:
        result = await _preview(portal)
    else:
        results = await fan_out(ctx, None, _preview, first_match=True)
        _, result = unwrap_first_match(results, bare_id, "Resource")

    field_info = [{"name": f["id"], "type": f.get("type")} for f in result.get("fields", []) if not f["id"].startswith("_")]
    clean_records = strip_internal_fields(result.get("records", []))

    return format_records(clean_records, row_count=len(clean_records), total=result.get("total", 0), preview=True, fields=field_info)
