"""
Shaped MCP server - exposes Shaped API as MCP tools and resources.
"""

import datetime
import hashlib
import json
import os
import time
from html.parser import HTMLParser
from typing import Any

import httpx
from shaped import Client
from shaped.autogen.models.engine_config_v2 import EngineConfigV2
from shaped.autogen.models.global_query_request import GlobalQueryRequest
from starlette.requests import Request
from starlette.responses import JSONResponse

from mcp.server.auth.provider import AccessToken, TokenVerifier
from mcp.server.auth.settings import AuthSettings
from mcp.server.fastmcp import FastMCP


class ShapedApiKeyVerifier(TokenVerifier):
    """Accepts Shaped API keys as Bearer tokens for HTTP transport."""

    async def verify_token(self, token: str) -> AccessToken | None:
        if not token or len(token) < 20:
            return None
        return AccessToken(
            token=token,
            client_id=hashlib.sha256(token.encode()).hexdigest()[:8],
            scopes=[],
            expires_at=int(time.time()) + 86400,
        )


mcp = FastMCP(
    "Shaped",
    instructions="MCP server for Shaped - create and manage tables, views, "
    "engines, and run retrieval queries.",
    host=os.environ.get("FASTMCP_HOST", "0.0.0.0"),
    port=int(os.environ.get("FASTMCP_PORT", "8080")),
    streamable_http_path="/",
    token_verifier=ShapedApiKeyVerifier(),
    auth=AuthSettings(
        issuer_url="https://api.shaped.ai",
        resource_server_url="https://api.shaped.ai",
        required_scopes=[],
    ),
)


@mcp.custom_route("/health", methods=["GET"])
async def health_check(request: Request) -> JSONResponse:
    """Health check for Kubernetes probes."""
    return JSONResponse({"status": "ok"})


def _get_client(api_key: str | None = None) -> Client:
    """Get Shaped client. API key from arg, auth context (HTTP), or env."""
    key = api_key
    if not key:
        try:
            from mcp.server.auth.middleware.auth_context import get_access_token

            access = get_access_token()
            if access:
                key = access.token
        except Exception:
            pass
    if not key:
        key = os.environ.get("SHAPED_API_KEY")
    if not key:
        raise ValueError(
            "SHAPED_API_KEY environment variable is required (STDIO). "
            "For HTTP, send Authorization: Bearer <api_key>."
        )
    base_url = os.environ.get("SHAPED_BASE_URL", "https://api.shaped.ai/v2")
    return Client(api_key=key, base_url=base_url)


def _to_json(obj: Any) -> str:
    """Serialize to JSON, handling non-serializable types."""

    def default(o: Any) -> Any:
        if isinstance(o, (datetime.datetime, datetime.date)):
            return o.isoformat()
        if isinstance(o, set):
            return list(o)
        if hasattr(o, "model_dump"):
            return o.model_dump(by_alias=True, exclude_none=True)
        if hasattr(o, "to_dict"):
            return o.to_dict()
        raise TypeError(f"Object of type {type(o)} is not JSON serializable")

    return json.dumps(obj, default=default, indent=2)


def _to_dict(obj: Any) -> Any:
    """Convert an SDK object to a plain dict/list, reusing the _to_json serialiser."""
    return json.loads(_to_json(obj))


def _unwrap_view_details(result: Any) -> Any:
    """Unwrap get_view anyOf responses to plain dicts for JSON."""
    if isinstance(result, dict):
        return result
    to_dict = getattr(result, "to_dict", None)
    if callable(to_dict):
        return to_dict()
    inner = getattr(result, "actual_instance", None)
    if inner is not None:
        return inner
    return result


# ---------------------------------------------------------------------------
# Query tools
# ---------------------------------------------------------------------------


@mcp.tool()
def shaped_query(
    engine_name: str,
    query: str,
    parameters: str = "{}",
    return_metadata: bool = False,
    return_explanation: bool = False,
    return_journey_explanations: bool = False,
) -> str:
    """
    Execute a ShapedQL or declarative query against an engine.
    query: SQL string or JSON query config.
    parameters: JSON object of query parameters.
    return_explanation: include per-stage pipeline breakdown
      (retrieve/filter/score/reorder) with timing and input/output counts
      — useful for debugging query behaviour.
    return_journey_explanations: include per-entity journey through the pipeline.
    """
    try:
        client = _get_client()
        params = json.loads(parameters) if parameters else {}
        result = client.execute_query(
            engine_name=engine_name,
            query=query,
            parameters=params or None,
            return_metadata=return_metadata,
            return_explanation=return_explanation,
            return_journey_explanations=return_journey_explanations,
        )
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_list_saved_queries(engine_name: str) -> str:
    """List all saved queries defined on an engine."""
    try:
        client = _get_client()
        result = client.list_saved_queries(engine_name)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_execute_saved_query(
    engine_name: str,
    query_name: str,
    parameters: str = "{}",
    return_metadata: bool = False,
    return_explanation: bool = False,
    return_journey_explanations: bool = False,
) -> str:
    """
    Execute a saved query by name.
    parameters: JSON object mapping param names to values (e.g. {"user_id": "u123"}).
    Use shaped_get_saved_query to discover parameter names first.
    return_explanation: include per-stage pipeline breakdown
      (retrieve/filter/score/reorder) with timing and input/output counts
      — useful for debugging query behaviour.
    return_journey_explanations: include per-entity journey through the pipeline.
    """
    try:
        client = _get_client()
        params = json.loads(parameters) if parameters else {}
        result = client.execute_saved_query(
            engine_name=engine_name,
            query_name=query_name,
            parameters=params or None,
            return_metadata=return_metadata,
            return_explanation=return_explanation,
            return_journey_explanations=return_journey_explanations,
        )
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_get_saved_query(engine_name: str, query_name: str) -> str:
    """Get details for a saved query including its parameter names."""
    try:
        client = _get_client()
        result = client.get_saved_query_info(engine_name, query_name)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_global_query(query: str, parameters: str = "{}") -> str:
    """
    Execute a query against the Shaped data catalog or a Shaped engine.
    Two namespace prefixes are supported, each with different syntax:

    ClickHouse SQL — use for data catalog namespaces:
      data.{table_name}               — query a table
      data.{view_name}                — query a view
      data.{engine_name}.items        — engine item data (ClickHouse)
      data.{engine_name}.users        — engine user data (ClickHouse)
      data.{engine_name}.interactions — engine interaction data (ClickHouse)
    Example: SELECT * FROM data.my_engine.items LIMIT 10

    ShapedQL — use for engine namespaces (items only):
      engine.{engine_name}.items             — scan item data
      engine.{engine_name}.retrieve(...)     — retrieval queries
                                               (similarity, text_search, etc.)
    Example: SELECT * FROM engine.my_engine.items LIMIT 10
    Example: SELECT * FROM engine.my_engine.retrieve(
               text_search(query='$q', mode='vector')) LIMIT 10

    parameters: JSON object of $variable substitutions.
    """
    try:
        client = _get_client()

        params = json.loads(parameters) if parameters else {}
        # shaped.Client has no public global-query method; access internals directly.
        result = client._query_api.post_global_query_query_post(
            x_api_key=client._api_key,
            global_query_request=GlobalQueryRequest(
                query=query, parameters=params or None
            ),
        )
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


# ---------------------------------------------------------------------------
# Table tools
# ---------------------------------------------------------------------------


@mcp.tool()
def shaped_create_table(table_config: str) -> str:
    """
    Create a table. table_config: JSON object with schema_type, name, column_schema.
    """
    try:
        client = _get_client()
        config = json.loads(table_config)
        result = client.create_table(config)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_update_table(table_config: str) -> str:
    """
    Update a table's schema or configuration.
    table_config: JSON with full table config including name.
    """
    try:
        client = _get_client()
        config = json.loads(table_config)
        result = client.update_table(config)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_list_tables() -> str:
    """List all tables in the Shaped account."""
    try:
        client = _get_client()
        result = client.list_tables()
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_insert_rows(table_name: str, rows: str) -> str:
    """
    Insert rows into a table. rows: JSON array of row objects.
    """
    try:
        client = _get_client()
        data = json.loads(rows)
        result = client.insert_table_rows(table_name, data)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_get_table(table_name: str) -> str:
    """Get details for a table by name."""
    try:
        client = _get_client()
        result = client._table_api.get_table_details_tables_table_name_get(
            table_name=table_name,
            x_api_key=client._api_key,
        )
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_delete_table(table_name: str) -> str:
    """Delete a table by name."""
    try:
        client = _get_client()
        result = client.delete_table(table_name)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


# ---------------------------------------------------------------------------
# View tools
# ---------------------------------------------------------------------------


@mcp.tool()
def shaped_create_view(view_config: str) -> str:
    """
    Create a view. view_config: JSON with type (sql_view_config or
    ai_enrichment_view_config), name, and config fields.
    """
    try:
        client = _get_client()
        config = json.loads(view_config)
        result = client.create_view(config)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_update_view(view_config: str) -> str:
    """
    Update a view's configuration.
    view_config: JSON with full view config including name.
    """
    try:
        client = _get_client()
        config = json.loads(view_config)
        result = client.update_view(config)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_list_views() -> str:
    """List all views in the Shaped account."""
    try:
        client = _get_client()
        result = client.list_views()
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_get_view(view_name: str) -> str:
    """Get details for a view by name."""
    try:
        client = _get_client()
        result = client.get_view(view_name)
        payload = _unwrap_view_details(result)
        if payload is None:
            return json.dumps({"error": "View not found", "detail": "Empty response"})
        return _to_json(payload)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_delete_view(view_name: str) -> str:
    """Delete a view by name."""
    try:
        client = _get_client()
        result = client.delete_view(view_name)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


# ---------------------------------------------------------------------------
# Engine tools
# ---------------------------------------------------------------------------


@mcp.tool()
def shaped_create_engine(engine_config: str) -> str:
    """
    Create an engine from a JSON config.
    Required fields: name, version.
    Optional: data (dataset sources), index (embeddings, lexical search),
              training, deployment, queries, description, tags.
    """
    try:
        client = _get_client()
        config = json.loads(engine_config)
        result = client.create_engine(EngineConfigV2.from_dict(config))
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_update_engine(engine_config: str) -> str:
    """Update an existing engine. engine_config: JSON with full config."""
    try:
        client = _get_client()
        config = json.loads(engine_config)
        result = client.update_engine(EngineConfigV2.from_dict(config))
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_list_engines() -> str:
    """List all engines in the Shaped account."""
    try:
        client = _get_client()
        result = client.list_engines()
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_delete_engine(engine_name: str) -> str:
    """Delete an engine by name."""
    try:
        client = _get_client()
        result = client.delete_engine(engine_name)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_get_engine(engine_name: str) -> str:
    """Get full engine details and config by name."""
    try:
        client = _get_client()
        result = client.get_engine(engine_name)
        return _to_json(result)
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_validate_engine_config(engine_config: str) -> str:
    """
    Validate an engine config JSON offline without making an API call.
    Returns the parsed/normalised config on success, or a structured error
    on failure — useful for catching schema mistakes before create/update.
    engine_config: JSON object matching the EngineConfigV2 schema.
    """
    try:
        config = json.loads(engine_config)
        parsed = EngineConfigV2.from_dict(config)
        return json.dumps({"valid": True, "parsed": _to_dict(parsed)})
    except Exception as e:
        return json.dumps({"valid": False, "error": str(e), "detail": type(e).__name__})


# ---------------------------------------------------------------------------
# Docs tool
# ---------------------------------------------------------------------------

_DOCS_BASE_URL = "https://docs.shaped.ai/docs"
_DOCS_SLUG_MAP = {
    "shapedql": "v2/query_reference/shapedql",
    "api/v2": "api/v2",
}
_docs_cache: dict[str, tuple[float, str]] = {}  # url -> (fetched_at, content)
_DOCS_CACHE_TTL = 3600


def _extract_text_from_html(html_text: str) -> str:
    """Strip HTML tags, scripts, and nav elements; return plain text."""

    class _TextExtractor(HTMLParser):
        def __init__(self) -> None:
            super().__init__()
            self._parts: list[str] = []
            self._skip_depth = 0
            self._skip_tags = {"script", "style", "nav", "header", "footer"}

        def handle_starttag(self, tag: str, attrs: list) -> None:
            if tag in self._skip_tags:
                self._skip_depth += 1

        def handle_endtag(self, tag: str) -> None:
            if tag in self._skip_tags and self._skip_depth > 0:
                self._skip_depth -= 1

        def handle_data(self, data: str) -> None:
            if self._skip_depth == 0:
                stripped = data.strip()
                if stripped:
                    self._parts.append(stripped)

        def get_text(self) -> str:
            return "\n".join(self._parts)

    extractor = _TextExtractor()
    extractor.feed(html_text)
    return extractor.get_text()


def _fetch_docs_page(page: str) -> str:
    """Fetch a docs.shaped.ai page and return its text content."""
    path = _DOCS_SLUG_MAP.get(page, page)
    url = f"{_DOCS_BASE_URL}/{path.lstrip('/')}"
    if not url.startswith("https://docs.shaped.ai/"):
        return json.dumps(
            {
                "error": (
                    f"Page '{page}' resolves outside docs.shaped.ai"
                    " — only docs.shaped.ai pages are allowed."
                )
            }
        )
    cached = _docs_cache.get(url)
    if cached and time.time() - cached[0] < _DOCS_CACHE_TTL:
        return cached[1]
    try:
        response = httpx.get(url, follow_redirects=True, timeout=15)
        response.raise_for_status()
        content_type = response.headers.get("content-type", "")
        if "text/html" in content_type:
            content = _extract_text_from_html(response.text)
        else:
            content = response.text
        _docs_cache[url] = (time.time(), content)
        return content
    except httpx.HTTPStatusError as e:
        return json.dumps({"error": f"HTTP {e.response.status_code}", "url": url})
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.tool()
def shaped_fetch_docs(page: str) -> str:
    """
    Fetch a Shaped documentation page and return its text content.
    page: a short slug ('shapedql', 'api/v2') or any sub-path under docs.shaped.ai.
    Use this mid-conversation when you need to look up ShapedQL syntax or API details
    and the docs resource has not been pre-attached.
    """
    return _fetch_docs_page(page)


# ---------------------------------------------------------------------------
# Resources
# ---------------------------------------------------------------------------


def _sample(client: Any, namespace: str, limit: int = 5) -> Any:
    """Run a SELECT * sample query and return the parsed result, or None on failure."""
    try:
        return client._query_api.post_global_query_query_post(
            x_api_key=client._api_key,
            global_query_request=GlobalQueryRequest(
                query=f"SELECT * FROM {namespace} LIMIT {limit}",
            ),
        )
    except Exception:
        return None


@mcp.resource("shaped://account")
def account_resource() -> str:
    """Account overview — all engines, tables, and views."""
    try:
        client = _get_client()
        return json.dumps(
            {
                "engines": _to_dict(client.list_engines()),
                "tables": _to_dict(client.list_tables()),
                "views": _to_dict(client.list_views()),
            }
        )
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.resource("shaped://engines/{engine_name}")
def engine_resource(engine_name: str) -> str:
    """
    Full config, schema, and status for a specific engine,
    with data samples from items/users/interactions.
    """
    try:
        client = _get_client()
        engine = client.get_engine(engine_name)
        return json.dumps(
            {
                "engine": _to_dict(engine),
                "sample": {
                    "items": _to_dict(_sample(client, f"data.{engine_name}.items")),
                    "users": _to_dict(_sample(client, f"data.{engine_name}.users")),
                    "interactions": _to_dict(
                        _sample(client, f"data.{engine_name}.interactions")
                    ),
                },
            }
        )
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.resource("shaped://tables/{table_name}")
def table_resource(table_name: str) -> str:
    """Schema, status, and a 5-row data sample for a specific table."""
    try:
        client = _get_client()
        table = client._table_api.get_table_details_tables_table_name_get(
            table_name=table_name,
            x_api_key=client._api_key,
        )
        return json.dumps(
            {
                "table": _to_dict(table),
                "sample": _to_dict(_sample(client, f"data.{table_name}")),
            }
        )
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.resource("shaped://views/{view_name}")
def view_resource(view_name: str) -> str:
    """Schema, status, and a 5-row data sample for a specific view."""
    try:
        client = _get_client()
        view = client.get_view(view_name)
        view_payload = _unwrap_view_details(view)
        if view_payload is None:
            return json.dumps({"error": "View not found", "detail": "Empty response"})
        return json.dumps(
            {
                "view": _to_dict(view_payload),
                "sample": _to_dict(_sample(client, f"data.{view_name}")),
            }
        )
    except Exception as e:
        return json.dumps({"error": str(e), "detail": type(e).__name__})


@mcp.resource("shaped://docs/shapedql")
def docs_shapedql_resource() -> str:
    """
    ShapedQL reference documentation — attach before a query-building
    session for full syntax coverage.
    """
    return _fetch_docs_page("shapedql")


# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------


@mcp.prompt(
    name="shaped-build-query",
    title="Build a Query",
    description=(
        "Design a ShapedQL query for a given use case by inspecting"
        " the engine config and choosing the right retrieval strategy."
    ),
)
def shaped_build_query_prompt(
    engine_name: str,
    use_case: str,
    user_id: str = "",
) -> str:
    """Prompt template for designing a ShapedQL query from a use case description."""
    personalization = (
        f"\nuser_id '{user_id}' has been provided — consider personalized ranking "
        "using a similarity + score pattern if the engine has a suitable model."
        if user_id
        else "\nNo user_id provided — focus on content-based or text search strategies."
    )
    return (
        f"Design a ShapedQL query for engine '{engine_name}'.\n"
        f"Use case: {use_case}{personalization}\n\n"
        "Steps:\n"
        "1. Gather context:\n"
        "   - ShapedQL docs: if shaped://docs/shapedql is already attached"
        " as context, use it for syntax.\n"
        "     Otherwise call shaped_fetch_docs('shapedql') to retrieve"
        " the reference, or rely on the\n"
        "     inline syntax guide in step 2.\n"
        "   - Engine config: skip if the engine resource is already"
        " attached as context.\n"
        "     Otherwise call shaped_get_engine to fetch it.\n"
        "   - Note the embedding names under config.index.embeddings.\n"
        "   - Note any model names under config.training or config.queries.\n"
        "   - Check shaped_list_saved_queries in case a named query"
        " already covers this use case.\n\n"
        "2. Choose the retrieval strategy that best fits the use case"
        " and available config:\n"
        "   - Text search    → FROM retrieve(text_search(query='$q', mode='vector'))\n"
        "   - Similar items  → FROM retrieve(similarity("
        "embedding_ref='<name>', encoder='precomputed_item',"
        " input_item_id='$id'))\n"
        "   - Personalized   → FROM retrieve(similarity("
        "embedding_ref='<name>', encoder='precomputed_user',"
        " input_user_id='$uid'))\n"
        "                        with SELECT score(expression='<model>')"
        " AS s ... ORDER BY s\n"
        "   - Rerank candidates → FROM retrieve(ids($candidate_ids))"
        " with score(expression='<model>')\n\n"
        "3. Consider WHERE clause filters:\n"
        "   - Standard SQL filters: WHERE price > 100 AND category = 'electronics'\n"
        "   - prebuilt('filter_name', input_user_id='$uid') — apply a"
        " pre-configured filter (e.g. exclude already-purchased items)\n"
        "   - truncate(n) — cap the candidate set before scoring"
        " to improve performance\n"
        "   Example: WHERE category = 'electronics'"
        " AND prebuilt('exclude_purchased', input_user_id='$uid')"
        " AND truncate(50)\n\n"
        "4. Consider REORDER clause for post-ranking adjustments:\n"
        "   - diversity(score=s, strength=0.3, diversity_lookback_window=30)"
        " — reduce consecutive similarity for variety\n"
        "   - exploration(score=s, retriever='explore', strength=0.3)"
        " — inject items from an exploration retriever\n"
        "   - boosted(score=s, retriever='boosted_items', strength=0.3)"
        " — inject items weighted by a boost column\n"
        "   Example: REORDER diversity(score=s, strength=0.3)\n\n"
        "5. Write the ShapedQL query, substituting real embedding/model"
        " names from step 1.\n\n"
        "6. Show the final query and the shaped_query call needed to run it, "
        "including the parameters JSON object."
    )


@mcp.prompt(
    name="shaped-engine-data-summary",
    title="Engine Data Summary",
    description=(
        "Sample and summarize the items, users, and interactions"
        " tables for an engine."
    ),
)
def shaped_engine_data_summary_prompt(engine_name: str, limit: int = 5) -> str:
    """Prompt template for summarising engine data."""
    return (
        f"Summarise the data in engine '{engine_name}'"
        " by sampling each of its three data tables.\n\n"
        f"Run the following three shaped_global_query calls"
        f" (limit={limit} each) using ClickHouse SQL syntax:\n\n"
        f"  SELECT * FROM data.{engine_name}.items LIMIT {limit}\n"
        f"  SELECT * FROM data.{engine_name}.users LIMIT {limit}\n"
        f"  SELECT * FROM data.{engine_name}.interactions LIMIT {limit}\n\n"
        "For each table report:\n"
        "- The column names and their types (inferred from the sample).\n"
        "- A brief description of what the data represents.\n"
        "- Any obvious data quality observations (nulls, unexpected values, etc.).\n\n"
        "Finish with a one-paragraph overview of how items, users,"
        " and interactions relate to each other."
    )


@mcp.prompt(
    name="shaped-query-table",
    title="Query a Table",
    description=(
        "Query a Shaped table — accepts a plain-English description"
        " (e.g. 'show me 100 blue items') or a SQL query to validate and run."
    ),
)
def shaped_query_table_prompt(table_name: str, query: str = "") -> str:
    """Prompt template for querying a table from a description or SQL query."""
    if not query:
        intent_section = (
            f"No query was provided. Write a sensible default:"
            f" SELECT * FROM data.{table_name} LIMIT 5"
        )
    else:
        intent_section = (
            f"The user supplied the following input:\n  {query}\n\n"
            "Determine whether this is:\n"
            "  a) A plain-English description — translate it into"
            " a ClickHouse SQL query.\n"
            "  b) A SQL query — check it for ClickHouse syntax errors"
            " or unsupported syntax and fix any issues before running.\n"
            "Either way, call shaped_get_table first to confirm"
            " the column schema, then produce a valid query."
        )
    return (
        f"Query the table '{table_name}' using shaped_global_query.\n\n"
        f"Namespace: data.{table_name}\n\n"
        f"{intent_section}\n\n"
        "Rules:\n"
        "- Use ClickHouse SQL syntax.\n"
        "- Pass $variable placeholders in the query string and supply"
        " values in the parameters argument.\n"
        "- Show the final SQL before executing it."
    )


@mcp.prompt(
    name="shaped-query-view",
    title="Query a View",
    description=(
        "Query a Shaped view — accepts a plain-English description"
        " (e.g. 'show me 100 blue items') or a SQL query to validate and run."
    ),
)
def shaped_query_view_prompt(view_name: str, query: str = "") -> str:
    """Prompt template for querying a view from a description or SQL query."""
    if not query:
        intent_section = (
            f"No query was provided. Write a sensible default:"
            f" SELECT * FROM data.{view_name} LIMIT 5"
        )
    else:
        intent_section = (
            f"The user supplied the following input:\n  {query}\n\n"
            "Determine whether this is:\n"
            "  a) A plain-English description — translate it into"
            " a ClickHouse SQL query.\n"
            "  b) A SQL query — check it for ClickHouse syntax errors"
            " or unsupported syntax and fix any issues before running.\n"
            "Either way, call shaped_get_view first to confirm"
            " the column schema, then produce a valid query."
        )
    return (
        f"Query the view '{view_name}' using shaped_global_query.\n\n"
        f"Namespace: data.{view_name}\n\n"
        f"{intent_section}\n\n"
        "Rules:\n"
        "- Use ClickHouse SQL syntax.\n"
        "- Pass $variable placeholders in the query string and supply"
        " values in the parameters argument.\n"
        "- Show the final SQL before executing it."
    )


def run(
    transport: str = "stdio", host: str | None = None, port: int | None = None
) -> None:
    """Run the MCP server.

    Args:
        transport: MCP transport type: `stdio` | `streamable-http`.
        host: Host to bind for HTTP transports (when applicable).
        port: Port to bind for HTTP transports (when applicable).
    """
    # FastMCP captures host/port in its internal settings at init time. Since the
    # CLI passes `host`/`port` as explicit args, update the live settings so
    # `mcp.run()` binds correctly.
    if host is not None:
        os.environ["FASTMCP_HOST"] = host
        if hasattr(mcp, "settings"):
            mcp.settings.host = host
    if port is not None:
        os.environ["FASTMCP_PORT"] = str(port)
        if hasattr(mcp, "settings"):
            mcp.settings.port = int(port)

    mcp.run(transport=transport)
