Pipeline/backend/app/schemas/provider_credentials.py

"""Schemas for AI provider credential CRUD endpoints."""

from __future__ import annotations

from datetime import datetime
from uuid import UUID

from sqlmodel import SQLModel

RUNTIME_ANNOTATION_TYPES = (datetime, UUID)


class ProviderCredentialCreate(SQLModel):
    provider: str
    account_key: str
    display_name: str = ""
    api_key: str | None = None
    session_key: str | None = None   # Claude.ai / ChatGPT session token for subscription usage
    base_url: str | None = None
    active: bool = True


class ProviderCredentialTestRequest(SQLModel):
    provider: str
    account_key: str = "test"
    api_key: str | None = None
    session_key: str | None = None
    base_url: str | None = None


class ProviderCredentialUpdate(SQLModel):
    display_name: str | None = None
    api_key: str | None = None        # None = keep existing; "" = clear it
    session_key: str | None = None    # None = keep existing; "" = clear it
    base_url: str | None = None
    active: bool | None = None


class SubscriptionWindowRead(SQLModel):
    """One subscription-plan usage window from a session/OAuth token."""

    key: str           # "five_hour" | "seven_day" | "seven_day_sonnet" | etc.
    label: str         # "Current session" | "All models" | "Sonnet" | "3h" | "Week"
    pct_used: float    # 0–100
    reset_in_ms: int | None = None
    source: str = "provider_native"
    confidence: str = "high"


class TokenWindowRead(SQLModel):
    limit: int | None = None
    remaining: int | None = None
    used: int | None = None
    pct_used: float | None = None
    reset_at: str | None = None      # ISO 8601 UTC
    reset_in_ms: int | None = None   # ms until reset


class RequestWindowRead(SQLModel):
    limit: int | None = None
    remaining: int | None = None
    reset_at: str | None = None
    reset_in_ms: int | None = None


class ProviderUsageLiveRead(SQLModel):
    """Real-time token usage and rate limit data fetched directly from the provider API."""

    provider: str
    account_key: str
    checked_at: str           # ISO 8601 UTC
    reachable: bool
    source: str  # provider_native | provider_api_rate_limit | local_jsonl_estimate | configured_limit
    confidence: str  # high | medium | low
    error: str | None = None
    tokens: TokenWindowRead
    input_tokens: TokenWindowRead   # Anthropic input-only window
    output_tokens: TokenWindowRead  # Anthropic output-only window
    requests: RequestWindowRead
    # Subscription usage windows — populated when session_key is configured
    subscription_windows: list[SubscriptionWindowRead] = []
    subscription_plan: str | None = None  # e.g. "pro", "plus ($15.00)"
    models: list[str] = []
    sample_model: str | None = None
    sample_input_tokens: int | None = None
    sample_output_tokens: int | None = None
    sample_latency_ms: int | None = None
    # Optional debugging aid: exact rate-limit header names returned by provider.
    debug_rate_limit_headers: list[str] | None = None


class ProviderCredentialRead(SQLModel):
    """Safe read schema — api_key and session_key are never included."""

    id: UUID
    organization_id: UUID
    provider: str
    account_key: str
    display_name: str
    api_key_last_four: str | None
    has_api_key: bool
    session_key_last_four: str | None = None
    has_session_key: bool = False
    base_url: str | None
    active: bool
    created_at: datetime
    updated_at: datetime