"""Phase 3: billing fields, tenant_llm_keys, stripe_events, audit index, agent budget

Revision ID: 005
Revises: 004
Create Date: 2026-03-24

This migration adds:

1. Billing columns on tenants table:
   - stripe_customer_id, stripe_subscription_id, stripe_subscription_item_id
   - subscription_status (TEXT, default 'none')
   - trial_ends_at (TIMESTAMPTZ, nullable)
   - agent_quota (INTEGER, default 0)

2. Budget column on agents table:
   - budget_limit_usd (FLOAT, nullable) — monthly spend cap per agent

3. tenant_llm_keys table:
   - Stores encrypted BYO API keys per tenant per provider
   - RLS enabled (same FORCE ROW LEVEL SECURITY pattern as agents)
   - UNIQUE(tenant_id, provider) constraint
   - key_hint column (VARCHAR(4)) for safe portal display without decryption
   - konstruct_app granted SELECT, INSERT, DELETE (no UPDATE — keys are immutable)

4. stripe_events table:
   - Idempotency guard for Stripe webhook event processing
   - Simple TEXT primary key (stripe event_id)
   - konstruct_app granted SELECT, INSERT (no UPDATE, no DELETE)

5. Composite index on audit_events:
   - idx_audit_events_tenant_type_created ON audit_events(tenant_id, action_type, created_at DESC)
   - Supports cost aggregation queries in usage.py endpoints

Design:
  - tenant_llm_keys intentionally grants DELETE (portal operators can remove keys)
  - audit_events immutability is NOT weakened — only a new covering index is added
  - stripe_events does not need RLS — it's platform-wide idempotency, not tenant-scoped
"""

from __future__ import annotations

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects.postgresql import UUID


# revision identifiers, used by Alembic.
revision: str = "005"
down_revision: Union[str, None] = "004"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # =========================================================================
    # 1. Billing columns on tenants
    # =========================================================================
    op.add_column("tenants", sa.Column(
        "stripe_customer_id",
        sa.String(255),
        nullable=True,
        comment="Stripe Customer ID (cus_...)",
    ))
    op.add_column("tenants", sa.Column(
        "stripe_subscription_id",
        sa.String(255),
        nullable=True,
        comment="Stripe Subscription ID (sub_...)",
    ))
    op.add_column("tenants", sa.Column(
        "stripe_subscription_item_id",
        sa.String(255),
        nullable=True,
        comment="Stripe Subscription Item ID (si_...) for quantity updates",
    ))
    op.add_column("tenants", sa.Column(
        "subscription_status",
        sa.String(50),
        nullable=False,
        server_default="none",
        comment="none | trialing | active | past_due | canceled | unpaid",
    ))
    op.add_column("tenants", sa.Column(
        "trial_ends_at",
        sa.DateTime(timezone=True),
        nullable=True,
        comment="Trial expiry timestamp (NULL for non-trial subscriptions)",
    ))
    op.add_column("tenants", sa.Column(
        "agent_quota",
        sa.Integer,
        nullable=False,
        server_default="0",
        comment="Number of active agents allowed under current subscription",
    ))

    # =========================================================================
    # 2. Budget column on agents
    # =========================================================================
    op.add_column("agents", sa.Column(
        "budget_limit_usd",
        sa.Float,
        nullable=True,
        comment="Monthly spend cap in USD. NULL means no limit.",
    ))

    # =========================================================================
    # 3. tenant_llm_keys — encrypted BYO API keys
    # =========================================================================
    op.create_table(
        "tenant_llm_keys",
        sa.Column(
            "id",
            UUID(as_uuid=True),
            primary_key=True,
            server_default=sa.text("gen_random_uuid()"),
        ),
        sa.Column(
            "tenant_id",
            UUID(as_uuid=True),
            sa.ForeignKey("tenants.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column(
            "provider",
            sa.Text,
            nullable=False,
            comment="LLM provider: openai | anthropic | cohere | groq | etc.",
        ),
        sa.Column(
            "label",
            sa.Text,
            nullable=False,
            comment="Human-readable label for the portal",
        ),
        sa.Column(
            "encrypted_key",
            sa.Text,
            nullable=False,
            comment="Fernet-encrypted API key ciphertext",
        ),
        sa.Column(
            "key_hint",
            sa.String(4),
            nullable=True,
            comment="Last 4 chars of plaintext key for portal display",
        ),
        sa.Column(
            "key_version",
            sa.Integer,
            nullable=False,
            server_default="1",
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("NOW()"),
        ),
        sa.UniqueConstraint("tenant_id", "provider", name="uq_tenant_llm_key_provider"),
    )

    op.create_index("ix_tenant_llm_keys_tenant", "tenant_llm_keys", ["tenant_id"])

    # RLS: only rows matching current tenant are visible
    op.execute("ALTER TABLE tenant_llm_keys ENABLE ROW LEVEL SECURITY")
    op.execute("ALTER TABLE tenant_llm_keys FORCE ROW LEVEL SECURITY")
    op.execute("""
        CREATE POLICY tenant_isolation ON tenant_llm_keys
        USING (tenant_id = current_setting('app.current_tenant', TRUE)::uuid)
    """)

    # SELECT, INSERT, DELETE — no UPDATE (keys are immutable; rotate by delete+insert)
    op.execute("GRANT SELECT, INSERT, DELETE ON tenant_llm_keys TO konstruct_app")

    # =========================================================================
    # 4. stripe_events — webhook idempotency guard
    # =========================================================================
    op.create_table(
        "stripe_events",
        sa.Column(
            "event_id",
            sa.Text,
            primary_key=True,
            comment="Stripe event ID — globally unique per Stripe account",
        ),
        sa.Column(
            "processed_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("NOW()"),
        ),
    )

    # No RLS — platform-wide idempotency table, not tenant-scoped
    op.execute("GRANT SELECT, INSERT ON stripe_events TO konstruct_app")

    # =========================================================================
    # 5. Composite index on audit_events for usage aggregation queries
    # =========================================================================
    # Covers: WHERE tenant_id = X AND action_type = 'llm_call' AND created_at >= ...
    # Used by: usage.py endpoints for per-agent and per-provider cost aggregation
    op.create_index(
        "idx_audit_events_tenant_type_created",
        "audit_events",
        ["tenant_id", "action_type", "created_at"],
        postgresql_ops={"created_at": "DESC"},
    )


def downgrade() -> None:
    # Remove composite index on audit_events
    op.drop_index("idx_audit_events_tenant_type_created", table_name="audit_events")

    # Remove stripe_events
    op.execute("REVOKE ALL ON stripe_events FROM konstruct_app")
    op.drop_table("stripe_events")

    # Remove tenant_llm_keys
    op.execute("REVOKE ALL ON tenant_llm_keys FROM konstruct_app")
    op.drop_table("tenant_llm_keys")

    # Remove budget column from agents
    op.drop_column("agents", "budget_limit_usd")

    # Remove billing columns from tenants
    op.drop_column("tenants", "agent_quota")
    op.drop_column("tenants", "trial_ends_at")
    op.drop_column("tenants", "subscription_status")
    op.drop_column("tenants", "stripe_subscription_item_id")
    op.drop_column("tenants", "stripe_subscription_id")
    op.drop_column("tenants", "stripe_customer_id")
