from functools import wraps
from flask import jsonify, request, g, session
from src.lib.jwt import verify_token
from src.services.principals import find_principal_by_id
from src.services.teachers import find_teacher_by_id
from src.services.schools import find_school_by_id
from src.utils.response import error_response

def handle_exceptions(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            return jsonify({
                "success": False,
                "error": str(e),
                "message": "Internal server error"
            }), 500
    return wrapper


def token_required(role=None, optional=False):  # Added 'optional' parameter
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Check if already authenticated via g
            if hasattr(g, "current_user") and hasattr(g, "current_role"):
                if not role or g.current_role == role:
                    return func(*args, **kwargs)

            # Check for session-based auth (for students)
            if session.get("student"):
                # If optional=True, allow session-based auth to pass through
                if optional:
                    return func(*args, **kwargs)

            # Try to get JWT token from cookies
            token = request.cookies.get("auth_token")
            if not token:
                if optional:  # If token is optional, continue without auth
                    return func(*args, **kwargs)
                return error_response("No active user found", 401)

            # Verify JWT token
            user_id, user_role = verify_token(token)
            if not user_id or (role and role != user_role):
                if optional:
                    return func(*args, **kwargs)
                return error_response("Invalid or unauthorized token", 403)

            # Check if user already loaded in g
            user = None
            if hasattr(g, "current_user") and hasattr(g, "current_role"):
                if g.current_role == user_role:
                    user = g.current_user

            # Fetch user from database if not in g
            if not user:
                if user_role == "teacher":
                    user = find_teacher_by_id(user_id)
                elif user_role == "principal":
                    user = find_principal_by_id(user_id)
                elif user_role == "school":
                    user = find_school_by_id(user_id)
                else:
                    if optional:
                        return func(*args, **kwargs)
                    return error_response("Invalid role", 403)

                if not user:
                    if optional:
                        return func(*args, **kwargs)
                    return error_response(f"{user_role.capitalize()} not found", 404)

                g.current_user = user
                g.current_role = user_role

            return func(*args, **kwargs)

        return wrapper
    return decorator