"""CCA rules tools for querying cloud compliance and optimization data."""

from pydantic import BaseModel, Field
from src.client import get_client


class GetCCASchemaInput(BaseModel):
    """Input for get_cca_schema tool."""

    pass


class QueryCCAInput(BaseModel):
    """Input for query_cca tool."""

    query: str = Field(description="Trino SQL query to execute on CCA rules data")


def register_cca_tools(mcp):
    """Register CCA rules tools."""

    @mcp.tool(
        "get_cca_schema",
        "Get the current schema for CCA rules table. "
        "Use this tool for general cloud optimization rules, idle resources, and compliance issues (e.g. stopped EC2, unused snapshots, idle endpoints). "
        "DO NOT use this for machine resize recommendations - use get_rightsizing_schema for that. "
        "Call this first to know the table name and available columns before querying.",
        GetCCASchemaInput,
    )
    async def get_cca_schema(args: GetCCASchemaInput) -> str:
        client_id = getattr(args, "_client_id", None)
        client_secret = getattr(args, "_client_secret", None)
        client = get_client(client_id, client_secret)
        schema = await client.get(f"/mcp/cca-checkers/schema")

        # Format schema for LLM
        result = f"Schema for cca-checkers:\\n\\n"
        result += f"Table name: {schema['table']['name']}\\n\\n"
        result += "Columns:\\n"
        for col in schema["columns"]:
            result += f"  - {col['name']} ({col['type']})\\n"

        if "meta" in schema:
            result += "\\nDescriptions:\\n"
            for meta in schema["meta"]:
                result += f"  - {meta['name']}: {meta['description']}\\n"

        result += "\\n⚠️ IMPORTANT QUERY RULES:\\n"
        result += "1. This dataset is daily partitioned. Always query the most recent data (yesterday, D-1).\\n"
        result += "   Example: WHERE year = 2024 AND month = 3 AND day = 5\\n"
        result += "2. All providers (aws, azure, gcp) are in the same table. Optionally filter by provider.\\n"
        result += "   Example: WHERE cca_provider = 'aws'\\n"
        result += "3. ALWAYS use double quotes for column names with special characters (/, -, etc).\\n"

        if "sample" in schema and schema["sample"]:
            result += "\\nSample data (first row):\\n"
            sample = schema["sample"][0]
            for key, value in sample.items():
                result += f"  - {key}: {value}\\n"

        return result

    @mcp.tool(
        "query_cca",
        "Execute a Trino SQL query on CCA rules data. "
        "Use this tool for general cloud optimization rules, idle resources, and compliance issues (e.g. stopped EC2, unused snapshots, idle endpoints). "
        "DO NOT use this for machine resize recommendations - use query_rightsizing for that. "
        "IMPORTANT: 1) Always filter for the most recent date (D-1) using partition columns (year, month, day). "
        "2) Optionally filter by provider (all providers are in the same table). "
        "3) Use double quotes for column names with special characters. "
        "Use get_cca_schema first to see available columns.",
        QueryCCAInput,
    )
    async def query_cca(args: QueryCCAInput) -> str:
        client_id = getattr(args, "_client_id", None)
        client_secret = getattr(args, "_client_secret", None)
        client = get_client(client_id, client_secret)
        result = await client.post(f"/mcp/cca-checkers/query", {"query": args.query})

        # Format result
        return str(result)
