"""Rebilling tools for querying cloud cost data."""

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


class GetSchemaInput(BaseModel):
    """Input for get_rebilling_schema tool."""

    provider: str = Field(description="Cloud provider: 'aws', 'azure', or 'gcp'")


class QueryRebillingInput(BaseModel):
    """Input for query_rebilling tool."""

    provider: str = Field(description="Cloud provider: 'aws', 'azure', or 'gcp'")
    query: str = Field(description="Trino SQL query to execute on rebilling data")


def register_rebilling_tools(mcp):
    """Register rebilling tools."""

    @mcp.tool(
        "get_rebilling_schema",
        "Get the current schema (table name, columns, types, descriptions) for rebilling table. "
        "Call this first to know the table name and available columns before querying.",
        GetSchemaInput,
    )
    async def get_rebilling_schema(args: GetSchemaInput) -> str:
        client = get_client()
        schema = await client.get(f"/mcp/{args.provider}-rebilling/schema")

        # Format schema for LLM
        result = f"Schema for {args.provider}-rebilling:\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"

        result += "\nDescriptions:\n"
        for meta in schema["meta"]:
            result += f"  - {meta['name']}: {meta['description']}\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_rebilling",
        "Execute a Trino SQL query on rebilling data. "
        "Use get_rebilling_schema first to see available columns.",
        QueryRebillingInput,
    )
    async def query_rebilling(args: QueryRebillingInput) -> str:
        client = get_client()
        result = await client.post(
            f"/mcp/{args.provider}-rebilling/query", {"query": args.query}
        )

        # Format result
        return str(result)
