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

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


class GetCCATablesInput(BaseModel):
    pass


class GetCCASchemaInput(BaseModel):
    connector_id: str = Field(description="Connector ID returned by get_cca_tables")


def register_cca_tools(mcp):

    @mcp.tool(
        "get_cca_tables",
        "List available CCA rules connectors and their table names. "
        "Use this 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_tables for that. "
        "Call this first to get the connector_id before fetching schema or querying.",
        GetCCATablesInput,
    )
    async def get_cca_tables(args: GetCCATablesInput) -> str:
        data = await get_client().get("/mcp/kinds/cca-checkers/tables")
        connectors = data.get("connectors", [])
        if not connectors:
            return "No CCA connectors found."
        result = "CCA connectors:\n"
        for c in connectors:
            result += f"  - ID: {c['connector_id']} | Name: {c['connector_name']} | Table: {c['table_name']}\n"
        return result

    @mcp.tool(
        "get_cca_schema",
        "Get the schema (table name, columns, types, descriptions) for a CCA rules connector. "
        "Use the connector_id returned by get_cca_tables. "
        "Use this for general cloud optimization rules, idle resources, and compliance issues. "
        "DO NOT use this for machine resize recommendations - use get_rightsizing_schema for that. "
        "IMPORTANT: 1) Dataset is daily partitioned - always filter by most recent date (D-1). "
        "2) All providers (aws, azure, gcp) are in the same table. "
        "3) Always use double quotes for column names with special characters.",
        GetCCASchemaInput,
    )
    async def get_cca_schema(args: GetCCASchemaInput) -> str:
        schema = await get_client().get(f"/mcp/connectors/{args.connector_id}/schema")

        result = f"Table: {schema['table']['name']}\n\nColumns:\n"
        for col in schema["columns"]:
            result += f"  - {col['name']} ({col['type']})\n"

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

        result += "\n⚠️ IMPORTANT QUERY RULES:\n"
        result += "1. Always filter for the most recent date (D-1) using partition columns (year, month, day).\n"
        result += "2. All providers (aws, azure, gcp) are in the same table. Filter with cca_provider if needed.\n"
        result += "3. Always use double quotes for column names with special characters (/, -, etc).\n"

        return result
