smart-org/apps/api/app/interfaces/http/routers/agy.py

"""agy account management — Google OAuth headless login for the AI provider.

The agy provider has no API keys: an *account* is a Google identity. This
router implements the headless paste-URL login loop for the platform owner:

1. ``POST /admin/agy/accounts/login/start`` — returns an ``auth_url`` the owner
   opens in a browser and completes the Google consent screen.
2. Google redirects to the loopback callback URL. Nothing is listening there
   on purpose — the parameters (``code``, ``state``) are already in the address
   bar, and the owner pastes the whole URL back.
3. ``POST /admin/agy/accounts/login/complete`` — exchanges the code against
   Google (PKCE proof stays server-side), fetches the account email and
   Antigravity project id, and stores the refresh token **encrypted at rest**.

Every response is deliberately free of credentials: the console sees email,
status and cooldown, never ciphertext. Same deny-by-default guard as the rest
of the admin console.
"""

from __future__ import annotations

import json
from datetime import datetime

from fastapi import APIRouter, Depends, Request, status
from pydantic import BaseModel, ConfigDict, Field
from sqlalchemy.ext.asyncio import AsyncSession

from app.application.services.audit_service import record_audit
from app.core.crypto import encrypt
from app.core.database import get_session
from app.core.errors import NotFoundError
from app.domain.ai_control.models import AiModel, ModelCredential, ProviderType
from app.domain.audit.models import AuditAction
from app.domain.identity.models import User
from app.infrastructure.repositories.ai_control import (
    ModelCredentialRepository,
    ProviderRepository,
)
from app.interfaces.http.deps import client_meta, require_platform_owner

router = APIRouter(
    prefix="/admin/agy",
    tags=["admin-agy"],
    # Deny-by-default, same rule as the admin console: an endpoint added later
    # is owner-only without its author having to remember.
    dependencies=[Depends(require_platform_owner)],
)

#: Name of the singleton agy provider row. The seed creates it; the login flow
#: creates it on demand so a bare deployment is never stuck.
AGY_PROVIDER_NAME = "agy"

_PROJECTED_OUT_FIELDS = (
    "id",
    "provider_id",
    "name",
    "email",
    "project_id",
    "status",
    "cooldown_until",
    "is_active",
    "last_error",
    "last_used_at",
    "created_at",
)


# ---------------------------------------------------------------------------
# Schemas — note there is no field anywhere capable of holding a token.
# ---------------------------------------------------------------------------


class AgyLoginStartOut(BaseModel):
    login_id: str
    auth_url: str


class AgyLoginCompleteIn(BaseModel):
    login_id: str = Field(min_length=8, max_length=64)
    callback_url: str = Field(min_length=20, max_length=2048)


class AgyAccountOut(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: str
    provider_id: str
    name: str
    email: str
    project_id: str
    status: str
    cooldown_until: datetime | None
    is_active: bool
    last_error: str
    last_used_at: datetime | None
    created_at: datetime


class AgyTestOut(BaseModel):
    ok: bool
    detail: str
    model_internal_name: str
    input_tokens: int
    output_tokens: int


# ---------------------------------------------------------------------------
# Provider helpers
# ---------------------------------------------------------------------------


async def _ensure_agy_provider(repo: ProviderRepository):
    provider = await repo.get_by_name(AGY_PROVIDER_NAME)
    if provider is not None:
        return provider
    from app.domain.ai_control.models import Provider

    return await repo.create(Provider(name=AGY_PROVIDER_NAME, type=ProviderType.AGY))


def _account_out(credential: ModelCredential) -> AgyAccountOut:
    return AgyAccountOut(
        id=credential.id,
        provider_id=credential.provider_id,
        name=credential.name,
        email=credential.email,
        project_id=credential.project_id,
        status=credential.status,
        cooldown_until=credential.cooldown_until,
        is_active=credential.is_active,
        last_error=credential.last_error,
        last_used_at=credential.last_used_at,
        created_at=credential.created_at,
    )


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------


@router.post(
    "/accounts/login/start", response_model=AgyLoginStartOut,
    status_code=status.HTTP_201_CREATED,
)
async def start_agy_login() -> AgyLoginStartOut:
    """Open a headless login and hand back the Google URL to open."""
    from app.infrastructure.adapters.agy import begin_login

    login_id, auth_url = await begin_login()
    return AgyLoginStartOut(login_id=login_id, auth_url=auth_url)


@router.post(
    "/accounts/login/complete", response_model=AgyAccountOut,
    status_code=status.HTTP_201_CREATED,
)
async def complete_agy_login(
    payload: AgyLoginCompleteIn,
    request: Request,
    owner: User = Depends(require_platform_owner),
    session: AsyncSession = Depends(get_session),
) -> AgyAccountOut:
    """Redeem a started login with the callback URL the owner pasted."""
    from app.core.errors import ValidationError
    from app.infrastructure.adapters.agy import redeem_login

    try:
        result = await redeem_login(payload.login_id, payload.callback_url)
    except Exception as exc:
        # Login failures are user-fixable input problems, never 500s.
        raise ValidationError(str(exc)) from exc
    if not result.email:
        raise ValidationError("Google نشانی ایمیل اکانت را برنگرداند.")

    providers = ProviderRepository(session)
    credentials = ModelCredentialRepository(session)
    provider = await _ensure_agy_provider(providers)

    envelope = json.dumps(
        {
            "refresh_token": result.refresh_token,
            "project_id": result.project_id,
            "email": result.email,
        }
    )
    existing = await credentials.find_by_email(provider.id, result.email)
    if existing is not None:
        # Same Google account re-login: rotate the key, keep the row.
        credential = await credentials.replace_material(
            existing.id, encrypted_api_key=encrypt(envelope),
            project_id=result.project_id,
        )
    else:
        credential = await credentials.create(
            ModelCredential(
                provider_id=provider.id,
                name=result.email,
                encrypted_api_key=encrypt(envelope),
                email=result.email,
                project_id=result.project_id,
            )
        )
    assert credential is not None  # narrowed for type checkers; create/replace return it

    ip, ua = client_meta(request)
    await record_audit(
        session, action=AuditAction.PROVIDER_CHANGE, actor_user_id=owner.id,
        target_type="agy_account", target_id=credential.id,
        metadata={"email": result.email, "provider_id": provider.id},
        ip=ip, user_agent=ua,
    )
    return _account_out(credential)


@router.get("/accounts", response_model=list[AgyAccountOut])
async def list_agy_accounts(
    session: AsyncSession = Depends(get_session),
) -> list[AgyAccountOut]:
    providers = ProviderRepository(session)
    provider = await providers.get_by_name(AGY_PROVIDER_NAME)
    if provider is None:
        return []
    rows = await ModelCredentialRepository(session).list_for_provider(provider.id)
    return [_account_out(c) for c in rows]


@router.post("/accounts/{credential_id}/deactivate", response_model=AgyAccountOut)
async def deactivate_agy_account(
    credential_id: str,
    request: Request,
    owner: User = Depends(require_platform_owner),
    session: AsyncSession = Depends(get_session),
) -> AgyAccountOut:
    return await _set_account_active(
        credential_id, active=False, request=request, owner=owner, session=session
    )


@router.post("/accounts/{credential_id}/activate", response_model=AgyAccountOut)
async def activate_agy_account(
    credential_id: str,
    request: Request,
    owner: User = Depends(require_platform_owner),
    session: AsyncSession = Depends(get_session),
) -> AgyAccountOut:
    return await _set_account_active(
        credential_id, active=True, request=request, owner=owner, session=session
    )


async def _set_account_active(
    credential_id: str, *, active: bool, request: Request, owner: User,
    session: AsyncSession,
) -> AgyAccountOut:
    repo = ModelCredentialRepository(session)
    credential = await repo.get(credential_id)
    if credential is None:
        raise NotFoundError("اکانت agy یافت نشد.")
    await repo.set_active(credential_id, active)
    credential = await repo.get(credential_id)
    assert credential is not None
    ip, ua = client_meta(request)
    await record_audit(
        session, action=AuditAction.PROVIDER_CHANGE, actor_user_id=owner.id,
        target_type="agy_account", target_id=credential_id,
        metadata={"active": active},
        ip=ip, user_agent=ua,
    )
    return _account_out(credential)


@router.delete("/accounts/{credential_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_agy_account(
    credential_id: str,
    request: Request,
    owner: User = Depends(require_platform_owner),
    session: AsyncSession = Depends(get_session),
) -> None:
    repo = ModelCredentialRepository(session)
    credential = await repo.get(credential_id)
    if credential is None:
        raise NotFoundError("اکانت agy یافت نشد.")
    await repo.delete(credential_id)
    ip, ua = client_meta(request)
    await record_audit(
        session, action=AuditAction.PROVIDER_CHANGE, actor_user_id=owner.id,
        target_type="agy_account", target_id=credential_id,
        metadata={"deleted": True, "email": credential.email},
        ip=ip, user_agent=ua,
    )


@router.post("/accounts/{credential_id}/test", response_model=AgyTestOut)
async def test_agy_account(
    credential_id: str,
    session: AsyncSession = Depends(get_session),
) -> AgyTestOut:
    """Run a one-shot completion through exactly this account."""
    from app.application.ports import ChatMessage
    from app.infrastructure.adapters.agy import AgyModelAdapter

    repo = ModelCredentialRepository(session)
    credential = await repo.get(credential_id)
    if credential is None:
        raise NotFoundError("اکانت agy یافت نشد.")

    adapter = AgyModelAdapter()  # stateless for a single forced account
    model = AiModel(
        provider_id=credential.provider_id,
        internal_name="gemini-3-flash",
        display_name_for_admin="تست اکانت",
    )
    try:
        response = await adapter.complete(
            model=model,
            credentials=[credential],
            messages=[ChatMessage(role="user", content="سلام، یک کلمه پاسخ بده.")],
            max_tokens=64,
        )
    except Exception as exc:  # noqa: BLE001 - the console wants the reason, not a 500
        return AgyTestOut(
            ok=False, detail=str(exc)[:300], model_internal_name="",
            input_tokens=0, output_tokens=0,
        )
    return AgyTestOut(
        ok=True, detail=response.content[:300],
        model_internal_name=response.model_internal_name,
        input_tokens=response.input_tokens, output_tokens=response.output_tokens,
    )


@router.get("/models", response_model=list[str])
async def list_agy_models() -> list[str]:
    """Model ids the logged-in agy subscription exposes (registry helper)."""
    from app.infrastructure.adapters.registry import get_model_provider

    provider = get_model_provider()
    lister = getattr(provider, "available_models", None)
    if lister is None:
        return []
    return await lister()