api/auth/auth_handler.py

"""
Authentication handler for patient and therapist portals
"""
import os
import jwt
import bcrypt
from datetime import datetime, timedelta, date
from typing import Optional, Dict
from sqlalchemy.orm import Session
from flask import request

from database.models.user import User, UserRole
from database.models.student import Student
from database.models.therapist import Therapist
from config.database import get_db
from config.settings import PAYMENT_SYSTEM_ENABLED
from api.services.email_service import email_service
from utils.session_manager import SessionManager
from utils.audit_logger import HIPAAAuditLogger
from utils.password_validator import PasswordValidator
from utils.consent_service import ConsentService

# JWT Configuration
SECRET_KEY = os.getenv("JWT_SECRET_KEY", "your-secret-key-change-in-production")
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 24  # 24 hours

class AuthException(Exception):
    """Custom exception for auth errors"""
    def __init__(self, message, status_code=401):
        self.message = message
        self.status_code = status_code
        super().__init__(self.message)


class AuthHandler:
    """Handle authentication for both patient and therapist portals"""
    
    def get_password_hash(self, password: str) -> str:
        """Hash a password using bcrypt"""
        return bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
    
    def verify_password(self, plain_password: str, hashed_password: str) -> bool:
        """Verify a password against hash"""
        return bcrypt.checkpw(plain_password.encode('utf-8'), hashed_password.encode('utf-8'))
    
    def create_access_token(self, user_id: int, role: str) -> str:
        """Create JWT token with session management"""
        # Use SessionManager for HIPAA-compliant session handling
        return SessionManager.create_token(user_id, role)

    def create_mfa_pending_token(self, user_id: int) -> str:
        """Create a short-lived token for MFA verification step"""
        payload = {
            'sub': str(user_id),
            'type': 'mfa_pending',
            'exp': datetime.utcnow() + timedelta(minutes=5)  # 5 minute expiry
        }
        return jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM)
    
    def decode_token(self, token: str) -> Dict:
        """Decode and validate JWT token with session checking"""
        try:
            # First check if session is still valid
            is_valid, message = SessionManager.validate_session(token)
            if not is_valid:
                raise AuthException(message, 401)

            # Decode the token using SessionManager's JWT secret
            payload = jwt.decode(
                token,
                os.environ.get('JWT_SECRET_KEY', 'your-secret-key-change-this'),
                algorithms=['HS256']
            )
            return payload
        except jwt.ExpiredSignatureError:
            raise AuthException("Session has expired", 401)
        except jwt.InvalidTokenError:
            raise AuthException("Invalid session token", 401)
    
    async def register_patient(self, db: Session, email: str, password: str,
                              username: str, first_name: str, last_name: str,
                              date_of_birth: str = None, age: int = None,
                              phone_number: str = None,
                              address_line1: str = None, address_line2: str = None,
                              city: str = None, state_province: str = None,
                              postal_code: str = None, country: str = 'USA',
                              timezone: str = 'America/New_York',
                              terms_accepted: bool = False,
                              privacy_accepted: bool = False,
                              health_data_accepted: bool = False,
                              parent_email: str = None,
                              parent_name: str = None,
                              access_code: str = None) -> Dict:
        """Register a new patient"""
        # Debug logging
        import logging
        logger = logging.getLogger(__name__)
        logger.info(f"Attempting to register patient with email: {email}, username: {username}")
        
        # Check if user exists
        # First check how many total users are in database
        total_users = db.query(User).count()
        logger.info(f"Total users in database: {total_users}")
        
        # List all users for debugging
        all_users = db.query(User).all()
        for u in all_users:
            logger.info(f"  User {u.id}: {u.email} ({u.username})")
        
        existing_user = db.query(User).filter(
            (User.email == email) | (User.username == username)
        ).first()
        
        if existing_user:
            logger.error(f"Found existing user: ID={existing_user.id}, email={existing_user.email}, username={existing_user.username}")
            # Double check by direct ID query
            user_by_id = db.query(User).filter(User.id == existing_user.id).first()
            if user_by_id:
                logger.error(f"Confirmed user exists by ID: {user_by_id.id}")
            else:
                logger.error(f"STRANGE: User ID {existing_user.id} not found when queried directly!")
            raise AuthException("User with this email or username already exists", 400)

        logger.info(f"No existing user found, proceeding with registration")

        # PIPEDA/HIPAA Compliance: Validate consent was given
        if not (terms_accepted and privacy_accepted and health_data_accepted):
            raise AuthException(
                "You must accept the Terms of Service, Privacy Policy, and Health Data Collection consent to create an account",
                400
            )

        # Check if parental consent required for minors
        if date_of_birth and ConsentService.is_minor(date_of_birth):
            if not (parent_email and parent_name):
                raise AuthException(
                    "Parental consent is required for users under 18. Please provide parent/guardian email and name.",
                    400
                )

        # Validate password strength for HIPAA compliance
        is_valid, password_errors = PasswordValidator.validate_password(
            password=password,
            username=username,
            email=email
        )

        if not is_valid:
            error_msg = "Password does not meet security requirements: " + "; ".join(password_errors)
            logger.warning(f"Weak password attempt for {email}: {error_msg}")
            raise AuthException(error_msg, 400)

        # Validate access code if provided (for beta testers)
        access_code_obj = None
        free_access_until = None
        if access_code:
            from database.models.access_code import AccessCode
            access_code_obj = db.query(AccessCode).filter(AccessCode.code == access_code).first()

            if not access_code_obj:
                raise AuthException("Invalid access code", 400)

            if not access_code_obj.is_valid():
                if access_code_obj.is_used:
                    raise AuthException("This access code has already been used", 400)
                else:
                    raise AuthException("This access code has expired", 400)

            # Only allow beta codes during signup if payment system is not enabled
            # Once payment system is implemented, set PAYMENT_SYSTEM_ENABLED=True in config
            if access_code_obj.code_type != 'beta' and not PAYMENT_SYSTEM_ENABLED:
                raise AuthException("This access code cannot be used during signup. Please contact support for payment options.", 400)

            # Calculate free access expiration
            free_access_until = datetime.utcnow() + timedelta(days=access_code_obj.access_duration_days)
            logger.info(f"Valid beta code provided: {access_code}. Granting free access until {free_access_until}")

        # Create user with verification token and code
        hashed_password = self.get_password_hash(password)
        verification_token = email_service.generate_verification_token()
        verification_code = email_service.generate_verification_code()
        
        user = User(
            email=email,
            username=username,
            hashed_password=hashed_password,
            first_name=first_name,
            last_name=last_name,
            role=UserRole.PATIENT,
            phone_number=phone_number,
            is_active=True,
            is_verified=False,
            email_verification_token=verification_token,
            email_verification_code=verification_code,
            verification_sent_at=datetime.utcnow(),
            free_access_until=free_access_until  # Set free access if beta tester code was used
        )
        db.add(user)
        db.flush()
        
        # Parse date_of_birth if provided
        dob = None
        calculated_age = age  # Use provided age as fallback
        if date_of_birth:
            try:
                # Handle different date formats
                if isinstance(date_of_birth, str):
                    dob = datetime.strptime(date_of_birth, '%Y-%m-%d').date()
                else:
                    dob = date_of_birth
                
                # Calculate age from birthdate
                today = date.today()
                calculated_age = today.year - dob.year
                if (today.month, today.day) < (dob.month, dob.day):
                    calculated_age -= 1
            except:
                pass  # If parsing fails, fall back to provided age
        
        # Create patient profile with location
        patient = Patient(
            user_id=user.id,
            date_of_birth=dob,
            age=calculated_age,  # Store calculated age or provided age
            address_line1=address_line1,
            address_line2=address_line2,
            city=city,
            state_province=state_province,
            postal_code=postal_code,
            country=country,
            timezone=timezone,
            therapy_goals=[],  # Initialize as empty array for JSON field
            medical_history="",
            target_sounds=[]  # Initialize as empty array for JSON field
        )
        db.add(patient)

        # Note: Patients will connect to therapists after signup using the ConnectTherapist page
        # No invitation code handling during signup

        db.commit()

        # Mark access code as used if provided
        if access_code_obj:
            try:
                access_code_obj.is_used = True
                access_code_obj.used_at = datetime.utcnow()
                access_code_obj.used_by_email = email
                access_code_obj.used_by_user_id = user.id
                db.commit()
                logger.info(f"Access code {access_code} marked as used by {email}")
            except Exception as e:
                logger.error(f"Failed to mark access code as used: {str(e)}")
                # Don't fail registration if this fails

        # PIPEDA/HIPAA Compliance: Record user consent
        try:
            ConsentService.record_all_signup_consents(
                db=db,
                user_id=user.id,
                terms_accepted=terms_accepted,
                privacy_accepted=privacy_accepted,
                health_data_accepted=health_data_accepted,
                parent_email=parent_email,
                parent_name=parent_name
            )
            logger.info(f"Consent recorded for user {user.id}")
        except Exception as e:
            logger.error(f"Failed to record consent for user {user.id}: {str(e)}")
            # Don't fail registration if consent recording fails, but log it

        # Send verification email
        email_sent = email_service.send_verification_email(
            to_email=email,
            user_name=first_name,
            verification_token=verification_token,
            verification_code=verification_code,
            user_type='patient'
        )
        
        # Create token (but user needs to verify email for full access)
        token = self.create_access_token(user.id, "patient")
        
        return {
            "access_token": token,
            "token_type": "bearer",
            "user": {
                "id": user.id,
                "email": user.email,
                "username": user.username,
                "first_name": user.first_name,
                "last_name": user.last_name,
                "role": "patient",
                "email_verified": False,
                "verification_email_sent": email_sent
            }
        }
    
    async def register_therapist(self, db: Session, email: str, password: str,
                                username: str, first_name: str, last_name: str,
                                job_title: str = None, institution: str = None,
                                country: str = None, province: str = None,
                                specialization: str = None,
                                terms_accepted: bool = False,
                                privacy_accepted: bool = False) -> Dict:
        """Register a new therapist"""
        # Check if user exists
        existing_user = db.query(User).filter(
            (User.email == email) | (User.username == username)
        ).first()

        if existing_user:
            raise AuthException("User with this email or username already exists", 400)

        # PIPEDA/HIPAA Compliance: Validate consent was given
        if not (terms_accepted and privacy_accepted):
            raise AuthException(
                "You must accept the Terms of Service and Privacy Policy to create an account",
                400
            )

        # Validate password strength for HIPAA compliance
        is_valid, password_errors = PasswordValidator.validate_password(
            password=password,
            username=username,
            email=email
        )

        if not is_valid:
            error_msg = "Password does not meet security requirements: " + "; ".join(password_errors)
            raise AuthException(error_msg, 400)

        # Create user with verification token and code
        hashed_password = self.get_password_hash(password)
        verification_token = email_service.generate_verification_token()
        verification_code = email_service.generate_verification_code()
        
        user = User(
            email=email,
            username=username,
            hashed_password=hashed_password,
            first_name=first_name,
            last_name=last_name,
            role=UserRole.THERAPIST,
            is_active=True,
            is_verified=False,
            email_verification_token=verification_token,
            email_verification_code=verification_code,
            verification_sent_at=datetime.utcnow()
        )
        db.add(user)
        db.flush()
        
        # Create therapist profile with additional fields
        therapist = Therapist(
            user_id=user.id,
            license_number=None,  # Optional now
            job_title=job_title,
            institution=institution,
            country=country,
            province=province,
            specialization=specialization
        )
        db.add(therapist)
        db.commit()

        # PIPEDA/HIPAA Compliance: Record user consent
        import logging
        logger = logging.getLogger(__name__)
        try:
            ConsentService.record_all_signup_consents(
                db=db,
                user_id=user.id,
                terms_accepted=terms_accepted,
                privacy_accepted=privacy_accepted,
                health_data_accepted=False  # Therapists don't need health data consent
            )
            logger.info(f"Consent recorded for therapist {user.id}")
        except Exception as e:
            logger.error(f"Failed to record consent for therapist {user.id}: {str(e)}")
            # Don't fail registration if consent recording fails, but log it

        # Send verification email
        email_sent = email_service.send_verification_email(
            to_email=email,
            user_name=first_name,
            verification_token=verification_token,
            verification_code=verification_code,
            user_type='therapist'
        )
        
        # Create token (but user needs to verify email for full access)
        token = self.create_access_token(user.id, "therapist")
        
        return {
            "access_token": token,
            "token_type": "bearer",
            "user": {
                "id": user.id,
                "email": user.email,
                "username": user.username,
                "first_name": user.first_name,
                "last_name": user.last_name,
                "role": "therapist",
                "email_verified": False,
                "verification_email_sent": email_sent
            }
        }
    
    async def login(self, db: Session, username_or_email: str, password: str) -> Dict:
        """Login user (patient or therapist)"""
        # Get IP address and user agent for audit logging
        ip_address = None
        user_agent = None
        try:
            if hasattr(request, 'remote_addr'):
                ip_address = request.remote_addr
            if hasattr(request, 'headers'):
                user_agent = request.headers.get('User-Agent')
        except:
            pass  # Flask request context might not be available in all cases

        # Find user by email or username
        user = db.query(User).filter(
            (User.email == username_or_email) | (User.username == username_or_email)
        ).first()

        if not user:
            # Log failed login attempt
            HIPAAAuditLogger.log_login(
                user_id=None,
                user_role=None,
                user_email=username_or_email,
                ip_address=ip_address,
                success=False,
                error_message="User not found",
                user_agent=user_agent
            )
            raise AuthException("Invalid credentials", 401)

        # Verify password
        if not self.verify_password(password, user.hashed_password):
            # Log failed login attempt
            HIPAAAuditLogger.log_login(
                user_id=user.id,
                user_role=user.role.value,
                user_email=user.email,
                ip_address=ip_address,
                success=False,
                error_message="Invalid password",
                user_agent=user_agent
            )
            raise AuthException("Invalid credentials", 401)

        # Check if account is active
        if not user.is_active:
            # Log failed login attempt
            HIPAAAuditLogger.log_login(
                user_id=user.id,
                user_role=user.role.value,
                user_email=user.email,
                ip_address=ip_address,
                success=False,
                error_message="Account inactive",
                user_agent=user_agent
            )
            raise AuthException("Account is inactive", 403)

        # Check if email is verified
        if not user.is_verified:
            # Log failed login attempt
            HIPAAAuditLogger.log_login(
                user_id=user.id,
                user_role=user.role.value,
                user_email=user.email,
                ip_address=ip_address,
                success=False,
                error_message="Email not verified",
                user_agent=user_agent
            )
            raise AuthException("Please verify your email before logging in. Check your inbox for the verification code.", 403)

        # HIPAA Compliance: Check if MFA is enabled
        if getattr(user, 'mfa_enabled', False) and user.mfa_enabled:
            # Don't update last_login yet - wait for MFA verification
            # Create a temporary MFA token (short-lived)
            mfa_token = self.create_mfa_pending_token(user.id)

            # Log partial login (MFA pending)
            HIPAAAuditLogger.log_login(
                user_id=user.id,
                user_role=user.role.value,
                user_email=user.email,
                ip_address=ip_address,
                success=True,
                user_agent=user_agent,
                additional_info="MFA verification pending"
            )

            return {
                "mfa_required": True,
                "user_id": user.id,
                "mfa_token": mfa_token,
                "message": "Please enter your MFA code to complete login"
            }

        # Update last login
        user.last_login = datetime.utcnow()
        db.commit()

        # Create token
        token = self.create_access_token(user.id, user.role.value)

        # Log successful login
        HIPAAAuditLogger.log_login(
            user_id=user.id,
            user_role=user.role.value,
            user_email=user.email,
            ip_address=ip_address,
            success=True,
            user_agent=user_agent
        )
        
        # Build user response
        user_data = {
            "id": user.id,
            "email": user.email,
            "username": user.username,
            "first_name": user.first_name,
            "last_name": user.last_name,
            "role": user.role.value
        }
        
        # If user is a patient, add their student_id and date_of_birth
        if user.role == UserRole.PATIENT:
            student = db.query(Student).filter_by(user_id=user.id).first()
            if student:
                user_data["student_id"] = student.id
                # Add date_of_birth if available
                if student.date_of_birth:
                    user_data["date_of_birth"] = student.date_of_birth.isoformat()
                # Also calculate and add age
                if student.age:
                    user_data["age"] = student.age
                elif student.date_of_birth:
                    from datetime import date
                    today = date.today()
                    age = today.year - student.date_of_birth.year - ((today.month, today.day) < (student.date_of_birth.month, student.date_of_birth.day))
                    user_data["age"] = age
                print(f"Login: Added student_id {student.id} for user_id {user.id}")
            else:
                print(f"WARNING: No student record found for patient user_id {user.id}")
        
        return {
            "access_token": token,
            "token_type": "bearer",
            "user": user_data
        }
    
    def get_current_user(self, token: str, db: Session) -> User:
        """Get current authenticated user from token"""
        payload = self.decode_token(token)
        
        user_id = int(payload["sub"])
        user = db.query(User).filter(User.id == user_id).first()
        
        if not user:
            raise AuthException("User not found", 404)
        
        if not user.is_active:
            raise AuthException("Account is inactive", 403)
        
        return user
    
    def require_patient(self, current_user: User) -> User:
        """Require current user to be a patient"""
        if current_user.role != UserRole.PATIENT:
            raise AuthException("Patient access required", 403)
        return current_user
    
    def require_therapist(self, current_user: User) -> User:
        """Require current user to be a therapist"""
        if current_user.role != UserRole.THERAPIST:
            raise AuthException("Therapist access required", 403)
        return current_user


# Create singleton instance
auth_handler = AuthHandler()