from __future__ import annotations

import asyncio
import re
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, TypeVar

from ontario_data.cache import CacheManager, InvalidQueryError  # noqa: F401
from ontario_data.portals import PORTALS, PortalType
from ontario_data.protocols import PortalClient

if TYPE_CHECKING:
    from fastmcp import Context

# Re-export for consumers
__all__ = [
    "ResourceNotCachedError", "DatastoreNotAvailableError", "SpatialExtensionError",
    "get_deps", "get_cache", "parse_portal_id", "fan_out", "unwrap_first_match",
    "resolve_dataset", "resolve_resource_portal", "strip_internal_fields",
    "make_table_name", "make_geo_table_name", "require_cached", "infer_portal_from_table",
    "arcgis_guard", "is_arcgis_portal", "get_lifespan_state",
]

T = TypeVar("T")


class ResourceNotCachedError(Exception):
    """Raised when a tool requires cached data that doesn't exist."""
    pass


class DatastoreNotAvailableError(Exception):
    """Raised when a resource has no datastore."""
    pass


class SpatialExtensionError(Exception):
    """Raised when DuckDB spatial extension is not available."""
    pass


def get_lifespan_state(ctx: Context) -> dict:
    return ctx.lifespan_context


def get_deps(ctx: Context, portal: str) -> tuple[PortalClient, CacheManager]:
    """Extract portal client and cache manager from MCP context.

    Lazily creates the client for the requested portal on first use.
    portal is required — there is no default.
    """
    state = get_lifespan_state(ctx)
    configs = state["portal_configs"]

    if portal not in configs:
        available = list(configs.keys())
        raise ValueError(f"Unknown portal '{portal}'. Available: {available}")

    clients = state["portal_clients"]
    if portal not in clients:
        config = configs[portal]
        if config.portal_type == PortalType.CKAN:
            from ontario_data.ckan_client import CKANClient
            clients[portal] = CKANClient(
                base_url=config.base_url,
                http_client=state["http_client"],
            )
        elif config.portal_type == PortalType.ARCGIS_HUB:
            from ontario_data.arcgis_client import ArcGISHubClient
            clients[portal] = ArcGISHubClient(
                base_url=config.base_url,
                http_client=state["http_client"],
                org_name=portal,
                org_title=config.name.replace(" Open Data", ""),
                owner_filter=config.owner_filter,
            )
        else:
            raise ValueError(f"Portal '{portal}' uses unknown type {config.portal_type}.")

    return clients[portal], state["cache"]


def parse_portal_id(id_str: str, known_portals: set[str]) -> tuple[str | None, str]:
    """Split 'portal:bare_id'. Returns (None, id_str) if no valid prefix."""
    if ":" in id_str:
        prefix, rest = id_str.split(":", 1)
        if prefix in known_portals:
            return prefix, rest
    return None, id_str


async def fan_out(
    ctx: Context,
    portal: str | None,
    fn: Callable[[str], Awaitable[T]],
    *,
    first_match: bool = False,
) -> list[tuple[str, T | None, str | None]]:
    """Run fn(portal_key) across portals.

    first_match=False (default): asyncio.gather all portals, collect all results.
      Returns [(portal_key, result_or_None, error_or_None), ...].
    first_match=True: try portals sequentially in PORTALS dict order.
      Return on first success. Swallow errors and continue.
      On all-fail, return all errors so caller can build a diagnostic message.

    If portal is specified, only run against that one portal.
    No timeout — relies on each client's own 30s timeout + retry.
    """
    configs = get_lifespan_state(ctx)["portal_configs"]

    if portal:
        keys = [portal]
    else:
        keys = list(configs.keys())

    if first_match:
        errors: list[tuple[str, None, str]] = []
        for key in keys:
            try:
                result = await fn(key)
                return [(key, result, None)]
            except Exception as exc:
                errors.append((key, None, str(exc)))
        return errors

    # Parallel fan-out
    async def _safe(key: str) -> tuple[str, T | None, str | None]:
        try:
            result = await fn(key)
            return (key, result, None)
        except Exception as exc:
            return (key, None, str(exc))

    return list(await asyncio.gather(*[_safe(k) for k in keys]))


_TRANSIENT_INDICATORS = ("429", "500", "502", "503", "504", "timeout", "timed out", "connect")


def unwrap_first_match(
    results: list[tuple[str, T | None, str | None]],
    bare_id: str,
    entity_type: str = "Dataset",
) -> tuple[str, T]:
    """Extract the first success from ``fan_out(first_match=True)`` results.

    Returns ``(portal, result)``.  Raises :class:`ValueError` with a
    consistent, user-friendly message when every portal failed.

    Distinguishes transient errors (rate limits, 5xx, timeouts) from genuine
    not-found responses so callers know whether to retry vs. search again.
    """
    if results and results[0][2] is None:
        return results[0][0], results[0][1]
    errors = (
        "; ".join(f"{pk}: {err}" for pk, _, err in results)
        if results
        else "no portals available"
    )
    has_transient = any(
        any(indicator in (err or "").lower() for indicator in _TRANSIENT_INDICATORS)
        for _, _, err in (results or [])
    )
    if has_transient:
        raise ValueError(
            f"{entity_type} '{bare_id}': one or more portals returned a transient error "
            f"(rate limit or server error) — retry, or use a prefixed ID "
            f"(e.g. 'ontario:{bare_id}') to target one portal. Details: {errors}"
        )
    raise ValueError(
        f"{entity_type} '{bare_id}' not found. Tried: {errors}. "
        f"Use search_datasets to find the correct prefixed ID."
    )


async def resolve_dataset(
    ctx: Context, dataset_id: str
) -> tuple[str, str, dict]:
    """Resolve a (possibly bare) dataset ID to ``(portal, bare_id, ds_dict)``.

    If *dataset_id* carries a portal prefix the call goes directly to that
    portal; otherwise every configured portal is tried sequentially via
    :func:`fan_out`.
    """
    configs = get_lifespan_state(ctx)["portal_configs"]
    portal, bare_id = parse_portal_id(dataset_id, set(configs.keys()))

    async def _show(pk: str):
        client, _ = get_deps(ctx, pk)
        return await client.package_show(bare_id)

    if portal:
        client, _ = get_deps(ctx, portal)
        ds = await client.package_show(bare_id)
    else:
        results = await fan_out(ctx, None, _show, first_match=True)
        portal, ds = unwrap_first_match(results, bare_id, "Dataset")

    return portal, bare_id, ds


async def resolve_resource_portal(
    ctx: Context, resource_id: str
) -> tuple[str, str]:
    """Resolve a (possibly bare) resource ID to ``(portal, bare_id)``.

    Unlike :func:`resolve_dataset` this does **not** return the resource
    dict — callers typically need to make their own follow-up API call
    (e.g. ``resource_show`` or ``datastore_search``) after knowing the portal.
    """
    configs = get_lifespan_state(ctx)["portal_configs"]
    portal, bare_id = parse_portal_id(resource_id, set(configs.keys()))

    if portal:
        return portal, bare_id

    async def _try(pk: str):
        client, _ = get_deps(ctx, pk)
        await client.resource_show(bare_id)
        return pk

    results = await fan_out(ctx, None, _try, first_match=True)
    portal, _ = unwrap_first_match(results, bare_id, "Resource")
    return portal, bare_id


def get_cache(ctx: Context) -> CacheManager:
    return get_lifespan_state(ctx)["cache"]


def strip_internal_fields(records: list[dict]) -> list[dict]:
    """Strip CKAN bookkeeping columns (_id, _full_text, etc.) that clutter
    results returned to the LLM."""
    return [{k: v for k, v in r.items() if not k.startswith("_")} for r in records]


def _slugify_table(name: str, fallback: str = "unknown", max_len: int = 40) -> str:
    """Lowercase, collapse non-alphanumerics to underscores, truncate.
    e.g. 'Ontario COVID-19 Cases' → 'ontario_covid_19_cases'."""
    slug = re.sub(r"[^a-z0-9]", "_", (name or fallback).lower())
    return re.sub(r"_+", "_", slug).strip("_")[:max_len]


def infer_portal_from_table(table_name: str) -> str:
    """Extract portal key from a table name like 'ds_toronto_...' or 'geo_ottawa_...'.

    Inverse of make_table_name / make_geo_table_name. Falls back to 'ontario'
    for legacy names or unexpected formats.
    """
    parts = table_name.split("_", 2)
    if len(parts) >= 3 and parts[0] in ("ds", "geo") and parts[1] in PORTALS:
        return parts[1]
    return "ontario"


def make_table_name(dataset_name: str, resource_id: str, portal: str = "ontario") -> str:
    """Build a deterministic table name like 'ds_ontario_covid_cases_a1b2c3d4'
    so the same resource always lands in the same table."""
    slug = _slugify_table(dataset_name)
    return f"ds_{portal}_{slug}_{resource_id[:8]}"


def make_geo_table_name(dataset_name: str, resource_id: str, portal: str = "ontario") -> str:
    """Like make_table_name but prefixed 'geo_' so spatial tables are
    distinguishable in cache listings."""
    slug = _slugify_table(dataset_name, fallback="geo")
    return f"geo_{portal}_{slug}_{resource_id[:8]}"


def arcgis_guard(resource_id: str, alternative: str = "download_resource + query_cached") -> str:
    """Return a user-friendly 'not available' message for ArcGIS Hub resources.

    Use in tools that rely on CKAN datastore APIs (query_resource, sql_query,
    preview_data, get_resource_schema).
    """
    return (
        "**Not available:** ArcGIS Hub has no remote datastore API.\n\n"
        f"Use {alternative} instead."
    )


def is_arcgis_portal(ctx: Context, portal: str) -> bool:
    """Check if a portal is an ArcGIS Hub portal."""
    configs = get_lifespan_state(ctx)["portal_configs"]
    return portal in configs and configs[portal].portal_type == PortalType.ARCGIS_HUB


def require_cached(cache: CacheManager, resource_id: str) -> str:
    """Return table name, or raise with a user-friendly message that includes
    the download_resource command to run.

    Strips any portal prefix (e.g. 'toronto:abc123' -> 'abc123') before
    cache lookup, since the cache stores bare IDs only.
    """
    _, bare_id = parse_portal_id(resource_id, set(PORTALS.keys()))
    table_name = cache.get_table_name(bare_id)
    if not table_name:
        raise ResourceNotCachedError(
            f"Resource {resource_id} is not cached. "
            f"Use download_resource(resource_id='{resource_id}') first."
        )
    return table_name


