auth.py

import os
import uuid
import logging
from datetime import datetime, timedelta
from urllib.parse import quote_plus
from typing import List, Optional, Any

from dotenv import load_dotenv
from fastapi import APIRouter, HTTPException, Depends, Request, status, UploadFile, File, Form
from fastapi.responses import JSONResponse
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from jose import JWTError, jwt
from passlib.context import CryptContext
from pydantic import BaseModel, EmailStr, Field, validator
from pymongo import MongoClient
import os.path

load_dotenv()

# Setup logging
logger = logging.getLogger("uvicorn")
logger.setLevel(logging.INFO)

# MongoDB setup for user management
password = quote_plus(os.getenv("MONGO_PASSWORD"))
MONGO_URL = os.getenv("CONNECTION_STRING").replace("${PASSWORD}", password)
client = MongoClient(MONGO_URL)
db = client.users_database
users_collection = db.users

# Rate limiter (applied to auth endpoints)
limiter = Limiter(key_func=get_remote_address, default_limits=["200 per day", "50 per hour"])

# OAuth2 setup
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

# Create an APIRouter instance
router = APIRouter()

# Pydantic models
class User(BaseModel):
    name: str = Field(..., min_length=3, max_length=50)
    email: EmailStr
    password: str

    @validator("password")
    def validate_password(cls, value):
        if len(value) < 8:
            raise ValueError("Password must be at least 8 characters long.")
        if not any(char.isdigit() for char in value):
            raise ValueError("Password must include at least one number.")
        if not any(char.isupper() for char in value):
            raise ValueError("Password must include at least one uppercase letter.")
        if not any(char.islower() for char in value):
            raise ValueError("Password must include at least one lowercase letter.")
        if not any(char in "!@#$%^&*()-_+=<>?/" for char in value):
            raise ValueError("Password must include at least one special character.")
        return value

class UserUpdate(BaseModel):
    name: Optional[str] = Field(None, min_length=3, max_length=50)
    email: Optional[EmailStr]
    password: Optional[str]

    @validator("password")
    def validate_password(cls, value):
        if value is not None:
            if len(value) < 8:
                raise ValueError("Password must be at least 8 characters long.")
            if not any(char.isdigit() for char in value):
                raise ValueError("Password must include at least one number.")
            if not any(char.isupper() for char in value):
                raise ValueError("Password must include at least one uppercase letter.")
            if not any(char.islower() for char in value):
                raise ValueError("Password must include at least one lowercase letter.")
            if not any(char in "!@#$%^&*()-_+=<>?/" for char in value):
                raise ValueError("Password must include at least one special character.")
        return value

class Token(BaseModel):
    access_token: str
    refresh_token: str
    token_type: str

class LoginResponse(Token):
    name: str
    avatar: Optional[str] = None

class TokenData(BaseModel):
    email: Optional[str] = None

# Password hashing
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

def verify_password(plain_password: str, hashed_password: str) -> bool:
    return pwd_context.verify(plain_password, hashed_password)

def get_password_hash(password: str) -> str:
    return pwd_context.hash(password)

def get_user(email: str) -> Optional[dict]:
    return users_collection.find_one({"email": email})

def authenticate_user(email: str, password: str) -> Optional[dict]:
    user = get_user(email)
    if not user or not verify_password(password, user["hashed_password"]):
        return None
    return user

def create_token(data: dict, expires_delta: timedelta = None) -> str:
    to_encode = data.copy()
    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=15))
    to_encode.update({"exp": expire})
    secret_key = os.getenv("SECRET_KEY")
    algorithm = "HS256"
    return jwt.encode(to_encode, secret_key, algorithm=algorithm)

def create_access_token(email: str) -> str:
    return create_token({"sub": email}, timedelta(minutes=int(os.getenv("ACCESS_TOKEN_EXPIRE_MINUTES", "90"))))

def create_refresh_token(email: str) -> str:
    return create_token({"sub": email}, timedelta(days=int(os.getenv("REFRESH_TOKEN_EXPIRE_DAYS", "7"))))

def get_current_user(token: str = Depends(oauth2_scheme)) -> dict:
    secret_key = os.getenv("SECRET_KEY")
    try:
        payload = jwt.decode(token, secret_key, algorithms=["HS256"])
        email: str = payload.get("sub")
        if not email:
            raise HTTPException(status_code=401, detail="Invalid credentials")
        user = get_user(email)
        if not user:
            raise HTTPException(status_code=401, detail="User not found")
        return user
    except JWTError:
        raise HTTPException(status_code=401, detail="Invalid token")

# Setup for avatar file saving
AVATAR_DIR = "avatars"
if not os.path.exists(AVATAR_DIR):
    os.makedirs(AVATAR_DIR)

def save_avatar_file(file: UploadFile) -> str:
    allowed_types = ["image/jpeg", "image/png", "image/gif"]
    if file.content_type not in allowed_types:
        logger.error(f"Unsupported file type: {file.content_type}")
        raise HTTPException(
            status_code=400,
            detail="Invalid image format. Only JPEG, PNG, and GIF are accepted."
        )
    file_extension = os.path.splitext(file.filename)[1]
    unique_filename = f"{uuid.uuid4()}{file_extension}"
    file_path = os.path.join(AVATAR_DIR, unique_filename)
    try:
        contents = file.file.read()
        with open(file_path, "wb") as f:
            f.write(contents)
        logger.info(f"Avatar saved as {file_path}")
    except Exception as e:
        logger.exception("Failed to save avatar file")
        raise HTTPException(status_code=500, detail="Could not save avatar file.")
    finally:
        file.file.close()
    return file_path

# ----- Auth Endpoints -----

@router.post("/signup", response_model=Token)
@limiter.limit("5/minute")
async def signup(
    request: Request,
    name: str = Form(...),
    email: EmailStr = Form(...),
    password: str = Form(...),
    avatar: Optional[UploadFile] = File(None)
):
    # Validate input using the User model
    try:
        _ = User(name=name, email=email, password=password)
    except Exception as e:
        logger.error(f"Validation error during signup: {e}")
        raise HTTPException(status_code=400, detail=str(e))
    if get_user(email):
        logger.warning(f"Attempt to register already existing email: {email}")
        raise HTTPException(status_code=400, detail="Email already registered")
    hashed_password = get_password_hash(password)
    user_data = {
        "name": name,
        "email": email,
        "hashed_password": hashed_password,
        "chat_histories": []  # Initialize an empty array for chat histories and future data.
    }
    if avatar:
        avatar_path = save_avatar_file(avatar)
        user_data["avatar"] = f"/avatars/{os.path.basename(avatar_path)}"
    users_collection.insert_one(user_data)
    logger.info(f"New user registered: {email}")
    return {
        "access_token": create_access_token(email),
        "refresh_token": create_refresh_token(email),
        "token_type": "bearer"
    }

@router.post("/login", response_model=LoginResponse)
@limiter.limit("10/minute")
async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends()):
    user = authenticate_user(form_data.username, form_data.password)
    if not user:
        logger.warning(f"Failed login attempt for: {form_data.username}")
        raise HTTPException(status_code=401, detail="Incorrect username or password")
    logger.info(f"User logged in: {user['email']}")
    return {
        "access_token": create_access_token(user["email"]),
        "refresh_token": create_refresh_token(user["email"]),
        "token_type": "bearer",
        "name": user["name"],
        "avatar": user.get("avatar")
    }

@router.get("/user/data")
@limiter.limit("20/minute")
async def get_user_data(request: Request, current_user: dict = Depends(get_current_user)):
    return {
        "name": current_user["name"],
        "email": current_user["email"],
        "avatar": current_user.get("avatar"),
        "chat_histories": current_user.get("chat_histories", [])
    }

@router.put("/user/update")
@limiter.limit("10/minute")
async def update_user(
    request: Request,
    name: Optional[str] = Form(None),
    email: Optional[EmailStr] = Form(None),
    password: Optional[str] = Form(None),
    avatar: Optional[UploadFile] = File(None),
    current_user: dict = Depends(get_current_user)
):
    update_data = {}
    if name is not None:
        update_data["name"] = name
    if email is not None:
        update_data["email"] = email
    if password is not None:
        try:
            _ = User(name=current_user["name"], email=current_user["email"], password=password)
        except Exception as e:
            logger.error(f"Password validation error during update: {e}")
            raise HTTPException(status_code=400, detail=str(e))
        update_data["hashed_password"] = get_password_hash(password)
    if avatar:
        avatar_path = save_avatar_file(avatar)
        update_data["avatar"] = f"/avatars/{os.path.basename(avatar_path)}"
    if not update_data:
        logger.info("No update parameters provided")
        raise HTTPException(status_code=400, detail="No update parameters provided")
    users_collection.update_one({"email": current_user["email"]}, {"$set": update_data})
    logger.info(f"User updated: {current_user['email']}")
    return {"message": "User updated successfully"}

@router.post("/logout")
@limiter.limit("20/minute")
async def logout(request: Request, current_user: dict = Depends(get_current_user)):
    logger.info(f"User logged out: {current_user['email']}")
    return {"message": "User logged out successfully"}