kite/backend/api/admin/ai.py

from __future__ import annotations

from typing import Optional

from fastapi import APIRouter, Depends, HTTPException, Request, status
from pydantic import BaseModel, ConfigDict, field_validator
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from ai import get_provider
from ai.provider_config import ProviderConfig, PROVIDER_DEFAULTS
from db.models import SystemSettings, User
from deps.auth import get_current_admin
from deps.db import get_db
from deps.utils import get_client_ip
from services.ai_config import encrypt_api_key, load_provider_config_by_id, validate_provider_id
from services.audit import write_audit_log

router = APIRouter()  # NO prefix — parent __init__.py carries /api/admin (D-04)


# ── Local helper (only ai.py uses this — not in shared.py) ───────────────────

def _ai_config_to_dict(row: SystemSettings) -> dict:
    """Return a safe subset of SystemSettings fields — explicitly excludes api_key_enc.

    has_api_key is the ONLY indicator that a key is stored (T-07-01 mitigated).
    The raw encrypted value and any decrypted plaintext are NEVER returned.
    """
    return {
        "provider_id": row.provider_id,
        "base_url": row.base_url,
        "model_name": row.model_name,
        "context_chars": row.context_chars,
        "is_active": row.is_active,
        "has_api_key": row.api_key_enc is not None,
        "updated_at": row.updated_at.isoformat() if row.updated_at else None,
    }


# ── Request models ────────────────────────────────────────────────────────────

class SystemAiConfigUpdate(BaseModel):
    """Request model for PUT /api/admin/ai-config (system-level provider configuration).

    Security: extra="forbid" prevents mass-assignment of unexpected fields (T-07-13).
    provider_id is validated against PROVIDER_DEFAULTS keys (T-07-13).
    api_key is write-only: when None the existing api_key_enc is left untouched,
    when "" the api_key_enc is cleared, when a non-empty string it is encrypted.
    """

    model_config = ConfigDict(extra="forbid")

    provider_id: str
    api_key: Optional[str] = None
    base_url: Optional[str] = None
    model_name: Optional[str] = None
    context_chars: Optional[int] = None
    is_active: Optional[bool] = None

    @field_validator("provider_id")
    @classmethod
    def provider_must_be_known(cls, v: str) -> str:
        return validate_provider_id(v)


class TestConnectionRequest(BaseModel):
    """Request body for POST /api/admin/ai-config/test-connection.

    Unsaved form values (api_key, base_url, model_name) override the DB row so
    admins can verify credentials before saving.  All override fields are optional;
    omitting them falls back to whatever is stored in system_settings.
    """

    model_config = ConfigDict(extra="forbid")

    provider_id: str
    api_key: Optional[str] = None      # If non-empty, used instead of stored api_key_enc
    base_url: Optional[str] = None     # If non-None, overrides DB base_url
    model_name: Optional[str] = None   # If non-empty, overrides DB model_name

    @field_validator("provider_id")
    @classmethod
    def provider_must_be_known(cls, v: str) -> str:
        return validate_provider_id(v)


# ── Endpoints ─────────────────────────────────────────────────────────────────

@router.get("/ai-config/models")
async def get_ai_config_models(
    provider_id: str,
    session: AsyncSession = Depends(get_db),
    _admin: User = Depends(get_current_admin),
) -> dict:
    """Return the list of model IDs available from a provider's API (D-08).

    Calls the provider's standard GET /models endpoint using the stored
    config (base_url + api_key from system_settings).  Always returns 200
    with {"models": [...]} — never 5xx on provider failure (returns empty list).

    Security: requires get_current_admin; provider_id from query param only;
    decrypted api_key never appears in the response.
    """
    import httpx  # noqa: PLC0415 — local import keeps startup fast

    config = await load_provider_config_by_id(session, provider_id)

    # Resolve base_url: prefer DB row, fall back to PROVIDER_DEFAULTS
    if config and config.base_url:
        base_url = config.base_url.rstrip("/")
    else:
        base_url = (PROVIDER_DEFAULTS.get(provider_id, {}).get("base_url") or "").rstrip("/")

    if not base_url:
        return {"models": [], "provider_id": provider_id}

    api_key = config.api_key if config else ""

    # Anthropic uses x-api-key; all others use Bearer (different auth schemes)
    if provider_id == "anthropic":
        headers = {
            "x-api-key": api_key,
            "anthropic-version": "2023-06-01",
        }
        models_url = "https://api.anthropic.com/v1/models"
    else:
        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
        models_url = f"{base_url}/models"

    try:
        async with httpx.AsyncClient(timeout=8.0) as client:
            resp = await client.get(models_url, headers=headers)
            resp.raise_for_status()
            data = resp.json()

        raw_list = data.get("data") or data.get("models") or []
        model_ids: list[str] = sorted(
            {
                item["id"] if isinstance(item, dict) else str(item)
                for item in raw_list
                if item
            }
        )
        return {"models": model_ids, "provider_id": provider_id}
    except Exception as exc:
        return {"models": [], "provider_id": provider_id, "error": str(exc)[:120]}


@router.post("/ai-config/test-connection")
async def test_ai_connection(
    body: TestConnectionRequest,
    session: AsyncSession = Depends(get_db),
    _admin: User = Depends(get_current_admin),
) -> dict:
    """Test connectivity for an AI provider, optionally with unsaved form values (D-08).

    Loads the stored system_settings row for body.provider_id, then overlays any
    non-empty values from the request body so admins can verify credentials before
    saving them to the database.

    Override priority (highest -> lowest):
      1. body.api_key / base_url / model_name  (unsaved form values)
      2. system_settings DB row               (previously saved config)
      3. PROVIDER_DEFAULTS                    (built-in fallback)

    Returns {"ok": true/false, "provider_id": str} — never raises 5xx for
    provider-side failures; surfaces as ok=False so the UI shows a clear status.

    Security: requires get_current_admin; api_key from body is used only for the
    in-flight health_check() call and is never stored or logged.
    """
    provider_id = body.provider_id
    stored = await load_provider_config_by_id(session, provider_id)
    defaults = PROVIDER_DEFAULTS.get(provider_id, {})

    # Resolve effective values: body overrides DB, DB overrides PROVIDER_DEFAULTS
    effective_api_key = (
        body.api_key
        if body.api_key
        else (stored.api_key if stored else "")
    )
    effective_base_url = (
        body.base_url
        if body.base_url is not None
        else (stored.base_url if stored else defaults.get("base_url"))
    )
    effective_model = (
        body.model_name
        if body.model_name
        else (stored.model if stored else defaults.get("model", ""))
    )

    effective_config = ProviderConfig(
        provider_id=provider_id,
        api_key=effective_api_key,
        base_url=effective_base_url,
        model=effective_model,
    )

    try:
        provider = get_provider(effective_config)
        ok = await provider.health_check()
        return {"ok": ok, "provider_id": provider_id}
    except Exception:
        return {"ok": False, "provider_id": provider_id, "reason": "health_check_failed"}


@router.get("/ai-config")
async def get_ai_config(
    session: AsyncSession = Depends(get_db),
    _admin: User = Depends(get_current_admin),
) -> dict:
    """Return all AI provider configurations for the admin panel (D-08).

    Includes DB rows for providers that have been saved AND synthesised stubs
    for providers that only exist in PROVIDER_DEFAULTS (so the admin UI always
    shows all 10 providers even before any have been configured).

    Security invariant: api_key_enc is NEVER returned (T-07-01).
    Use has_api_key (bool) as the only indicator that a key is stored.
    """
    result = await session.execute(select(SystemSettings))
    db_rows = result.scalars().all()

    # Build a lookup for DB rows
    db_by_provider: dict[str, SystemSettings] = {r.provider_id: r for r in db_rows}

    providers_out = []
    for pid in PROVIDER_DEFAULTS:
        if pid in db_by_provider:
            providers_out.append(_ai_config_to_dict(db_by_provider[pid]))
        else:
            # Synthesise a stub entry for providers with no DB row yet
            defaults = PROVIDER_DEFAULTS[pid]
            providers_out.append({
                "provider_id": pid,
                "base_url": defaults.get("base_url"),
                "model_name": defaults.get("model", ""),
                "context_chars": defaults.get("context_chars", 8000),
                "is_active": False,
                "has_api_key": False,
                "updated_at": None,
            })

    return {"providers": providers_out}


@router.put("/ai-config")
async def update_system_ai_config(
    body: SystemAiConfigUpdate,
    request: Request,
    session: AsyncSession = Depends(get_db),
    _admin: User = Depends(get_current_admin),
) -> dict:
    """Create or update a system-level AI provider configuration (D-08, D-15).

    Upsert semantics: if no row exists for body.provider_id, one is created using
    PROVIDER_DEFAULTS for any omitted fields.

    API key handling (T-07-01 mitigated):
      - body.api_key is None -> leave existing api_key_enc untouched
      - body.api_key == "" -> clear api_key_enc (set to NULL)
      - body.api_key is a non-empty string -> HKDF-encrypt and store

    is_active=True handling (T-07-03 mitigated):
      When body.is_active is True, a single atomic UPDATE flips all rows:
      SET is_active = (provider_id = :target_id)
      This guarantees COUNT(WHERE is_active) == 1 with no read-then-write race.

    Audit log (T-07-14 mitigated):
      metadata_ contains only provider_id + fields_changed list — never the
      api_key value itself.
    """
    from config import settings as _settings  # noqa: PLC0415

    stmt = select(SystemSettings).where(SystemSettings.provider_id == body.provider_id)
    result = await session.execute(stmt)
    row = result.scalar_one_or_none()

    is_new = row is None
    if is_new:
        defaults = PROVIDER_DEFAULTS[body.provider_id]
        row = SystemSettings(
            provider_id=body.provider_id,
            model_name=defaults.get("model", ""),
            context_chars=defaults.get("context_chars", 8000),
            base_url=defaults.get("base_url"),
            is_active=False,
            api_key_enc=None,
        )

    fields_changed: list[str] = []

    if body.api_key is not None:
        fields_changed.append("api_key")
        if body.api_key == "":
            row.api_key_enc = None
        else:
            master_key_str = _settings.cloud_creds_key
            master_key_bytes = (
                master_key_str.encode("utf-8")
                if isinstance(master_key_str, str)
                else master_key_str
            )
            row.api_key_enc = encrypt_api_key(master_key_bytes, body.provider_id, body.api_key)

    if body.base_url is not None:
        row.base_url = body.base_url
        fields_changed.append("base_url")

    if body.model_name is not None:
        row.model_name = body.model_name
        fields_changed.append("model_name")

    if body.context_chars is not None:
        row.context_chars = body.context_chars
        fields_changed.append("context_chars")

    if body.is_active is not None:
        fields_changed.append("is_active")

    if is_new:
        session.add(row)
        await session.flush()  # ensure row has an id before UPDATE

    # Atomic is_active flip: SET is_active = (provider_id = :target) on ALL rows.
    # Single UPDATE statement prevents dual-active race condition (T-07-03).
    if body.is_active is True:
        await session.execute(
            update(SystemSettings).values(
                is_active=(SystemSettings.provider_id == body.provider_id)
            )
        )
        # Reflect the flip on the in-memory row
        row.is_active = True

    _ip_addr = get_client_ip(request)
    await write_audit_log(
        session,
        event_type="admin.ai_config_changed",
        user_id=None,
        actor_id=_admin.id,
        resource_id=None,
        ip_address=_ip_addr,
        metadata_={"provider_id": body.provider_id, "fields_changed": fields_changed},
    )

    await session.commit()

    await session.refresh(row)

    return _ai_config_to_dict(row)