backend/app/api/auth.py

import json
import logging
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, status, Query
from fastapi.responses import RedirectResponse
from fastapi.security import OAuth2PasswordRequestForm
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from pydantic import BaseModel, EmailStr
from google_auth_oauthlib.flow import Flow

from app.config import settings
from app.database.postgres import get_db, User
from app.auth.security import get_password_hash, verify_password, create_access_token, encrypt_token

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/auth", tags=["auth"])

class UserRegister(BaseModel):
    email: EmailStr
    name: Optional[str] = None
    password: str

class TokenResponse(BaseModel):
    access_token: str
    token_type: str
    user_id: str
    email: str

@router.post("/register", response_model=TokenResponse)
async def register(user_in: UserRegister, db: AsyncSession = Depends(get_db)):
    # Check if user already exists
    result = await db.execute(select(User).where(User.email == user_in.email))
    existing_user = result.scalars().first()
    if existing_user:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="User with this email already registered"
        )

    # Create new user
    hashed_pwd = get_password_hash(user_in.password)
    user = User(
        email=user_in.email,
        name=user_in.name or user_in.email.split("@")[0],
        hashed_password=hashed_pwd
    )
    db.add(user)
    await db.commit()
    await db.refresh(user)

    access_token = create_access_token(data={"sub": str(user.id)})
    return {
        "access_token": access_token,
        "token_type": "bearer",
        "user_id": str(user.id),
        "email": user.email
    }

@router.post("/login", response_model=TokenResponse)
async def login(form_data: OAuth2PasswordRequestForm = Depends(), db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(User).where(User.email == form_data.username))
    user = result.scalars().first()
    
    if not user or not user.hashed_password or not verify_password(form_data.password, user.hashed_password):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect email or password",
            headers={"WWW-Authenticate": "Bearer"},
        )

    access_token = create_access_token(data={"sub": str(user.id)})
    return {
        "access_token": access_token,
        "token_type": "bearer",
        "user_id": str(user.id),
        "email": user.email
    }

@router.get("/google-login")
async def google_login():
    """Build and return Google OAuth login URL or mock url."""
    if settings.MOCK_EMAIL_PROVIDER:
        # Mock redirect URI for offline testing
        mock_auth_url = f"http://localhost:8000/api/auth/callback?code=mock_code_123&state=mock_state"
        return {"url": mock_auth_url}
        
    try:
        import urllib.parse
        import secrets
        
        scopes = [
            "https://www.googleapis.com/auth/gmail.readonly",
            "https://www.googleapis.com/auth/gmail.send",
            "https://www.googleapis.com/auth/gmail.compose",
            "https://www.googleapis.com/auth/userinfo.email",
            "openid"
        ]
        
        params = {
            "client_id": settings.GOOGLE_CLIENT_ID,
            "redirect_uri": settings.GOOGLE_REDIRECT_URI,
            "response_type": "code",
            "scope": " ".join(scopes),
            "access_type": "offline",
            "prompt": "consent",
            "state": secrets.token_urlsafe(16)
        }
        
        auth_url = "https://accounts.google.com/o/oauth2/v2/auth?" + urllib.parse.urlencode(params)
        return {"url": auth_url}
    except Exception as e:
        logger.error(f"Error creating Google Auth Flow: {e}")
        # Fallback to mock login url for seamless testing
        mock_auth_url = f"http://localhost:8000/api/auth/callback?code=mock_code_123&state=mock_state"
        return {"url": mock_auth_url}

@router.get("/callback")
async def auth_callback(code: str = Query(...), db: AsyncSession = Depends(get_db)):
    """Handle callback from Google (or the mock setup) and exchange codes for access tokens."""
    try:
        if settings.MOCK_EMAIL_PROVIDER or code.startswith("mock_"):
            # Mock exchange flow
            email = "demo-user@example.com"
            name = "Demo User"
            tokens = {
                "access_token": "mock_access_token",
                "refresh_token": "mock_refresh_token",
                "expires_in": 3600
            }
        else:
            # Manual HTTPX code exchange to bypass PKCE/code verifier issues
            import httpx
            async with httpx.AsyncClient() as client:
                token_resp = await client.post(
                    "https://oauth2.googleapis.com/token",
                    data={
                        "client_id": settings.GOOGLE_CLIENT_ID,
                        "client_secret": settings.GOOGLE_CLIENT_SECRET,
                        "code": code,
                        "grant_type": "authorization_code",
                        "redirect_uri": settings.GOOGLE_REDIRECT_URI
                    }
                )
                if token_resp.status_code != 200:
                    raise HTTPException(
                        status_code=400,
                        detail=f"Token exchange failed: {token_resp.text}"
                    )
                
                token_data = token_resp.json()
                access_token = token_data.get("access_token")
                
                # Fetch user details using access_token
                user_info_resp = await client.get(
                    "https://www.googleapis.com/oauth2/v3/userinfo",
                    headers={"Authorization": f"Bearer {access_token}"}
                )
                if user_info_resp.status_code != 200:
                    raise HTTPException(
                        status_code=400,
                        detail="Failed to fetch Google user info"
                    )
                    
                user_info = user_info_resp.json()
                email = user_info.get("email", "oauth-user@example.com")
                name = user_info.get("name", "OAuth User")
                
                tokens = {
                    "access_token": access_token,
                    "refresh_token": token_data.get("refresh_token"),
                    "expires_in": token_data.get("expires_in"),
                    "token_uri": "https://oauth2.googleapis.com/token",
                    "client_id": settings.GOOGLE_CLIENT_ID,
                    "client_secret": settings.GOOGLE_CLIENT_SECRET,
                    "scopes": token_data.get("scope", "").split(" ")
                }

        # Check if user already exists
        result = await db.execute(select(User).where(User.email == email))
        user = result.scalars().first()
        
        # Serialize and encrypt tokens dict
        encrypted_creds = encrypt_token(json.dumps(tokens))
        
        if not user:
            user = User(
                email=email,
                name=name,
                oauth_token_encrypted=encrypted_creds
            )
            db.add(user)
        else:
            user.oauth_token_encrypted = encrypted_creds
            
        await db.commit()
        await db.refresh(user)

        # Generate app JWT token
        jwt_token = create_access_token(data={"sub": str(user.id)})
        
        # Redirect back to the frontend with token params
        frontend_redirect_url = f"http://localhost:3000/login?token={jwt_token}&email={user.email}&userId={user.id}"
        return RedirectResponse(url=frontend_redirect_url)
    except Exception as e:
        logger.error(f"Failed in oauth callback handler: {e}")
        # Redirect to frontend login with error flag
        return RedirectResponse(url="http://localhost:3000/login?error=oauth_failed")