app/services/google_auth.py

import uuid
from datetime import datetime, timezone, timedelta
from decimal import Decimal

from fastapi import HTTPException
from google.oauth2 import id_token as google_id_token
from google.auth.transport import requests as google_requests
import jwt
from sqlalchemy import select
from sqlalchemy.orm import Session

from app.models.database import Tenant
from app.utils.api_key import generate_api_key
from app.utils.registration_rate_limit import enforce_registration_rate_limit

PARTIAL_TOKEN_SCOPE = "totp_pending"
PARTIAL_TOKEN_EXPIRE_MINUTES = 5

FULL_TOKEN_SCOPE = "admin"
FULL_TOKEN_EXPIRE_HOURS = 24

SUPER_ADMIN_TOKEN_SCOPE = "super_admin"
SUPER_ADMIN_TOKEN_EXPIRE_HOURS = 1


def verify_google_token(token: str, google_client_id: str) -> dict:
    """Verify a Google ID token; return {"email", "google_id", "name"}.

    Raises ValueError when the token is invalid OR its email is not
    verified by Google — `verify_oauth2_token` checks signature/audience/
    expiry only, so an unverified email would otherwise be able to
    register a tenant squatting someone else's address.
    """
    payload = google_id_token.verify_oauth2_token(
        token,
        google_requests.Request(),
        google_client_id,
    )
    if payload.get("email_verified") is not True:
        raise ValueError("Google email not verified")
    return {
        # Lowercased at the boundary so every consumer (login lookup,
        # collision pre-check, linking, registration store) compares and
        # stores one canonical form (B9.1).
        "email": payload["email"].lower(),
        "google_id": payload["sub"],
        "name": payload.get("name"),
    }


def lookup_tenant_for_google(
    email: str,
    google_id: str,
    db: Session,
) -> Tenant | None:
    """Find the tenant for a verified Google login WITHOUT registering one.

    Returns the matched tenant (binding google_id on a provisioned/gifted
    account that has none yet), or None when no tenant exists for either the
    email or the google_id. The caller decides what to do with None — B12
    Layer 1 requires explicit confirmation before a new tenant is created, so
    google_auth must NOT auto-create here. The caller must run collision
    detection BEFORE calling this function.
    """
    email = email.lower()  # defensive — canonical form regardless of caller
    tenant = db.scalars(select(Tenant).where(Tenant.owner_email == email)).first()

    if tenant is None:
        # The Google account's email may have changed since registration —
        # google_id is the stable identifier (and unique on tenants, so a
        # blind create would violate the constraint).
        tenant = db.scalars(
            select(Tenant).where(Tenant.google_id == google_id)
        ).first()

    if tenant is not None:
        if tenant.google_id is not None and tenant.google_id != google_id:
            raise HTTPException(status_code=401, detail="Unauthorized")
        if tenant.google_id is None:
            tenant.google_id = google_id
            db.commit()
            db.refresh(tenant)
        return tenant

    return None


def register_tenant_for_google(
    email: str,
    google_id: str,
    name: str | None,
    db: Session,
    client_ip: str,
) -> Tenant:
    """Create a new tenant for a verified Google identity (rate-limited).

    Caller must have already confirmed registration intent (B12 Layer 1) and
    run collision detection. Creates a credits-mode, zero-balance tenant.
    """
    email = email.lower()
    enforce_registration_rate_limit(client_ip)

    _, hashed_key = generate_api_key()
    tenant = Tenant(
        name=name or email.split("@")[0],
        owner_email=email,
        google_id=google_id,
        api_key=hashed_key,
        is_active=True,
        billing_mode="credits",
        balance_usd=Decimal("0"),
    )
    db.add(tenant)
    db.commit()
    db.refresh(tenant)
    return tenant


def lookup_or_register_tenant(
    email: str,
    google_id: str,
    name: str | None,
    db: Session,
    client_ip: str,
) -> tuple[Tenant, bool]:
    """Find the tenant for a verified Google login, registering one if none exists.

    Retained as a composition of lookup + register for callers that want the
    legacy auto-create behavior. B12 Layer 1 deliberately does NOT use this in
    google_auth anymore — login looks up only; creation goes through the
    confirm-register endpoint. The caller must run collision detection first.
    """
    tenant = lookup_tenant_for_google(email, google_id, db)
    if tenant is not None:
        return tenant, False
    tenant = register_tenant_for_google(email, google_id, name, db, client_ip)
    return tenant, True


def issue_partial_jwt(tenant_id: uuid.UUID, jwt_secret_key: str) -> str:
    """Issue a 5-minute JWT with scope=totp_pending."""
    payload = {
        "sub": str(tenant_id),
        "scope": PARTIAL_TOKEN_SCOPE,
        "exp": datetime.now(timezone.utc) + timedelta(minutes=PARTIAL_TOKEN_EXPIRE_MINUTES),
    }
    return jwt.encode(payload, jwt_secret_key, algorithm="HS256")


def decode_partial_jwt(token: str, jwt_secret_key: str) -> str:
    """Decode a partial JWT; return tenant_id as str. Raise ValueError on any failure."""
    try:
        payload = jwt.decode(token, jwt_secret_key, algorithms=["HS256"])
    except jwt.PyJWTError as exc:
        raise ValueError("Invalid or expired token") from exc
    if payload.get("scope") != PARTIAL_TOKEN_SCOPE:
        raise ValueError("Invalid scope")
    return payload["sub"]


def issue_full_jwt(tenant_id: uuid.UUID, jwt_secret_key: str) -> str:
    """Issue a 24-hour admin JWT on successful TOTP verification."""
    payload = {
        "sub": str(tenant_id),
        "scope": FULL_TOKEN_SCOPE,
        "exp": datetime.now(timezone.utc) + timedelta(hours=FULL_TOKEN_EXPIRE_HOURS),
    }
    return jwt.encode(payload, jwt_secret_key, algorithm="HS256")


def decode_full_jwt(token: str, jwt_secret_key: str) -> str:
    """Decode an admin JWT; return tenant_id as str. Raise ValueError on any failure."""
    try:
        payload = jwt.decode(token, jwt_secret_key, algorithms=["HS256"])
    except jwt.PyJWTError as exc:
        raise ValueError("Invalid or expired token") from exc
    if payload.get("scope") != FULL_TOKEN_SCOPE:
        raise ValueError("Invalid scope")
    return payload["sub"]


def issue_super_admin_jwt(email: str, jwt_secret_key: str) -> str:
    """Issue a 1-hour super-admin JWT with sub=email."""
    payload = {
        "sub": email,
        "scope": SUPER_ADMIN_TOKEN_SCOPE,
        "exp": datetime.now(timezone.utc) + timedelta(hours=SUPER_ADMIN_TOKEN_EXPIRE_HOURS),
    }
    return jwt.encode(payload, jwt_secret_key, algorithm="HS256")


def decode_super_admin_jwt(token: str, jwt_secret_key: str) -> str:
    """Decode a super-admin JWT; return email as str. Raise ValueError on any failure."""
    try:
        payload = jwt.decode(token, jwt_secret_key, algorithms=["HS256"])
    except jwt.PyJWTError as exc:
        raise ValueError("Invalid or expired token") from exc
    if payload.get("scope") != SUPER_ADMIN_TOKEN_SCOPE:
        raise ValueError("Invalid scope")
    email = payload.get("sub")
    if not email:
        raise ValueError("Missing sub claim")
    return email