"""KAPSARC Data Portal MCP Server.

Wraps the KAPSARC Opendatasoft Explore API v2.1 so an LLM can discover
and query energy/economic datasets from https://datasource.kapsarc.org.
"""

from __future__ import annotations

import os
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass

import httpx
from mcp.server.fastmcp import Context, FastMCP

from kapsarc_mcp_server.charts import render_chart

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
BASE_URL = "https://datasource.kapsarc.org/api/explore/v2.1"
DEFAULT_LIMIT = 10
MAX_RECORD_LIMIT = 100
ALLOWED_EXPORT_FORMATS = {"json", "csv"}

# ---------------------------------------------------------------------------
# Application context (holds the shared httpx client)
# ---------------------------------------------------------------------------


@dataclass
class AppContext:
    client: httpx.AsyncClient


@asynccontextmanager
async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
    """Create a shared httpx client for the lifetime of the server."""
    api_key = os.environ.get("KAPSARC_API_KEY")
    params = {}
    if api_key:
        params["apikey"] = api_key
    async with httpx.AsyncClient(
        base_url=BASE_URL,
        params=params,
        timeout=30.0,
        headers={"Accept": "application/json"},
    ) as client:
        yield AppContext(client=client)


mcp = FastMCP("KAPSARC Data Portal", lifespan=app_lifespan)

# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


def _build_params(**kwargs: object) -> dict[str, object]:
    """Build a query-param dict, dropping keys whose value is None."""
    return {k: v for k, v in kwargs.items() if v is not None}


async def _api_get(
    ctx: Context,
    path: str,
    params: dict[str, object] | None = None,
) -> dict | list | str:
    """Perform a GET request against the KAPSARC API.

    Returns the parsed JSON on success or an error string on failure.
    """
    app: AppContext = ctx.request_context.lifespan_context
    try:
        resp = await app.client.get(path, params=params)
    except httpx.HTTPError as exc:
        return f"HTTP error: {exc}"

    if resp.status_code == 200:
        return resp.json()

    status = resp.status_code
    if status == 400:
        return f"Bad request (400): check ODSQL syntax. {resp.text}"
    if status == 401:
        return "Unauthorized (401): set the KAPSARC_API_KEY environment variable."
    if status == 404:
        return f"Not found (404): {resp.text}"
    if status == 429:
        return "Rate limited (429): too many requests — try again shortly."
    return f"API error ({status}): {resp.text}"


# ---------------------------------------------------------------------------
# Helper: format a dataset summary for display
# ---------------------------------------------------------------------------


def _format_dataset_summary(ds: dict) -> str:
    """Return a concise multi-line summary of a dataset entry."""
    meta = ds.get("dataset", ds) if "dataset" in ds else ds
    dataset_id = meta.get("dataset_id", "?")
    metas = meta.get("metas", {})
    default = metas.get("default", {})
    title = default.get("title", dataset_id)
    description = default.get("description", "")
    theme = default.get("theme", [])
    keyword = default.get("keyword", [])
    records = default.get("records_count", "?")
    modified = default.get("modified", "?")

    lines = [
        f"**{title}**",
        f"  ID: `{dataset_id}`",
        f"  Records: {records} | Modified: {modified}",
    ]
    if theme:
        lines.append(f"  Themes: {', '.join(theme)}")
    if keyword:
        lines.append(f"  Keywords: {', '.join(keyword[:10])}")
    if description:
        desc_short = description[:200].replace("\n", " ")
        if len(description) > 200:
            desc_short += "…"
        lines.append(f"  Description: {desc_short}")
    return "\n".join(lines)


# ===========================================================================
# Tier 1 — Core tools
# ===========================================================================


@mcp.tool()
async def search_datasets(
    ctx: Context,
    query: str | None = None,
    where: str | None = None,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
) -> str:
    """Search and browse KAPSARC datasets by keyword, filter, or theme.

    Args:
        query: Free-text search (searches title, description, keywords).
               Internally wrapped in the ODSQL search() function.
        where: ODSQL filter expression for structured queries. Examples:
               - 'default.theme = "Energy"'
               - 'default.keyword = "oil"'
               Query and where combine with AND when both are provided.
        limit: Max datasets to return (default 10, max 100).
        offset: Pagination offset.

    Returns:
        Formatted list of matching datasets with IDs, titles, themes, and descriptions.
    """
    limit = min(max(1, limit), MAX_RECORD_LIMIT)
    params = _build_params(limit=limit, offset=offset, where=where)
    if query:
        search_clause = f"search({query!r})"
        if where:
            params["where"] = f"{where} AND {search_clause}"
        else:
            params["where"] = search_clause

    data = await _api_get(ctx, "/catalog/datasets", params)
    if isinstance(data, str):
        return data

    total = data.get("total_count", "?")
    datasets = data.get("results", [])
    if not datasets:
        return "No datasets found."

    parts = [f"Found {total} dataset(s) (showing {offset + 1}–{offset + len(datasets)}):\n"]
    for ds in datasets:
        parts.append(_format_dataset_summary(ds))
        parts.append("")  # blank line separator
    return "\n".join(parts)


@mcp.tool()
async def get_dataset(ctx: Context, dataset_id: str) -> str:
    """Get full metadata and field definitions for a specific dataset.

    Args:
        dataset_id: The dataset identifier (e.g. 'king-abdullah-petroleum-studies-and-research-center-kapsarc').

    Returns:
        Dataset title, description, themes, keywords, field list with types, and record count.
    """
    data = await _api_get(ctx, f"/catalog/datasets/{dataset_id}")
    if isinstance(data, str):
        return data

    meta = data.get("dataset", data)
    dataset_uid = meta.get("dataset_id", dataset_id)
    metas = meta.get("metas", {})
    default = metas.get("default", {})
    title = default.get("title", dataset_uid)
    description = default.get("description", "")
    theme = default.get("theme", [])
    keyword = default.get("keyword", [])
    records = default.get("records_count", "?")
    modified = default.get("modified", "?")
    publisher = default.get("publisher", "?")

    lines = [
        f"# {title}",
        f"ID: `{dataset_uid}`",
        f"Publisher: {publisher}",
        f"Records: {records} | Modified: {modified}",
    ]
    if theme:
        lines.append(f"Themes: {', '.join(theme)}")
    if keyword:
        lines.append(f"Keywords: {', '.join(keyword)}")
    if description:
        lines.append(f"\n## Description\n{description}")

    fields = meta.get("fields", [])
    if fields:
        lines.append("\n## Fields")
        for f in fields:
            name = f.get("name", "?")
            ftype = f.get("type", "?")
            label = f.get("label", name)
            desc = f.get("description", "")
            entry = f"- **{label}** (`{name}`, {ftype})"
            if desc:
                entry += f" — {desc}"
            lines.append(entry)

    return "\n".join(lines)


@mcp.tool()
async def get_records(
    ctx: Context,
    dataset_id: str,
    where: str | None = None,
    select: str | None = None,
    group_by: str | None = None,
    order_by: str | None = None,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
) -> dict | str:
    """Query records from a KAPSARC dataset using ODSQL.

    Args:
        dataset_id: The dataset identifier.
        where: ODSQL filter expression. Operators: =, !=, >, <, >=, <=,
               IN, LIKE, AND, OR, NOT. Functions: search(), year(), month(),
               date(), date_format(). Examples:
               - "year >= 2020"
               - "country = 'Saudi Arabia' AND year >= 2015"
               - "search('renewable energy')"
        select: Comma-separated fields to return, or aggregation expressions.
                Examples: "country, year, value", "avg(value) as avg_val",
                "count(*) as total", "sum(production) as total_prod".
        group_by: Fields to group by when using aggregations.
                  Example: "country" (with select="country, avg(value) as avg_val").
        order_by: Sort expression. Example: "year DESC", "value ASC".
        limit: Max records to return (default 10, max 100).
        offset: Pagination offset.

    Returns:
        JSON object with total_count and results array.
    """
    limit = min(max(1, limit), MAX_RECORD_LIMIT)
    params = _build_params(
        where=where,
        select=select,
        group_by=group_by,
        order_by=order_by,
        limit=limit,
        offset=offset,
    )
    return await _api_get(ctx, f"/catalog/datasets/{dataset_id}/records", params)


@mcp.tool()
async def get_record(ctx: Context, dataset_id: str, record_id: str) -> dict | str:
    """Retrieve a single record by its ID.

    Args:
        dataset_id: The dataset identifier.
        record_id: The record identifier (from a previous get_records result).

    Returns:
        The full record as JSON.
    """
    return await _api_get(ctx, f"/catalog/datasets/{dataset_id}/records/{record_id}")


# ===========================================================================
# Tier 2 — Discovery helpers
# ===========================================================================


@mcp.tool()
async def get_dataset_facets(
    ctx: Context,
    dataset_id: str,
    facet: str | None = None,
) -> str:
    """List categorical field values and their counts for a dataset.

    Useful for discovering what values exist in a field before filtering.

    Args:
        dataset_id: The dataset identifier.
        facet: Restrict to a specific facet name (field). If omitted, all
               facetable fields are returned.

    Returns:
        Formatted list of facet names, values, and record counts.
    """
    params = _build_params(facet=facet)
    data = await _api_get(ctx, f"/catalog/datasets/{dataset_id}/facets", params)
    if isinstance(data, str):
        return data

    facets = data.get("facets", data) if isinstance(data, dict) else data
    if not facets:
        return "No facets available for this dataset."

    lines = []
    for group in facets if isinstance(facets, list) else [facets]:
        name = group.get("name", "?")
        lines.append(f"## {name}")
        for item in group.get("facets", []):
            val = item.get("name", item.get("value", "?"))
            count = item.get("count", "?")
            lines.append(f"  - {val} ({count})")
        lines.append("")
    return "\n".join(lines)


@mcp.tool()
async def get_catalog_facets(ctx: Context, facet: str | None = None) -> str:
    """List portal-wide themes, keywords, and publishers across all datasets.

    Useful for discovering what topics are available before searching.

    Args:
        facet: Restrict to a specific catalog facet (e.g. 'theme', 'keyword',
               'publisher'). If omitted, all catalog facets are returned.

    Returns:
        Formatted list of catalog facet names, values, and dataset counts.
    """
    params = _build_params(facet=facet)
    data = await _api_get(ctx, "/catalog/facets", params)
    if isinstance(data, str):
        return data

    facets = data.get("facets", data) if isinstance(data, dict) else data
    if not facets:
        return "No catalog facets available."

    lines = []
    for group in facets if isinstance(facets, list) else [facets]:
        name = group.get("name", "?")
        lines.append(f"## {name}")
        for item in group.get("facets", []):
            val = item.get("name", item.get("value", "?"))
            count = item.get("count", "?")
            lines.append(f"  - {val} ({count})")
        lines.append("")
    return "\n".join(lines)


# ===========================================================================
# Tier 3 — Exports & attachments
# ===========================================================================


@mcp.tool()
async def export_dataset(
    ctx: Context,
    dataset_id: str,
    format: str = "json",
    where: str | None = None,
    select: str | None = None,
    order_by: str | None = None,
    limit: int | None = None,
) -> dict | list | str:
    """Export a dataset without the 100-record limit of get_records.

    Args:
        dataset_id: The dataset identifier.
        format: Export format — 'json' or 'csv' only.
        where: ODSQL filter expression (same syntax as get_records).
        select: Comma-separated fields to include.
        order_by: Sort expression.
        limit: Max records to export. Omit for all records (may be large!).

    Returns:
        JSON array of records (format=json) or CSV text (format=csv).
    """
    fmt = format.lower()
    if fmt not in ALLOWED_EXPORT_FORMATS:
        return f"Unsupported format '{format}'. Use 'json' or 'csv'."

    params = _build_params(where=where, select=select, order_by=order_by, limit=limit)
    app: AppContext = ctx.request_context.lifespan_context
    try:
        resp = await app.client.get(
            f"/catalog/datasets/{dataset_id}/exports/{fmt}",
            params=params,
        )
    except httpx.HTTPError as exc:
        return f"HTTP error: {exc}"

    if resp.status_code != 200:
        return f"API error ({resp.status_code}): {resp.text}"

    if fmt == "csv":
        return resp.text
    return resp.json()


@mcp.tool()
async def get_dataset_attachments(ctx: Context, dataset_id: str) -> str:
    """List documentation and file attachments for a dataset.

    Args:
        dataset_id: The dataset identifier.

    Returns:
        Formatted list of attachments with titles, MIME types, and URLs.
    """
    data = await _api_get(ctx, f"/catalog/datasets/{dataset_id}/attachments")
    if isinstance(data, str):
        return data

    attachments = data.get("attachments", data.get("results", data)) if isinstance(data, dict) else data
    if not attachments:
        return "No attachments found for this dataset."

    lines = [f"Attachments for `{dataset_id}`:\n"]
    items = attachments if isinstance(attachments, list) else [attachments]
    for att in items:
        title = att.get("title", att.get("id", "?"))
        mime = att.get("mimetype", "?")
        url = att.get("href", att.get("url", ""))
        lines.append(f"- **{title}** ({mime})")
        if url:
            lines.append(f"  URL: {url}")
    return "\n".join(lines)


# ===========================================================================
# Tier 4 — Catalog-level
# ===========================================================================


@mcp.tool()
async def export_catalog(
    ctx: Context,
    format: str = "json",
    where: str | None = None,
    limit: int | None = None,
) -> dict | list | str:
    """Bulk export of all dataset metadata in the catalog.

    Args:
        format: Export format — 'json' or 'csv' only.
        where: ODSQL filter on catalog metadata.
        limit: Max datasets to export. Omit for all.

    Returns:
        JSON array of dataset metadata (format=json) or CSV text (format=csv).
    """
    fmt = format.lower()
    if fmt not in ALLOWED_EXPORT_FORMATS:
        return f"Unsupported format '{format}'. Use 'json' or 'csv'."

    params = _build_params(where=where, limit=limit)
    app: AppContext = ctx.request_context.lifespan_context
    try:
        resp = await app.client.get(f"/catalog/exports/{fmt}", params=params)
    except httpx.HTTPError as exc:
        return f"HTTP error: {exc}"

    if resp.status_code != 200:
        return f"API error ({resp.status_code}): {resp.text}"

    if fmt == "csv":
        return resp.text
    return resp.json()


@mcp.tool()
async def list_export_formats(ctx: Context, dataset_id: str) -> str:
    """Check which export formats are available for a dataset.

    Args:
        dataset_id: The dataset identifier.

    Returns:
        List of available export formats with their descriptions.
    """
    data = await _api_get(ctx, f"/catalog/datasets/{dataset_id}/exports")
    if isinstance(data, str):
        return data

    links = data.get("links", data) if isinstance(data, dict) else data
    if not links:
        return "No export format information available."

    items = links if isinstance(links, list) else [links]
    lines = [f"Export formats for `{dataset_id}`:\n"]
    for link in items:
        href = link.get("href", "")
        rel = link.get("rel", "?")
        lines.append(f"- {rel}: {href}")
    return "\n".join(lines)


# ===========================================================================
# Tier 5 — Visualization tools
# ===========================================================================

ALLOWED_CHART_TYPES = {"bar", "hbar", "line", "pie", "scatter"}


@mcp.tool()
async def visualize_dataset(
    ctx: Context,
    dataset_id: str,
    select: str | None = None,
    where: str | None = None,
    group_by: str | None = None,
    order_by: str | None = None,
    limit: int = 50,
    chart_type: str | None = None,
    title: str | None = None,
) -> list:
    """Fetch data from a KAPSARC dataset and render a chart image.

    Auto-detects the best chart type based on data shape. Returns both a
    text summary and a PNG chart image.

    Args:
        dataset_id: The dataset identifier.
        select: Comma-separated fields or aggregation expressions.
        where: ODSQL filter expression.
        group_by: Fields to group by when using aggregations.
        order_by: Sort expression.
        limit: Max records to fetch (default 50, max 100).
        chart_type: Force a chart type: bar, hbar, line, pie, or scatter.
                    If omitted, auto-detected from data shape.
        title: Chart title. Auto-generated if omitted.

    Returns:
        A text summary and PNG chart image.
    """
    if chart_type and chart_type not in ALLOWED_CHART_TYPES:
        return f"Invalid chart_type '{chart_type}'. Use one of: {', '.join(sorted(ALLOWED_CHART_TYPES))}"

    limit = min(max(1, limit), MAX_RECORD_LIMIT)
    params = _build_params(
        where=where,
        select=select,
        group_by=group_by,
        order_by=order_by,
        limit=limit,
    )
    data = await _api_get(ctx, f"/catalog/datasets/{dataset_id}/records", params)
    if isinstance(data, str):
        return data

    records = data.get("results", [])
    if not records:
        return "No records returned — nothing to visualize."

    try:
        image = render_chart(records, chart_type=chart_type, title=title)
    except ValueError as exc:
        return f"Chart error: {exc}"

    total = data.get("total_count", "?")
    summary = f"Chart generated from `{dataset_id}` ({len(records)} of {total} records)."
    return [summary, image]


@mcp.tool()
async def visualize_data(
    ctx: Context,
    records: list[dict],
    x_field: str | None = None,
    y_field: str | None = None,
    chart_type: str | None = None,
    title: str | None = None,
) -> list:
    """Render a chart from pre-fetched records (no API call needed).

    Use this after calling get_records or export_dataset to visualize data
    you already have.

    Args:
        records: List of record dicts (from a prior get_records call).
        x_field: Field for x-axis (auto-detected if omitted).
        y_field: Field for y-axis (auto-detected if omitted).
        chart_type: Force a chart type: bar, hbar, line, pie, or scatter.
        title: Chart title. Auto-generated if omitted.

    Returns:
        A text summary and PNG chart image.
    """
    if chart_type and chart_type not in ALLOWED_CHART_TYPES:
        return f"Invalid chart_type '{chart_type}'. Use one of: {', '.join(sorted(ALLOWED_CHART_TYPES))}"

    if not records:
        return "No records provided — nothing to visualize."

    try:
        image = render_chart(
            records,
            x_field=x_field,
            y_field=y_field,
            chart_type=chart_type,
            title=title,
        )
    except ValueError as exc:
        return f"Chart error: {exc}"

    summary = f"Chart generated from {len(records)} records."
    return [summary, image]


@mcp.tool()
async def visualize_comparison(
    ctx: Context,
    dataset_id: str,
    x_field: str,
    y_field: str,
    series_field: str,
    where: str | None = None,
    limit: int = 100,
    chart_type: str | None = None,
    title: str | None = None,
) -> list:
    """Fetch data and render a multi-series comparison chart.

    Creates grouped bar charts or multi-line charts comparing different
    categories (e.g. oil production by country over time).

    Args:
        dataset_id: The dataset identifier.
        x_field: Field for x-axis (e.g. 'year').
        y_field: Field for y-axis values (e.g. 'production').
        series_field: Field to split into series (e.g. 'country').
        where: ODSQL filter expression.
        limit: Max records to fetch (default 100, max 100).
        chart_type: Force a chart type. If omitted, auto-detected.
        title: Chart title. Auto-generated if omitted.

    Returns:
        A text summary and PNG chart image.
    """
    limit = min(max(1, limit), MAX_RECORD_LIMIT)
    params = _build_params(where=where, limit=limit)
    data = await _api_get(ctx, f"/catalog/datasets/{dataset_id}/records", params)
    if isinstance(data, str):
        return data

    records = data.get("results", [])
    if not records:
        return "No records returned — nothing to visualize."

    try:
        image = render_chart(
            records,
            x_field=x_field,
            y_field=y_field,
            series_field=series_field,
            chart_type=chart_type,
            title=title,
        )
    except ValueError as exc:
        return f"Chart error: {exc}"

    total = data.get("total_count", "?")
    series_vals = set(str(r.get(series_field, "?")) for r in records)
    summary = (
        f"Comparison chart from `{dataset_id}` — {len(series_vals)} series, "
        f"{len(records)} of {total} records."
    )
    return [summary, image]


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
