"""Billing tools for querying cloud cost data from a single organization.

Billing vs Rebilling:
- BILLING: Cost data for a single organization, sourced directly from cloud providers
  (AWS CUR, Azure Cost APIs, GCP Billing Export). Use this for regular customers.
- REBILLING: Aggregated cost data across multiple customer organizations, exclusive
  to Pier Cloud partners. Use this only when the user is a partner managing multiple orgs.
"""

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


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


class GetBillingSchemaInput(BaseModel):
    connector_id: str = Field(description="Connector ID returned by get_billing_tables")


def register_billing_tools(mcp):

    @mcp.tool(
        "get_billing_tables",
        "List available billing connectors and their table names for a given provider. "
        "Use this for detailed, resource-level cloud cost data for a single organization, "
        "sourced directly from cloud providers (AWS CUR, Azure Cost APIs, GCP Billing Export). "
        "Billing data has high granularity: individual resources, tags, usage types, line items. "
        "Use this for questions like: cost by resource, cost by tag, service breakdown, usage details. "
        "⚠️ This is NOT rebilling - if the user is a Pier Cloud partner wanting a centralized view across multiple customer orgs, use get_rebilling_tables instead. "
        "Call this first to get the connector_id before fetching schema or querying.",
        GetBillingTablesInput,
    )
    async def get_billing_tables(args: GetBillingTablesInput) -> str:
        data = await get_client().get(f"/mcp/kinds/{args.provider}-billing/tables")
        connectors = data.get("connectors", [])
        if not connectors:
            return f"No billing connectors found for provider '{args.provider}'."
        result = f"Billing 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_billing_schema",
        "Get the schema (table name, columns, types, descriptions) for a billing connector. "
        "Use the connector_id returned by get_billing_tables. "
        "This schema reflects single-organization cost data from the cloud provider directly "
        "(e.g. AWS CUR columns, Azure cost fields, GCP billing export fields). "
        "⚠️ This is NOT rebilling schema - for partner aggregated data use get_rebilling_schema. "
        "Always use double quotes for column names with special characters (/, -, etc).",
        GetBillingSchemaInput,
    )
    async def get_billing_schema(args: GetBillingSchemaInput) -> 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
