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

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


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


class GetRightsizingSchemaInput(BaseModel):
    connector_id: str = Field(description="Connector ID returned by get_rightsizing_tables")


def register_rightsizing_tools(mcp):

    @mcp.tool(
        "get_rightsizing_tables",
        "List available rightsizing connectors and their table names for a given provider. "
        "Use this 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 - use get_cca_tables for that. "
        "Only available for 'aws' and 'azure'. "
        "Call this first to get the connector_id before fetching schema or querying.",
        GetRightsizingTablesInput,
    )
    async def get_rightsizing_tables(args: GetRightsizingTablesInput) -> str:
        if args.provider not in ("aws", "azure"):
            return "Error: rightsizing is only available for 'aws' and 'azure'."
        data = await get_client().get(f"/mcp/kinds/cca-rightsizing-{args.provider}/tables")
        connectors = data.get("connectors", [])
        if not connectors:
            return f"No rightsizing connectors found for provider '{args.provider}'."
        result = f"Rightsizing connectors for '{args.provider}':\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_rightsizing_schema",
        "Get the schema (table name, columns, types, descriptions) for a rightsizing connector. "
        "Use the connector_id returned by get_rightsizing_tables. "
        "Use this for machine/instance resize recommendations and rightsizing. "
        "DO NOT use this for general cloud optimization rules - use get_cca_schema for that. "
        "Only available for 'aws' and 'azure'.",
        GetRightsizingSchemaInput,
    )
    async def get_rightsizing_schema(args: GetRightsizingSchemaInput) -> 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. Always use double quotes for column names with special characters (/, -, etc).\n"

        return result
