app/routers/admin_auth.py

import hashlib
import hmac as hmac_mod
import re
import time
import uuid
from decimal import Decimal
from typing import Optional

from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from sqlalchemy import func, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from app.dependencies.admin_auth import require_admin_jwt
from app.models.database import (
    Conversation,
    CreditEvent,
    Tenant,
    TenantAuthProvider,
    User,
    get_db,
)
from app.services.google_auth import (
    decode_full_jwt,
    decode_partial_jwt,
    issue_full_jwt,
    issue_partial_jwt,
    lookup_tenant_for_google,
    register_tenant_for_google,
    verify_google_token,
)
from app.models.smart_models import KnowledgeBase
from app.services.totp_service import (
    consume_backup_code,
    encode_backup_codes,
    encrypt_secret,
    generate_backup_codes,
    generate_totp_secret,
    hash_backup_code,
    make_provisioning_uri,
    verify_totp_code,
)
from app.utils.api_key import generate_api_key
from app.utils.config import get_config
from app.utils.registration_rate_limit import enforce_registration_rate_limit

config = get_config()
router = APIRouter(prefix="/admin/auth", tags=["admin-auth"])

_TOTP_RE = re.compile(r"^\d{6}$")
_BACKUP_RE = re.compile(r"^[0-9a-f]{16}$")


class GoogleAuthRequest(BaseModel):
    id_token: str


class GoogleAuthResponse(BaseModel):
    totp_enabled: bool
    partial_token: Optional[str] = None
    registered: bool = False
    # B12 Layer 1: True when the Google identity is unknown and no tenant was
    # created — the consumer must ask the user to confirm before registering.
    new_account: bool = False


class TOTPSetupResponse(BaseModel):
    provisioning_uri: str
    backup_codes: list[str]  # plaintext — shown once, not retrievable again


class TOTPVerifyRequest(BaseModel):
    code: str


class TelegramAuthRequest(BaseModel):
    id: int
    first_name: str
    last_name: Optional[str] = None
    username: Optional[str] = None
    photo_url: Optional[str] = None
    auth_date: int
    hash: str
    phone_number: Optional[str] = None


def _verify_telegram_hmac(data: dict, bot_token: str) -> bool:
    """Verify Telegram Login Widget data using HMAC-SHA256."""
    received_hash = data.get("hash", "")
    check_pairs = sorted(
        (k, str(v)) for k, v in data.items() if k != "hash" and v is not None
    )
    data_check_string = "\n".join(f"{k}={v}" for k, v in check_pairs)
    secret_key = hashlib.sha256(bot_token.encode()).digest()
    expected_hash = hmac_mod.new(
        secret_key, data_check_string.encode(), hashlib.sha256
    ).hexdigest()
    return hmac_mod.compare_digest(expected_hash, received_hash)


def _set_access_cookie(response: Response, tenant_id: uuid.UUID) -> None:
    """Issue a full JWT and set it as an HttpOnly cookie on the response."""
    access_token = issue_full_jwt(tenant_id, config.jwt_secret_key)
    response.set_cookie(
        key="access_token",
        value=access_token,
        httponly=True,
        secure=config.env.name == "production",
        samesite="strict",
        path="/api",
        max_age=86400,
    )


def _resolve_enrollment_tenant_id(request: Request) -> tuple[uuid.UUID, bool]:
    """Resolve the tenant for a TOTP enrollment request via dual auth (B11).

    Two callers reach the enrollment endpoints:
    - Login-flow gate: a partial token (Authorization: Bearer, scope=totp_pending),
      issued by google_auth when 2FA is already enabled.
    - Settings enrollment: an already-logged-in tenant holding a full admin JWT in
      the HttpOnly access_token cookie (scope=admin) — JS cannot read it, so it can
      never be sent as a Bearer header.

    Returns (tenant_id, from_full_cookie). Partial token is tried first; on its
    absence/invalidity we fall back to the full cookie. Raises 401 if neither
    authenticates. Scope checks inside decode_* keep a full token from passing as a
    partial one and vice versa.
    """
    auth_header = request.headers.get("Authorization", "")
    if auth_header.startswith("Bearer "):
        token = auth_header.removeprefix("Bearer ")
        try:
            tenant_id_str = decode_partial_jwt(token, config.jwt_secret_key)
            return uuid.UUID(tenant_id_str), False
        except (ValueError, AttributeError):
            pass  # fall through to the full-cookie path

    cookie_token = request.cookies.get("access_token")
    if cookie_token:
        try:
            tenant_id_str = decode_full_jwt(cookie_token, config.jwt_secret_key)
            return uuid.UUID(tenant_id_str), True
        except (ValueError, AttributeError):
            pass

    raise HTTPException(status_code=401, detail="Unauthorized")


@router.post("/telegram")
async def telegram_auth(
    body: TelegramAuthRequest, request: Request, db: Session = Depends(get_db)
) -> JSONResponse:
    # 1. Guard: bot token must be configured
    if not config.telegram_bot_token:
        raise HTTPException(status_code=401, detail="Unauthorized")

    # 2. Verify HMAC
    data = body.model_dump(exclude_none=True)
    if not _verify_telegram_hmac(data, config.telegram_bot_token):
        raise HTTPException(status_code=401, detail="Unauthorized")

    # 3. Check auth_date freshness (24 hours)
    if time.time() - body.auth_date > 86400:
        raise HTTPException(status_code=401, detail="Unauthorized")

    telegram_user_id = str(body.id)
    phone = body.phone_number

    # 4. Direct match via auth_providers — telegram_id is the only lookup key.
    #    Phone matching was removed with the pre-provisioning model (B8): it
    #    carried a phone-recycling risk and contradicted self-registration.
    ap_row = db.scalars(
        select(TenantAuthProvider).where(
            TenantAuthProvider.provider == "telegram",
            TenantAuthProvider.provider_user_id == telegram_user_id,
        )
    ).first()

    registered = False
    if ap_row:
        tenant = db.get(Tenant, ap_row.tenant_id)
        if not tenant or not tenant.is_active:
            raise HTTPException(status_code=401, detail="Unauthorized")
    else:
        # 5. B8 self-registration: unknown Telegram identity creates a tenant.
        #    Telegram provides no email, so there is no cross-provider signal
        #    to collide on — collision detection reduces to the direct match.
        client_ip = request.client.host if request.client else "unknown"
        enforce_registration_rate_limit(client_ip)

        name = body.first_name + (f" {body.last_name}" if body.last_name else "")
        _, hashed_key = generate_api_key()
        tenant = Tenant(
            name=name,
            api_key=hashed_key,
            is_active=True,
            billing_mode="credits",
            balance_usd=Decimal("0"),
            phone=phone,
        )
        db.add(tenant)
        db.flush()
        db.add(TenantAuthProvider(
            tenant_id=tenant.id,
            provider="telegram",
            provider_user_id=telegram_user_id,
        ))
        try:
            db.commit()
            registered = True
        except IntegrityError:
            # B14: a concurrent request registered the same Telegram identity
            # between our lookup and insert. The (provider, provider_user_id)
            # unique constraint rejected our duplicate — recover by logging into
            # the tenant the winner created instead of creating a second one.
            db.rollback()
            ap_row = db.scalars(
                select(TenantAuthProvider).where(
                    TenantAuthProvider.provider == "telegram",
                    TenantAuthProvider.provider_user_id == telegram_user_id,
                )
            ).first()
            if ap_row is None:
                # Not a lost race — a genuine integrity failure. Don't swallow it.
                raise
            tenant = db.get(Tenant, ap_row.tenant_id)
            if not tenant or not tenant.is_active:
                raise HTTPException(status_code=401, detail="Unauthorized")

    # 6. Store phone on tenant if provided and not yet set
    if phone and not tenant.phone:
        tenant.phone = phone

    db.commit()

    # 7. Issue full JWT via cookie (set cookie on the returned JSONResponse, not the injected Response)
    result = JSONResponse(content={"ok": True, "registered": registered})
    _set_access_cookie(result, tenant.id)
    return result


def _check_google_collision(email: str, db: Session) -> None:
    """Raise 409 if a non-Google account already owns this email (B12 guard).

    A Telegram-only tenant whose owner_email matches must not be silently
    bound to a Google identity — surface the collision so the consumer routes
    the user to link instead.
    """
    pre_tenant = db.scalars(select(Tenant).where(Tenant.owner_email == email)).first()
    if pre_tenant and pre_tenant.google_id is None:
        existing_ap = db.scalars(
            select(TenantAuthProvider).where(
                TenantAuthProvider.tenant_id == pre_tenant.id,
                TenantAuthProvider.provider != "google",
            )
        ).first()
        if existing_ap:
            raise HTTPException(
                status_code=409,
                detail={"collision": True, "existing_provider": existing_ap.provider},
            )


def _finalize_google_login(
    tenant: Tenant, google_id: str, registered: bool, response: Response, db: Session
) -> GoogleAuthResponse:
    """Register the Google auth-provider row on first login, then issue the
    cookie (no 2FA) or a partial token (2FA enabled). Shared by login + register."""
    has_google_ap = db.scalars(
        select(TenantAuthProvider).where(
            TenantAuthProvider.tenant_id == tenant.id,
            TenantAuthProvider.provider == "google",
        )
    ).first()
    if not has_google_ap:
        db.add(TenantAuthProvider(
            tenant_id=tenant.id,
            provider="google",
            provider_user_id=google_id,
        ))
        db.commit()

    if not tenant.totp_enabled:
        _set_access_cookie(response, tenant.id)
        return GoogleAuthResponse(totp_enabled=False, registered=registered)

    partial_token = issue_partial_jwt(tenant.id, config.jwt_secret_key)
    return GoogleAuthResponse(
        totp_enabled=True, partial_token=partial_token, registered=registered
    )


@router.post("/google", response_model=GoogleAuthResponse)
async def google_auth(
    body: GoogleAuthRequest,
    request: Request,
    response: Response,
    db: Session = Depends(get_db),
) -> GoogleAuthResponse:
    try:
        google_payload = verify_google_token(body.id_token, config.google_client_id)
    except ValueError:
        raise HTTPException(status_code=401, detail="Unauthorized")

    email = google_payload["email"]
    google_id = google_payload["google_id"]

    _check_google_collision(email, db)

    tenant = lookup_tenant_for_google(email=email, google_id=google_id, db=db)

    # B12 Layer 1: unknown Google identity is NOT auto-registered. Signal a new
    # account (no cookie, no tenant created) so the consumer can require an
    # explicit "create workspace" confirmation before calling /google/register.
    if tenant is None:
        return GoogleAuthResponse(totp_enabled=False, registered=False, new_account=True)

    return _finalize_google_login(tenant, google_id, False, response, db)


@router.post("/google/register", response_model=GoogleAuthResponse)
async def google_register(
    body: GoogleAuthRequest,
    request: Request,
    response: Response,
    db: Session = Depends(get_db),
) -> GoogleAuthResponse:
    """Confirmed registration for a new Google identity (B12 Layer 1).

    Reached only after the user explicitly confirms "create a new workspace"
    in response to google_auth's new_account signal. Idempotent: if a tenant
    already exists for the identity (e.g. confirm clicked twice), it logs in
    rather than creating a duplicate."""
    try:
        google_payload = verify_google_token(body.id_token, config.google_client_id)
    except ValueError:
        raise HTTPException(status_code=401, detail="Unauthorized")

    email = google_payload["email"]
    google_id = google_payload["google_id"]
    profile_name = google_payload.get("name")

    _check_google_collision(email, db)

    tenant = lookup_tenant_for_google(email=email, google_id=google_id, db=db)
    registered = False
    if tenant is None:
        client_ip = request.client.host if request.client else "unknown"
        try:
            tenant = register_tenant_for_google(
                email=email, google_id=google_id, name=profile_name, db=db, client_ip=client_ip
            )
            registered = True
        except IntegrityError:
            # B13: a concurrent request registered this identity between our
            # lookup and insert (TOCTOU race). owner_email/google_id are unique,
            # so the DB rejected our duplicate — recover by logging into the
            # tenant the winner created instead of surfacing a 500.
            db.rollback()
            tenant = lookup_tenant_for_google(email=email, google_id=google_id, db=db)
            if tenant is None:
                # Not a lost race — a genuine integrity failure. Don't swallow it.
                raise

    return _finalize_google_login(tenant, google_id, registered, response, db)


@router.post("/link/telegram")
async def link_telegram(
    body: TelegramAuthRequest,
    tenant_id: uuid.UUID = Depends(require_admin_jwt),
    db: Session = Depends(get_db),
) -> JSONResponse:
    if not config.telegram_bot_token:
        raise HTTPException(status_code=401, detail="Unauthorized")

    data = body.model_dump(exclude_none=True)
    if not _verify_telegram_hmac(data, config.telegram_bot_token):
        raise HTTPException(status_code=401, detail="Unauthorized")

    if time.time() - body.auth_date > 86400:
        raise HTTPException(status_code=401, detail="Unauthorized")

    telegram_user_id = str(body.id)

    # Already linked to this tenant
    if db.scalars(
        select(TenantAuthProvider).where(
            TenantAuthProvider.tenant_id == tenant_id,
            TenantAuthProvider.provider == "telegram",
        )
    ).first():
        raise HTTPException(status_code=409, detail={"already_linked": True})

    # Telegram ID already used by another tenant
    if db.scalars(
        select(TenantAuthProvider).where(
            TenantAuthProvider.provider == "telegram",
            TenantAuthProvider.provider_user_id == telegram_user_id,
            TenantAuthProvider.tenant_id != tenant_id,
        )
    ).first():
        raise HTTPException(status_code=409, detail={"already_linked": True, "used_by_other": True})

    db.add(TenantAuthProvider(
        tenant_id=tenant_id,
        provider="telegram",
        provider_user_id=telegram_user_id,
    ))

    tenant = db.get(Tenant, tenant_id)
    if body.phone_number and tenant and not tenant.phone:
        tenant.phone = body.phone_number

    db.commit()
    return JSONResponse(content={"linked": True})


@router.post("/link/google")
async def link_google(
    body: GoogleAuthRequest,
    tenant_id: uuid.UUID = Depends(require_admin_jwt),
    db: Session = Depends(get_db),
) -> JSONResponse:
    try:
        google_payload = verify_google_token(body.id_token, config.google_client_id)
    except ValueError:
        raise HTTPException(status_code=401, detail="Unauthorized")

    google_id = google_payload["google_id"]
    email = google_payload["email"]

    # Already linked to this tenant
    if db.scalars(
        select(TenantAuthProvider).where(
            TenantAuthProvider.tenant_id == tenant_id,
            TenantAuthProvider.provider == "google",
        )
    ).first():
        raise HTTPException(status_code=409, detail={"already_linked": True})

    # Google ID already used by another tenant
    if db.scalars(
        select(TenantAuthProvider).where(
            TenantAuthProvider.provider == "google",
            TenantAuthProvider.provider_user_id == google_id,
            TenantAuthProvider.tenant_id != tenant_id,
        )
    ).first():
        raise HTTPException(status_code=409, detail={"already_linked": True, "used_by_other": True})

    db.add(TenantAuthProvider(
        tenant_id=tenant_id,
        provider="google",
        provider_user_id=google_id,
    ))

    tenant = db.get(Tenant, tenant_id)
    if tenant:
        if not tenant.owner_email:
            tenant.owner_email = email
        if not tenant.google_id:
            tenant.google_id = google_id

    db.commit()
    return JSONResponse(content={"linked": True})


@router.post("/totp/setup", response_model=TOTPSetupResponse)
@router.post("/totp/enroll", response_model=TOTPSetupResponse)
async def totp_setup(request: Request, db: Session = Depends(get_db)) -> TOTPSetupResponse:
    tenant_id, _ = _resolve_enrollment_tenant_id(request)

    tenant = db.get(Tenant, tenant_id)
    if not tenant:
        raise HTTPException(status_code=401, detail="Unauthorized")

    if tenant.totp_enabled:
        raise HTTPException(status_code=403, detail="TOTP already configured")

    raw_secret = generate_totp_secret()
    plain_codes = generate_backup_codes()

    tenant.totp_secret = encrypt_secret(raw_secret, config.fernet_key)
    tenant.backup_codes = encode_backup_codes([hash_backup_code(c) for c in plain_codes])
    tenant.totp_enabled = False
    db.commit()

    return TOTPSetupResponse(
        provisioning_uri=make_provisioning_uri(raw_secret, tenant.owner_email or ""),
        backup_codes=plain_codes,
    )


@router.post("/totp/verify")
async def totp_verify(body: TOTPVerifyRequest, request: Request, db: Session = Depends(get_db)) -> JSONResponse:
    tenant_id, from_full_cookie = _resolve_enrollment_tenant_id(request)

    tenant = db.get(Tenant, tenant_id)
    if not tenant or not tenant.totp_secret:
        raise HTTPException(status_code=401, detail="Unauthorized")

    code = body.code.strip().lower()
    verified = False

    if _TOTP_RE.match(code):
        verified = verify_totp_code(tenant.totp_secret, code, config.fernet_key)
    elif _BACKUP_RE.match(code) and tenant.backup_codes:
        updated = consume_backup_code(tenant.backup_codes, code)
        if updated is not None:
            tenant.backup_codes = encode_backup_codes(updated)
            verified = True

    if not verified:
        raise HTTPException(status_code=401, detail="Unauthorized")

    tenant.totp_enabled = True
    db.commit()

    # Settings enrollment (full-cookie path): the tenant already holds a valid
    # admin cookie — just confirm enrollment, do not re-issue it.
    if from_full_cookie:
        return JSONResponse(content={"ok": True})

    # Login-flow gate (partial-token path): issue the full admin cookie now.
    access_token = issue_full_jwt(tenant.id, config.jwt_secret_key)
    response = JSONResponse(content={"ok": True})
    # SameSite=Strict provides CSRF protection without a CSRF token
    response.set_cookie(
        key="access_token",
        value=access_token,
        httponly=True,
        secure=config.env.name == "production",
        samesite="strict",
        path="/api",
        max_age=86400,
    )
    return response


@router.delete("/totp")
async def totp_disable(
    body: TOTPVerifyRequest,
    tenant_id: uuid.UUID = Depends(require_admin_jwt),
    db: Session = Depends(get_db),
) -> JSONResponse:
    tenant = db.get(Tenant, tenant_id)
    if not tenant:
        raise HTTPException(status_code=401, detail="Unauthorized")

    if not tenant.totp_enabled or not tenant.totp_secret:
        raise HTTPException(status_code=400, detail="TOTP not configured")

    code = body.code.strip().lower()
    verified = False

    if _TOTP_RE.match(code):
        verified = verify_totp_code(tenant.totp_secret, code, config.fernet_key)
    elif _BACKUP_RE.match(code) and tenant.backup_codes:
        updated = consume_backup_code(tenant.backup_codes, code)
        if updated is not None:
            tenant.backup_codes = encode_backup_codes(updated)
            verified = True

    if not verified:
        raise HTTPException(status_code=401, detail="Unauthorized")

    tenant.totp_enabled = False
    tenant.totp_secret = None
    tenant.backup_codes = None
    db.commit()

    return JSONResponse(content={"totp_enabled": False})


@router.delete("/workspace")
async def delete_workspace(
    response: Response,
    tenant_id: uuid.UUID = Depends(require_admin_jwt),
    db: Session = Depends(get_db),
) -> JSONResponse:
    """B12 Layer 3: self-service deletion of an EMPTY workspace.

    Lets a user stranded with an accidental orphan account clear it themselves
    so they can link/reuse the identity. Only an empty workspace qualifies:
    zero balance, no KB entries, no conversations. A workspace with real data
    is refused (409) — bulk deletion of real content is out of scope here.

    The guard also refuses when chat users or credit-event history exist: those
    rows carry NOT-NULL tenant_id FKs with no cascade, so deleting past them
    would raise a database FK error. Refusing keeps deletion safe and matches
    the "truly empty" intent (a never-used orphan has none of these).
    """
    tenant = db.get(Tenant, tenant_id)
    if not tenant:
        raise HTTPException(status_code=401, detail="Unauthorized")

    def _count(model) -> int:
        return db.scalar(
            select(func.count()).select_from(model).where(model.tenant_id == tenant_id)
        )

    kb_count = _count(KnowledgeBase)
    conv_count = _count(Conversation)
    user_count = _count(User)
    credit_event_count = _count(CreditEvent)

    if (
        tenant.balance_usd != 0
        or kb_count > 0
        or conv_count > 0
        or user_count > 0
        or credit_event_count > 0
    ):
        raise HTTPException(
            status_code=409,
            detail={
                "not_empty": True,
                "balance_usd": str(tenant.balance_usd),
                "kb_entries": kb_count,
                "conversations": conv_count,
                "users": user_count,
                "credit_events": credit_event_count,
            },
        )

    db.delete(tenant)  # auth_providers cascade via delete-orphan
    db.commit()
    # Clear the now-stale session cookie for the deleted workspace.
    result = JSONResponse(content={"deleted": True})
    result.delete_cookie(key="access_token", path="/api")
    return result


@router.post("/logout")
async def logout(response: Response) -> dict:
    response.delete_cookie(key="access_token", path="/api")
    return {"ok": True}