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

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


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


class GetRebillingSchemaInput(BaseModel):
    connector_id: str = Field(
        description="Connector ID returned by get_rebilling_tables"
    )


def register_rebilling_tools(mcp):

    @mcp.tool(
        "get_rebilling_tables",
        "List available rebilling connectors for a given provider. "
        "Rebilling is EXCLUSIVE to Pier Cloud partners - use this when the user wants a centralized cost overview across ALL their managed customer organizations. "
        "Rebilling data is aggregated by day (no resource-level detail). "
        "Use this for questions like: total cost per customer, cost trends across organizations, partner-level spend overview. "
        "When the user asks for a 'general account', 'overall costs', 'all customers', or any consolidated/aggregated view, try rebilling FIRST before billing. "
        "⚠️ This is NOT regular billing - if the user wants detailed/resource-level cost data for a single org, use get_billing_tables instead. "
        "If this returns empty, the user is not a Pier Cloud partner - fall back to get_billing_tables. "
        "Call this first to get the connector_id before fetching schema or querying. "
        "Use get_rebilling_schema after this to know available columns.",
        GetRebillingTablesInput,
    )
    async def get_rebilling_tables(args: GetRebillingTablesInput) -> str:
        data = await get_client().get(
            f"/mcp/connectors-kinds/{args.provider}-rebilling"
        )
        connectors = data.get("connectors", [])
        if not connectors:
            return (
                f"No rebilling connectors found for provider '{args.provider}'. "
                "This user is NOT a Pier Cloud partner. "
                "Use get_billing_tables instead to query single-organization cost data."
            )
        result = f"Rebilling 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_rebilling_schema",
        "Get the schema (table name, columns, types, descriptions) for a rebilling connector. "
        "Use the connector_id returned by get_rebilling_tables. "
        "Rebilling schema contains aggregated cost data across multiple customer organizations, "
        "available only to Pier Cloud partners. "
        "⚠️ This is NOT billing schema - for single-org cost data use get_billing_schema. "
        "After getting the schema, call query_connector with BOTH connector_id AND query parameters. "
        "Always use double quotes for column names with special characters (/, -, etc).",
        GetRebillingSchemaInput,
    )
    async def get_rebilling_schema(args: GetRebillingSchemaInput) -> 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⚠️ Always use double quotes for column names with special characters (/, -, etc).\n"

        return result
