import inspect
import logging

import re
from functools import wraps
from typing import Dict, List, Optional, Any, Callable, Union, Tuple

from google.auth.exceptions import RefreshError
from googleapiclient.discovery import build
from fastmcp.server.dependencies import get_access_token, get_context
from auth.google_auth import get_authenticated_google_service, GoogleAuthenticationError
from auth.oauth21_session_store import (
    get_auth_provider,
    get_oauth21_session_store,
    ensure_session_from_access_token,
)
from auth.oauth_config import is_oauth21_enabled, get_oauth_config
from core.context import set_fastmcp_session_id
from auth.scopes import (
    GMAIL_READONLY_SCOPE,
    GMAIL_SEND_SCOPE,
    GMAIL_COMPOSE_SCOPE,
    GMAIL_MODIFY_SCOPE,
    GMAIL_LABELS_SCOPE,
    DRIVE_READONLY_SCOPE,
    DRIVE_FILE_SCOPE,
    DOCS_READONLY_SCOPE,
    DOCS_WRITE_SCOPE,
    CALENDAR_READONLY_SCOPE,
    CALENDAR_EVENTS_SCOPE,
    SHEETS_READONLY_SCOPE,
    SHEETS_WRITE_SCOPE,
    CHAT_READONLY_SCOPE,
    CHAT_WRITE_SCOPE,
    CHAT_SPACES_SCOPE,
    FORMS_BODY_SCOPE,
    FORMS_BODY_READONLY_SCOPE,
    FORMS_RESPONSES_READONLY_SCOPE,
    SLIDES_SCOPE,
    SLIDES_READONLY_SCOPE,
    TASKS_SCOPE,
    TASKS_READONLY_SCOPE,
    CUSTOM_SEARCH_SCOPE,
)

logger = logging.getLogger(__name__)

# Authentication helper functions
def _get_auth_context(
    tool_name: str,
) -> Tuple[Optional[str], Optional[str], Optional[str]]:
    """
    Get authentication context from FastMCP.

    Returns:
        Tuple of (authenticated_user, auth_method, mcp_session_id)
    """
    try:
        ctx = get_context()
        if not ctx:
            return None, None, None

        authenticated_user = ctx.get_state("authenticated_user_email")
        auth_method = ctx.get_state("authenticated_via")
        mcp_session_id = ctx.session_id if hasattr(ctx, "session_id") else None

        if mcp_session_id:
            set_fastmcp_session_id(mcp_session_id)

        logger.debug(
            f"[{tool_name}] Auth from middleware: {authenticated_user} via {auth_method}"
        )
        return authenticated_user, auth_method, mcp_session_id

    except Exception as e:
        logger.debug(f"[{tool_name}] Could not get FastMCP context: {e}")
        return None, None, None


def _detect_oauth_version(
    authenticated_user: Optional[str], mcp_session_id: Optional[str], tool_name: str
) -> bool:
    """
    Detect whether to use OAuth 2.1 based on configuration and context.

    Returns:
        True if OAuth 2.1 should be used, False otherwise
    """
    if not is_oauth21_enabled():
        return False

    # When OAuth 2.1 is enabled globally, ALWAYS use OAuth 2.1 for authenticated users
    if authenticated_user:
        logger.info(
            f"[{tool_name}] OAuth 2.1 mode: Using OAuth 2.1 for authenticated user '{authenticated_user}'"
        )
        return True

    # Only use version detection for unauthenticated requests
    config = get_oauth_config()
    request_params = {}
    if mcp_session_id:
        request_params["session_id"] = mcp_session_id

    oauth_version = config.detect_oauth_version(request_params)
    use_oauth21 = oauth_version == "oauth21"
    logger.info(
        f"[{tool_name}] OAuth version detected: {oauth_version}, will use OAuth 2.1: {use_oauth21}"
    )
    return use_oauth21


def _update_email_in_args(args: tuple, index: int, new_email: str) -> tuple:
    """Update email at specific index in args tuple."""
    if index < len(args):
        args_list = list(args)
        args_list[index] = new_email
        return tuple(args_list)
    return args


def _override_oauth21_user_email(
    use_oauth21: bool,
    authenticated_user: Optional[str],
    current_user_email: str,
    args: tuple,
    kwargs: dict,
    param_names: List[str],
    tool_name: str,
    service_type: str = "",
) -> Tuple[str, tuple]:
    """
    Override user_google_email with authenticated user when using OAuth 2.1.

    Returns:
        Tuple of (updated_user_email, updated_args)
    """
    if not (use_oauth21 and authenticated_user and current_user_email != authenticated_user):
        return current_user_email, args

    service_suffix = f" for service '{service_type}'" if service_type else ""
    logger.info(
        f"[{tool_name}] OAuth 2.1: Overriding user_google_email from '{current_user_email}' to authenticated user '{authenticated_user}'{service_suffix}"
    )

    # Update in kwargs if present
    if "user_google_email" in kwargs:
        kwargs["user_google_email"] = authenticated_user

    # Update in args if user_google_email is passed positionally
    try:
        user_email_index = param_names.index("user_google_email")
        args = _update_email_in_args(args, user_email_index, authenticated_user)
    except ValueError:
        pass  # user_google_email not in positional parameters

    return authenticated_user, args


async def _authenticate_service(
    use_oauth21: bool,
    service_name: str,
    service_version: str,
    tool_name: str,
    user_google_email: str,
    resolved_scopes: List[str],
    mcp_session_id: Optional[str],
    authenticated_user: Optional[str],
) -> Tuple[Any, str]:
    """
    Authenticate and get Google service using appropriate OAuth version.

    Returns:
        Tuple of (service, actual_user_email)
    """
    if use_oauth21:
        logger.debug(f"[{tool_name}] Using OAuth 2.1 flow")
        return await get_authenticated_google_service_oauth21(
            service_name=service_name,
            version=service_version,
            tool_name=tool_name,
            user_google_email=user_google_email,
            required_scopes=resolved_scopes,
            session_id=mcp_session_id,
            auth_token_email=authenticated_user,
            allow_recent_auth=False,
        )
    else:
        logger.debug(f"[{tool_name}] Using legacy OAuth 2.0 flow")
        return await get_authenticated_google_service(
            service_name=service_name,
            version=service_version,
            tool_name=tool_name,
            user_google_email=user_google_email,
            required_scopes=resolved_scopes,
            session_id=mcp_session_id,
        )


async def get_authenticated_google_service_oauth21(
    service_name: str,
    version: str,
    tool_name: str,
    user_google_email: str,
    required_scopes: List[str],
    session_id: Optional[str] = None,
    auth_token_email: Optional[str] = None,
    allow_recent_auth: bool = False,
) -> tuple[Any, str]:
    """
    OAuth 2.1 authentication using the session store with security validation.
    """
    provider = get_auth_provider()
    access_token = get_access_token()

    if provider and access_token:
        token_email = None
        if getattr(access_token, "claims", None):
            token_email = access_token.claims.get("email")

        resolved_email = token_email or auth_token_email or user_google_email
        if not resolved_email:
            raise GoogleAuthenticationError(
                "Authenticated user email could not be determined from access token."
            )

        if auth_token_email and token_email and token_email != auth_token_email:
            raise GoogleAuthenticationError(
                "Access token email does not match authenticated session context."
            )

        if token_email and user_google_email and token_email != user_google_email:
            raise GoogleAuthenticationError(
                f"Authenticated account {token_email} does not match requested user {user_google_email}."
            )

        credentials = ensure_session_from_access_token(access_token, resolved_email, session_id)
        if not credentials:
            raise GoogleAuthenticationError(
                "Unable to build Google credentials from authenticated access token."
            )

        scopes_available = set(credentials.scopes or [])
        if not scopes_available and getattr(access_token, "scopes", None):
            scopes_available = set(access_token.scopes)

        if not all(scope in scopes_available for scope in required_scopes):
            raise GoogleAuthenticationError(
                f"OAuth credentials lack required scopes. Need: {required_scopes}, Have: {sorted(scopes_available)}"
            )

        service = build(service_name, version, credentials=credentials)
        logger.info(f"[{tool_name}] Authenticated {service_name} for {resolved_email}")
        return service, resolved_email

    store = get_oauth21_session_store()

    # Use the validation method to ensure session can only access its own credentials
    credentials = store.get_credentials_with_validation(
        requested_user_email=user_google_email,
        session_id=session_id,
        auth_token_email=auth_token_email,
        allow_recent_auth=allow_recent_auth,
    )

    if not credentials:
        raise GoogleAuthenticationError(
            f"Access denied: Cannot retrieve credentials for {user_google_email}. "
            f"You can only access credentials for your authenticated account."
        )

    if not credentials.scopes:
        scopes_available = set(required_scopes)
    else:
        scopes_available = set(credentials.scopes)

    if not all(scope in scopes_available for scope in required_scopes):
        raise GoogleAuthenticationError(
            f"OAuth 2.1 credentials lack required scopes. Need: {required_scopes}, Have: {sorted(scopes_available)}"
        )

    service = build(service_name, version, credentials=credentials)
    logger.info(f"[{tool_name}] Authenticated {service_name} for {user_google_email}")

    return service, user_google_email


def _extract_oauth21_user_email(authenticated_user: Optional[str], func_name: str) -> str:
    """
    Extract user email for OAuth 2.1 mode.

    Args:
        authenticated_user: The authenticated user from context
        func_name: Name of the function being decorated (for error messages)

    Returns:
        User email string

    Raises:
        Exception: If no authenticated user found in OAuth 2.1 mode
    """
    if not authenticated_user:
        raise Exception(
            f"OAuth 2.1 mode requires an authenticated user for {func_name}, but none was found."
        )
    return authenticated_user


def _extract_oauth20_user_email(
    args: tuple,
    kwargs: dict,
    wrapper_sig: inspect.Signature
) -> str:
    """
    Extract user email for OAuth 2.0 mode from function arguments.

    Args:
        args: Positional arguments passed to wrapper
        kwargs: Keyword arguments passed to wrapper
        wrapper_sig: Function signature for parameter binding

    Returns:
        User email string

    Raises:
        Exception: If user_google_email parameter not found
    """
    bound_args = wrapper_sig.bind(*args, **kwargs)
    bound_args.apply_defaults()

    user_google_email = bound_args.arguments.get("user_google_email")
    if not user_google_email:
        raise Exception(
            "'user_google_email' parameter is required but was not found."
        )
    return user_google_email



def _remove_user_email_arg_from_docstring(docstring: str) -> str:
    """
    Remove user_google_email parameter documentation from docstring.

    Args:
        docstring: The original function docstring

    Returns:
        Modified docstring with user_google_email parameter removed
    """
    if not docstring:
        return docstring

    # Pattern to match user_google_email parameter documentation
    # Handles various formats like:
    # - user_google_email (str): The user's Google email address. Required.
    # - user_google_email: Description
    # - user_google_email (str) - Description
    patterns = [
        r'^\s*user_google_email\s*\([^)]*\)\s*:\s*[^\n]*\.?\s*(?:Required\.?)?\s*\n',
        r'^\s*user_google_email\s*:\s*[^\n]*\n',
        r'^\s*user_google_email\s*\([^)]*\)\s*-\s*[^\n]*\n',
    ]

    modified_docstring = docstring
    for pattern in patterns:
        modified_docstring = re.sub(pattern, '', modified_docstring, flags=re.MULTILINE)

    # Clean up any sequence of 3 or more newlines that might have been created
    modified_docstring = re.sub(r'\n{3,}', '\n\n', modified_docstring)
    return modified_docstring

# Service configuration mapping
SERVICE_CONFIGS = {
    "gmail": {"service": "gmail", "version": "v1"},
    "drive": {"service": "drive", "version": "v3"},
    "calendar": {"service": "calendar", "version": "v3"},
    "docs": {"service": "docs", "version": "v1"},
    "sheets": {"service": "sheets", "version": "v4"},
    "chat": {"service": "chat", "version": "v1"},
    "forms": {"service": "forms", "version": "v1"},
    "slides": {"service": "slides", "version": "v1"},
    "tasks": {"service": "tasks", "version": "v1"},
    "customsearch": {"service": "customsearch", "version": "v1"},
}


# Scope group definitions for easy reference
SCOPE_GROUPS = {
    # Gmail scopes
    "gmail_read": GMAIL_READONLY_SCOPE,
    "gmail_send": GMAIL_SEND_SCOPE,
    "gmail_compose": GMAIL_COMPOSE_SCOPE,
    "gmail_modify": GMAIL_MODIFY_SCOPE,
    "gmail_labels": GMAIL_LABELS_SCOPE,
    # Drive scopes
    "drive_read": DRIVE_READONLY_SCOPE,
    "drive_file": DRIVE_FILE_SCOPE,
    # Docs scopes
    "docs_read": DOCS_READONLY_SCOPE,
    "docs_write": DOCS_WRITE_SCOPE,
    # Calendar scopes
    "calendar_read": CALENDAR_READONLY_SCOPE,
    "calendar_events": CALENDAR_EVENTS_SCOPE,
    # Sheets scopes
    "sheets_read": SHEETS_READONLY_SCOPE,
    "sheets_write": SHEETS_WRITE_SCOPE,
    # Chat scopes
    "chat_read": CHAT_READONLY_SCOPE,
    "chat_write": CHAT_WRITE_SCOPE,
    "chat_spaces": CHAT_SPACES_SCOPE,
    # Forms scopes
    "forms": FORMS_BODY_SCOPE,
    "forms_read": FORMS_BODY_READONLY_SCOPE,
    "forms_responses_read": FORMS_RESPONSES_READONLY_SCOPE,
    # Slides scopes
    "slides": SLIDES_SCOPE,
    "slides_read": SLIDES_READONLY_SCOPE,
    # Tasks scopes
    "tasks": TASKS_SCOPE,
    "tasks_read": TASKS_READONLY_SCOPE,
    # Custom Search scope
    "customsearch": CUSTOM_SEARCH_SCOPE,
}



def _resolve_scopes(scopes: Union[str, List[str]]) -> List[str]:
    """Resolve scope names to actual scope URLs."""
    if isinstance(scopes, str):
        if scopes in SCOPE_GROUPS:
            return [SCOPE_GROUPS[scopes]]
        else:
            return [scopes]

    resolved = []
    for scope in scopes:
        if scope in SCOPE_GROUPS:
            resolved.append(SCOPE_GROUPS[scope])
        else:
            resolved.append(scope)
    return resolved


def _handle_token_refresh_error(
    error: RefreshError, user_email: str, service_name: str
) -> str:
    """
    Handle token refresh errors gracefully, particularly expired/revoked tokens.

    Args:
        error: The RefreshError that occurred
        user_email: User's email address
        service_name: Name of the Google service

    Returns:
        A user-friendly error message with instructions for reauthentication
    """
    error_str = str(error)

    if (
        "invalid_grant" in error_str.lower()
        or "expired or revoked" in error_str.lower()
    ):
        logger.warning(
            f"Token expired or revoked for user {user_email} accessing {service_name}"
        )


        service_display_name = f"Google {service_name.title()}"

        return (
            f"**Authentication Required: Token Expired/Revoked for {service_display_name}**\n\n"
            f"Your Google authentication token for {user_email} has expired or been revoked. "
            f"This commonly happens when:\n"
            f"- The token has been unused for an extended period\n"
            f"- You've changed your Google account password\n"
            f"- You've revoked access to the application\n\n"
            f"**To resolve this, please:**\n"
            f"1. Run `start_google_auth` with your email ({user_email}) and service_name='{service_display_name}'\n"
            f"2. Complete the authentication flow in your browser\n"
            f"3. Retry your original command\n\n"
            f"The application will automatically use the new credentials once authentication is complete."
        )
    else:
        # Handle other types of refresh errors
        logger.error(f"Unexpected refresh error for user {user_email}: {error}")
        return (
            f"Authentication error occurred for {user_email}. "
            f"Please try running `start_google_auth` with your email and the appropriate service name to reauthenticate."
        )


def require_google_service(
    service_type: str,
    scopes: Union[str, List[str]],
    version: Optional[str] = None,
):
    """
    Decorator that automatically handles Google service authentication and injection.

    Args:
        service_type: Type of Google service ("gmail", "drive", "calendar", etc.)
        scopes: Required scopes (can be scope group names or actual URLs)
        version: Service version (defaults to standard version for service type)

    Usage:
        @require_google_service("gmail", "gmail_read")
        async def search_messages(service, user_google_email: str, query: str):
            # service parameter is automatically injected
            # Original authentication logic is handled automatically
    """

    def decorator(func: Callable) -> Callable:
        original_sig = inspect.signature(func)
        params = list(original_sig.parameters.values())

        # The decorated function must have 'service' as its first parameter.
        if not params or params[0].name != "service":
            raise TypeError(
                f"Function '{func.__name__}' decorated with @require_google_service "
                "must have 'service' as its first parameter."
            )

        # Create a new signature for the wrapper that excludes the 'service' parameter.
        # In OAuth 2.1 mode, also exclude 'user_google_email' since it's automatically determined.
        if is_oauth21_enabled():
            # Remove both 'service' and 'user_google_email' parameters
            filtered_params = [
                p for p in params[1:]
                if p.name != 'user_google_email'
            ]
            wrapper_sig = original_sig.replace(parameters=filtered_params)
        else:
            # Only remove 'service' parameter for OAuth 2.0 mode
            wrapper_sig = original_sig.replace(parameters=params[1:])

        @wraps(func)
        async def wrapper(*args, **kwargs):
            # Note: `args` and `kwargs` are now the arguments for the *wrapper*,
            # which does not include 'service'.

            # Get authentication context early to determine OAuth mode
            authenticated_user, auth_method, mcp_session_id = _get_auth_context(
                func.__name__
            )

            # Extract user_google_email based on OAuth mode
            if is_oauth21_enabled():
                user_google_email = _extract_oauth21_user_email(authenticated_user, func.__name__)
            else:
                user_google_email = _extract_oauth20_user_email(args, kwargs, wrapper_sig)

            # Get service configuration from the decorator's arguments
            if service_type not in SERVICE_CONFIGS:
                raise Exception(f"Unknown service type: {service_type}")

            config = SERVICE_CONFIGS[service_type]
            service_name = config["service"]
            service_version = version or config["version"]

            # Resolve scopes
            resolved_scopes = _resolve_scopes(scopes)

            try:
                tool_name = func.__name__

                # Log authentication status
                logger.debug(
                    f"[{tool_name}] Auth: {authenticated_user or 'none'} via {auth_method or 'none'} (session: {mcp_session_id[:8] if mcp_session_id else 'none'})"
                )

                # Detect OAuth version
                use_oauth21 = _detect_oauth_version(
                    authenticated_user, mcp_session_id, tool_name
                )

                # In OAuth 2.1 mode, user_google_email is already set to authenticated_user
                # In OAuth 2.0 mode, we may need to override it
                if not is_oauth21_enabled():
                    wrapper_params = list(wrapper_sig.parameters.keys())
                    user_google_email, args = _override_oauth21_user_email(
                        use_oauth21,
                        authenticated_user,
                        user_google_email,
                        args,
                        kwargs,
                        wrapper_params,
                        tool_name,
                    )

                # Authenticate service
                service, actual_user_email = await _authenticate_service(
                    use_oauth21,
                    service_name,
                    service_version,
                    tool_name,
                    user_google_email,
                    resolved_scopes,
                    mcp_session_id,
                    authenticated_user,
                )
            except GoogleAuthenticationError as e:
                logger.error(
                    f"[{tool_name}] GoogleAuthenticationError during authentication. "
                    f"Method={auth_method or 'none'}, User={authenticated_user or 'none'}, "
                    f"Service={service_name} v{service_version}, MCPSessionID={mcp_session_id or 'none'}: {e}"
                )
                # Re-raise the original error without wrapping it
                raise

            try:
                # In OAuth 2.1 mode, we need to add user_google_email to kwargs since it was removed from signature
                if is_oauth21_enabled():
                    kwargs["user_google_email"] = user_google_email

                # Prepend the fetched service object to the original arguments
                return await func(service, *args, **kwargs)
            except RefreshError as e:
                error_message = _handle_token_refresh_error(
                    e, actual_user_email, service_name
                )
                raise Exception(error_message)

        # Set the wrapper's signature to the one without 'service'
        wrapper.__signature__ = wrapper_sig

        # Conditionally modify docstring to remove user_google_email parameter documentation
        if is_oauth21_enabled():
            logger.debug('OAuth 2.1 mode enabled, removing user_google_email from docstring')
            if func.__doc__:
                wrapper.__doc__ = _remove_user_email_arg_from_docstring(func.__doc__)

        return wrapper

    return decorator


def require_multiple_services(service_configs: List[Dict[str, Any]]):
    """
    Decorator for functions that need multiple Google services.

    Args:
        service_configs: List of service configurations, each containing:
            - service_type: Type of service
            - scopes: Required scopes
            - param_name: Name to inject service as (e.g., 'drive_service', 'docs_service')
            - version: Optional version override

    Usage:
        @require_multiple_services([
            {"service_type": "drive", "scopes": "drive_read", "param_name": "drive_service"},
            {"service_type": "docs", "scopes": "docs_read", "param_name": "docs_service"}
        ])
        async def get_doc_with_metadata(drive_service, docs_service, user_google_email: str, doc_id: str):
            # Both services are automatically injected
    """

    def decorator(func: Callable) -> Callable:
        original_sig = inspect.signature(func)

        service_param_names = {config["param_name"] for config in service_configs}
        params = list(original_sig.parameters.values())

        # Remove injected service params from the wrapper signature; drop user_google_email only for OAuth 2.1.
        filtered_params = [
            p for p in params
            if p.name not in service_param_names
        ]
        if is_oauth21_enabled():
            filtered_params = [
                p for p in filtered_params
                if p.name != 'user_google_email'
            ]

        wrapper_sig = original_sig.replace(parameters=filtered_params)
        wrapper_param_names = [p.name for p in filtered_params]

        @wraps(func)
        async def wrapper(*args, **kwargs):
            # Get authentication context early
            tool_name = func.__name__
            authenticated_user, _, mcp_session_id = _get_auth_context(tool_name)

            # Extract user_google_email based on OAuth mode
            if is_oauth21_enabled():
                user_google_email = _extract_oauth21_user_email(authenticated_user, tool_name)
            else:
                user_google_email = _extract_oauth20_user_email(args, kwargs, wrapper_sig)

            # Authenticate all services
            for config in service_configs:
                service_type = config["service_type"]
                scopes = config["scopes"]
                param_name = config["param_name"]
                version = config.get("version")

                if service_type not in SERVICE_CONFIGS:
                    raise Exception(f"Unknown service type: {service_type}")

                service_config = SERVICE_CONFIGS[service_type]
                service_name = service_config["service"]
                service_version = version or service_config["version"]
                resolved_scopes = _resolve_scopes(scopes)

                try:
                    # Detect OAuth version (simplified for multiple services)
                    use_oauth21 = (
                        is_oauth21_enabled() and authenticated_user is not None
                    )

                    # In OAuth 2.0 mode, we may need to override user_google_email
                    if not is_oauth21_enabled():
                        user_google_email, args = _override_oauth21_user_email(
                            use_oauth21,
                            authenticated_user,
                            user_google_email,
                            args,
                            kwargs,
                            wrapper_param_names,
                            tool_name,
                            service_type,
                        )

                    # Authenticate service
                    service, _ = await _authenticate_service(
                        use_oauth21,
                        service_name,
                        service_version,
                        tool_name,
                        user_google_email,
                        resolved_scopes,
                        mcp_session_id,
                        authenticated_user,
                    )

                    # Inject service with specified parameter name
                    kwargs[param_name] = service

                except GoogleAuthenticationError as e:
                    logger.error(
                        f"[{tool_name}] GoogleAuthenticationError for service '{service_type}' (user: {user_google_email}): {e}"
                    )
                    # Re-raise the original error without wrapping it
                    raise

            # Call the original function with refresh error handling
            try:
                # In OAuth 2.1 mode, we need to add user_google_email to kwargs since it was removed from signature
                if is_oauth21_enabled():
                    kwargs["user_google_email"] = user_google_email

                return await func(*args, **kwargs)
            except RefreshError as e:
                # Handle token refresh errors gracefully
                error_message = _handle_token_refresh_error(
                    e, user_google_email, "Multiple Services"
                )
                raise Exception(error_message)

        # Set the wrapper's signature
        wrapper.__signature__ = wrapper_sig

        # Conditionally modify docstring to remove user_google_email parameter documentation
        if is_oauth21_enabled():
            logger.debug('OAuth 2.1 mode enabled, removing user_google_email from docstring')
            if func.__doc__:
                wrapper.__doc__ = _remove_user_email_arg_from_docstring(func.__doc__)

        return wrapper

    return decorator


