"""CCA rightsizing tools for querying machine rightsizing recommendations."""

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


class GetRightsizingSchemaInput(BaseModel):
    provider: str = Field(description="Cloud provider: 'aws' or 'azure'")


class QueryRightsizingInput(BaseModel):
    provider: str = Field(description="Cloud provider: 'aws' or 'azure'")
    query: str = Field(description="Trino SQL query to execute on CCA rightsizing data")


def register_rightsizing_tools(mcp):
    """Register CCA rightsizing tools."""

    @mcp.tool(
        "get_rightsizing_schema",
        "Get the current schema for CCA rightsizing table. "
        "Use this tool when the user asks about machine/instance resize recommendations, rightsizing, or changing instance types to save costs. "
        "DO NOT use this for general cloud optimization rules or idle resources - use get_cca_schema for that. "
        "Only available for 'aws' and 'azure'. "
        "Call this first to know the table name and available columns before querying.",
        GetRightsizingSchemaInput,
    )
    async def get_rightsizing_schema(args: GetRightsizingSchemaInput) -> str:
        if args.provider not in ("aws", "azure"):
            return "Error: rightsizing is only available for 'aws' and 'azure'."

        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-rightsizing-{args.provider}/schema")

        result = f"Schema for cca-rightsizing-{args.provider}:\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. 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"
            for key, value in schema["sample"][0].items():
                result += f"  - {key}: {value}\n"

        return result

    @mcp.tool(
        "query_rightsizing",
        "Execute a Trino SQL query on CCA rightsizing data. "
        "Use this tool when the user asks about machine/instance resize recommendations, rightsizing, or changing instance types to save costs. "
        "DO NOT use this for general cloud optimization rules or idle resources - use query_cca for that. "
        "Only available for 'aws' and 'azure'. "
        "IMPORTANT: 1) Always filter for the most recent data (D-1) using partition columns (year, month, day). "
        "2) Use double quotes for column names with special characters. "
        "Use get_rightsizing_schema first to see available columns.",
        QueryRightsizingInput,
    )
    async def query_rightsizing(args: QueryRightsizingInput) -> str:
        if args.provider not in ("aws", "azure"):
            return "Error: rightsizing is only available for 'aws' and 'azure'."

        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-rightsizing-{args.provider}/query", {"query": args.query}
        )
        return str(result)
