routers/auth.py

from fastapi import APIRouter, Depends, status, HTTPException, Request
from passlib.context import CryptContext
from pydantic import BaseModel, field_validator
from typing import Annotated
from fastapi.security import OAuth2PasswordBearer
import re
import os
from jose import jwt, JWTError
from datetime import datetime, timedelta, timezone
from models.tables import Users
from config.db import SessionLocal
from sqlalchemy.orm import Session
from fastapi.responses import JSONResponse
import logging
from models.tables import Company
import json

router = APIRouter(prefix='/auth', tags=['Auth'])

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


crypt = CryptContext(schemes=['bcrypt'], deprecated='auto')
oauth2_bearer = OAuth2PasswordBearer(tokenUrl='auth/login')

SECRET_KEY = os.getenv("JWT_SECRET_KEY", "CHANGE_ME_IN_PRODUCTION")
ALGORITHM = "HS256"

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

class ValidationException(HTTPException):
    def __init__(self, detail: str):
        super().__init__(status_code=status.HTTP_400_BAD_REQUEST, detail=detail)

class LoginUserData(BaseModel):
    username: str
    password: str

    @field_validator('password')
    def validate_password(cls, v):
        if len(v) < 8:
            raise ValidationException({"status": False, "mssg": 'Password must be at least 8 characters long.'})       
        if not re.search(r'[A-Z]', v):
            raise ValidationException({"status": False, "mssg": 'Password must contain at least one uppercase letter.'})       
        if not re.search(r'[a-z]', v):
            raise ValidationException({"status": False, "mssg": 'Password must contain at least one lowercase letter.'})
        if not re.search(r'[0-9]', v):
            raise ValidationException({"status": False, "mssg": 'Password must contain at least one numeric digit.'})
        if not re.search(r'[\W_]', v):  
            raise ValidationException({"status": False, "mssg": 'Password must contain at least one special character.'})
        return v

class UserBase(BaseModel):
    username: str
    email: str
    first_name: str 
    last_name: str 
    password: str
    role: str

    @field_validator('password')
    def validate_password(cls, v):
        if len(v) < 8:
            raise ValidationException({"status": False, "mssg": 'Password must be at least 8 characters long.'})       
        if not re.search(r'[A-Z]', v):
            raise ValidationException({"status": False, "mssg": 'Password must contain at least one uppercase letter.'})       
        if not re.search(r'[a-z]', v):
            raise ValidationException({"status": False, "mssg": 'Password must contain at least one lowercase letter.'})
        if not re.search(r'[0-9]', v):
            raise ValidationException({"status": False, "mssg": 'Password must contain at least one numeric digit.'})
        if not re.search(r'[\W_]', v):  
            raise ValidationException({"status": False, "mssg": 'Password must contain at least one special character.'})
        return v

def get_db():
    db = SessionLocal()
    try:
        yield db 
    finally:
        db.close()

db_dependency = Annotated[Session, Depends(get_db)]

def get_token(request: Request):
    auth_header = request.headers.get("Authorization")
    if not auth_header:
        raise ValidationException({"status": False, "mssg": 'Authorization Header missing.'})   
    token_type, token = auth_header.split()
    if token_type.lower() != "bearer":
        raise ValidationException({"status": False, "mssg": 'Token type is not bearer.'})   
    return token

def authenticate_user(username: str, password: str, db):
    user = db.query(Users).filter(Users.username == username).first()
    if not user:
        logger.warning("Authentication failed: Username does not exist: %s", username)
        raise ValidationException({"status": False, "mssg": 'Username does not exist.'})       
    if not crypt.verify(password, user.hashed_password):
        logger.warning("Authentication failed: Incorrect password for user: %s", username)
        raise ValidationException({"status": False, "mssg": 'Password does not match.'})       
    logger.info("User authenticated successfully: %s", username)
    return user

def create_access_token(username: str, user_id: int, role: str, expire_delta: timedelta):
    encode = {"sub": username, 'id': user_id, 'user_role': role}
    expires = datetime.now(timezone.utc) + expire_delta
    encode.update({'exp': expires})
    return jwt.encode(encode, SECRET_KEY, algorithm=ALGORITHM)

async def get_current_user_with_token(token: Annotated[str, Depends(get_token)]):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get('sub')
        user_id: int = payload.get('id')

        if username is None or user_id is None:
            logger.error("Failed to verify user: username or user_id is None.")
            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not verify user.")
        logger.info("Current user verified: %s", username)
        return {"username": username, 'id': user_id}
    except JWTError:
        logger.error("JWT error during user verification.")
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not verify user")

async def get_current_user(token: Annotated[str, Depends(oauth2_bearer)]):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get('sub')
        user_id: int = payload.get('id')

        if username is None or user_id is None:
            logger.error("Failed to verify user: username or user_id is None.")
            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not verify user.")
        logger.info("Current user verified: %s", username)
        return {"username": username, 'id': user_id}
    except JWTError:
        logger.error("JWT error during user verification.")
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not verify user")

@router.post('/signup', status_code=status.HTTP_201_CREATED)
async def signup(db: db_dependency, user: UserBase):
    existing_user = db.query(Users).filter((Users.username == user.username) | (Users.email == user.email)).first()

    if existing_user:
        if existing_user.username == user.username:
            logger.warning("Signup failed: Username already exists: %s", user.username)
            raise ValidationException({"status": False, "mssg": 'Username already exists.'})       

        if existing_user.email == user.email:
            logger.warning("Signup failed: User email already exists: %s", user.email)
            raise ValidationException({"status": False, "mssg": 'User email already exists.'})       

    create_user_model = Users(
        email=user.email,
        username=user.username,
        first_name=user.first_name,
        last_name=user.last_name,
        hashed_password=crypt.hash(user.password),
        role=user.role,
        is_active=True
    )
    db.add(create_user_model)
    db.commit()
    logger.info("User account created successfully: %s", user.username)
    return JSONResponse(content={"detail": {"status": True, "mssg": "Account Created Successfully!"}})

user_dependency = Annotated[dict, Depends(get_current_user)]
@router.post('/login')
async def login(form_data: LoginUserData, db: db_dependency):
    user = authenticate_user(form_data.username, form_data.password, db)
    token = create_access_token(user.username, user.id, user.role, timedelta(days=30))
    logger.info("User logged in successfully: %s", user.username)
    if  user.role =='admin':
        return JSONResponse(content={"access_token": token, 'token_type': 'bearer','role':user.role})
    else:
        company = db.query(Company).filter(Company.email == user.email).first()
        company_json = {
            "email": company.email,
            "company_name": company.company_name,
            "base_url": company.base_url,
            "company_id": company.company_key
        }
        return JSONResponse(content={"access_token": token, 'token_type': 'bearer','role':user.role,'company_data':company_json})