api/auth.py

from typing import Optional, List
from fastapi import APIRouter, Depends, HTTPException, status, Request, Form, Body
from fastapi.security import OAuth2PasswordBearer, HTTPBearer, HTTPAuthorizationCredentials
from jose import JWTError, jwt
from sqlalchemy.orm import Session
from core.config import settings
from core.database import get_db
from models.schemas import UserCreate, Token, TokenData, UserResponse, UserLogin, ForgotPasswordRequest, ChangePasswordRequest, VerifyOTPRequest, ResetPasswordRequest
from models import db_models
from supabase import create_client, Client
import logging
import httpx
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/auth", tags=["auth"])

# Initialize Supabase Client
supabase: Client = create_client(settings.SUPABASE_URL, settings.SUPABASE_ANON_KEY)

# Global JWKS cache
JWKS_CACHE = None

async def get_jwks():
    """Fetches the JWKS public keys from Supabase."""
    global JWKS_CACHE
    if JWKS_CACHE is None:
        try:
            jwks_url = f"{settings.SUPABASE_URL}/auth/v1/.well-known/jwks.json"
            async with httpx.AsyncClient() as client:
                response = await client.get(jwks_url)
                response.raise_for_status()
                JWKS_CACHE = response.json()
                logger.info("Successfully fetched Supabase JWKS")
        except Exception as e:
            logger.error(f"Failed to fetch JWKS: {str(e)}")
            return None
    return JWKS_CACHE

# Schemes for Swagger
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/login", auto_error=False)
bearer_scheme = HTTPBearer(auto_error=False)

async def get_token(
    request: Request,
    oauth_token: Optional[str] = Depends(oauth2_scheme),
    bearer_token: Optional[HTTPAuthorizationCredentials] = Depends(bearer_scheme)
) -> str:
    token = oauth_token or (bearer_token.credentials if bearer_token else None)
    if not token:
        auth_header = request.headers.get("Authorization")
        if auth_header and auth_header.startswith("Bearer "):
            token = auth_header.split(" ")[1]
            
    if not token:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Not authenticated. Please Authorize with a Bearer token.",
            headers={"WWW-Authenticate": "Bearer"},
        )
    return token

async def get_current_user(token: str = Depends(get_token), db: Session = Depends(get_db)):
    """
    Verifies Supabase JWT using JWKS and returns the local user from Azure SQL.
    Supports both ES256 (asymmetric) and HS256 (symmetric) automatically.
    """
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    
    try:
        # 1. Get Token Header
        header = jwt.get_unverified_header(token)
        alg = header.get("alg")
        
        # 2. Logic based on Algorithm
        if alg == "ES256":
            # Asymmetric: Use JWKS Public Keys
            jwks = await get_jwks()
            if not jwks:
                raise credentials_exception
            
            payload = jwt.decode(
                token, 
                jwks, 
                algorithms=["ES256"],
                options={"verify_aud": False}
            )
        else:
            # Symmetric: Use standard JWT Secret
            payload = jwt.decode(
                token, 
                settings.SUPABASE_JWT_SECRET, 
                algorithms=["HS256", "HS384", "HS512"], 
                options={"verify_aud": False}
            )
            
        supabase_id: str = payload.get("sub")
        email: str = payload.get("email")
        
        if supabase_id is None:
            raise credentials_exception
            
    except JWTError as e:
        logger.error(f"JWT Verification Error: {str(e)}")
        raise credentials_exception
    
    # Standard DB Syncing Logic
    user = db.query(db_models.User).filter(db_models.User.supabase_id == supabase_id).first()
    if user is None:
        user = db.query(db_models.User).filter(db_models.User.email == email).first()
        if user:
            user.supabase_id = supabase_id
            db.commit()
            db.refresh(user)
        else:
            user = db_models.User(email=email, supabase_id=supabase_id, is_active=True)
            db.add(user)
            db.commit()
            db.refresh(user)
            
    return user

async def get_current_user_ws(token: str, db: Session):
    """WS Auth helper (Internal only)."""
    try:
        header = jwt.get_unverified_header(token)
        alg = header.get("alg")
        if alg == "ES256":
            # Use cached JWKS (might fail if cache empty, usually fine after first request)
            jwks = JWKS_CACHE or {} 
            payload = jwt.decode(token, jwks, algorithms=["ES256"], options={"verify_aud": False})
        else:
            payload = jwt.decode(token, settings.SUPABASE_JWT_SECRET, algorithms=["HS256"], options={"verify_aud": False})
        
        supabase_id = payload.get("sub")
        return db.query(db_models.User).filter(db_models.User.supabase_id == supabase_id).first()
    except:
        return None

@router.post("/register", response_model=UserResponse)
async def register(user_in: UserCreate, db: Session = Depends(get_db)):
    try:
        auth_response = supabase.auth.sign_up({"email": user_in.email, "password": user_in.password})
        if not auth_response.user: raise HTTPException(status_code=400, detail="Registration failed")
        
        supabase_id = auth_response.user.id
        db_user = db.query(db_models.User).filter(db_models.User.email == user_in.email).first()
        if db_user: db_user.supabase_id = supabase_id
        else:
            db_user = db_models.User(email=user_in.email, supabase_id=supabase_id, is_active=True)
            db.add(db_user)
        db.commit()
        db.refresh(db_user)
        return db_user
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.post("/login", response_model=Token)
async def login(
    request: Request,
    email: Optional[str] = Body(None),
    password: Optional[str] = Body(None),
    username: Optional[str] = Form(None),
    password_form: Optional[str] = Form(None, alias="password"),
    db: Session = Depends(get_db)
):
    final_email = email or username
    final_password = password or password_form
    if not final_email or not final_password:
        raise HTTPException(status_code=422, detail="Email and password required")

    # 1. Check if user exists locally first for better error messaging
    db_user = db.query(db_models.User).filter(db_models.User.email == final_email).first()
    if not db_user:
        raise HTTPException(status_code=404, detail="User does not exist")

    try:
        response = supabase.auth.sign_in_with_password({"email": final_email, "password": final_password})
        if not response.session: 
            raise HTTPException(status_code=401, detail="Invalid login credentials")

        # Sync supabase_id if it's missing or different (though it should matches)
        if not db_user.supabase_id or db_user.supabase_id != response.user.id:
            db_user.supabase_id = response.user.id
            db.commit()

        return {"access_token": response.session.access_token, "token_type": "bearer"}
    except Exception as e:
        error_str = str(e).lower()
        if "invalid login credentials" in error_str or "invalid_credentials" in error_str:
            raise HTTPException(status_code=401, detail="Invalid login credentials")
        
        logger.error(f"Login failed unexpectedly: {str(e)}")
        raise HTTPException(status_code=401, detail=f"Login failed: {str(e)}")

@router.post("/forgot-password")
async def forgot_password(request: ForgotPasswordRequest):
    try:
        supabase.auth.reset_password_for_email(request.email)
        return {"message": "Verification code has been sent."}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.post("/verify-otp")
async def verify_otp(request: VerifyOTPRequest):
    try:
        verify_response = supabase.auth.verify_otp({
            "email": request.email,
            "token": request.otp,
            "type": "recovery"
        })
        if not verify_response.session: raise HTTPException(status_code=400, detail="Invalid code")
        return {
            "access_token": verify_response.session.access_token,
            "token_type": "bearer",
            "message": "OTP verified."
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.post("/reset-password")
async def reset_password(request: ResetPasswordRequest, token: str = Depends(get_token)):
    try:
        temp_supabase = create_client(settings.SUPABASE_URL, settings.SUPABASE_ANON_KEY)
        temp_supabase.auth.set_session(token, "") 
        temp_supabase.auth.update_user({"password": request.new_password})
        return {"message": "Password successfully reset."}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.post("/change-password")
async def change_password(request: ChangePasswordRequest, current_user: db_models.User = Depends(get_current_user)):
    try:
        supabase.auth.update_user({"password": request.new_password})
        return {"message": "Password updated successfully"}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))