import base64
import json
import uuid
import warnings
from datetime import datetime, timezone
from typing import Annotated, Literal, Optional
from urllib.parse import quote
from uuid import UUID as parse_uuid

import posthog
from fastapi import APIRouter, Header, HTTPException, Request, Response, status
from fastapi.responses import JSONResponse, RedirectResponse
from pydantic import SecretStr
from server.auth.auth_utils import user_verifier
from server.auth.constants import (
    KEYCLOAK_CLIENT_ID,
    KEYCLOAK_REALM_NAME,
    KEYCLOAK_SERVER_URL_EXT,
    RECAPTCHA_SITE_KEY,
    ROLE_CHECK_ENABLED,
)
from server.auth.domain_blocker import domain_blocker
from server.auth.gitlab_sync import schedule_gitlab_repo_sync
from server.auth.recaptcha_service import recaptcha_service
from server.auth.saas_user_auth import SaasUserAuth
from server.auth.token_manager import TokenManager
from server.config import sign_token
from server.constants import IS_FEATURE_ENV
from server.routes.event_webhook import _get_session_api_key, _get_user_id
from server.services.org_invitation_service import (
    EmailMismatchError,
    InvitationExpiredError,
    InvitationInvalidError,
    OrgInvitationService,
    UserAlreadyMemberError,
)
from storage.database import session_maker
from storage.user import User
from storage.user_store import UserStore

from openhands.core.logger import openhands_logger as logger
from openhands.integrations.provider import ProviderHandler
from openhands.integrations.service_types import ProviderType, TokenResponse
from openhands.server.services.conversation_service import create_provider_tokens_object
from openhands.server.shared import config
from openhands.server.user_auth import get_access_token
from openhands.server.user_auth.user_auth import get_user_auth

with warnings.catch_warnings():
    warnings.simplefilter('ignore')

api_router = APIRouter(prefix='/api')
oauth_router = APIRouter(prefix='/oauth')

token_manager = TokenManager()


def set_response_cookie(
    request: Request,
    response: Response,
    keycloak_access_token: str,
    keycloak_refresh_token: str,
    secure: bool = True,
    accepted_tos: bool = False,
):
    # Create a signed JWT token
    cookie_data = {
        'access_token': keycloak_access_token,
        'refresh_token': keycloak_refresh_token,
        'accepted_tos': accepted_tos,
    }
    signed_token = sign_token(cookie_data, config.jwt_secret.get_secret_value())  # type: ignore

    # Set secure cookie with signed token
    domain = get_cookie_domain(request)
    if domain:
        response.set_cookie(
            key='keycloak_auth',
            value=signed_token,
            domain=domain,
            httponly=True,
            secure=secure,
            samesite=get_cookie_samesite(request),
        )
    else:
        response.set_cookie(
            key='keycloak_auth',
            value=signed_token,
            httponly=True,
            secure=secure,
            samesite=get_cookie_samesite(request),
        )


def get_cookie_domain(request: Request) -> str | None:
    # for now just use the full hostname except for staging stacks.
    return (
        None
        if not request.url.hostname
        or request.url.hostname.endswith('staging.all-hands.dev')
        else request.url.hostname
    )


def get_cookie_samesite(request: Request) -> Literal['lax', 'strict']:
    # for localhost and feature/staging stacks we set it to 'lax' as the cookie domain won't allow 'strict'
    return (
        'lax'
        if request.url.hostname == 'localhost'
        or (request.url.hostname or '').endswith('staging.all-hands.dev')
        else 'strict'
    )


def _extract_oauth_state(state: str | None) -> tuple[str, str | None, str | None]:
    """Extract redirect URL, reCAPTCHA token, and invitation token from OAuth state.

    Returns:
        Tuple of (redirect_url, recaptcha_token, invitation_token).
        Tokens may be None.
    """
    if not state:
        return '', None, None

    try:
        # Try to decode as JSON (new format with reCAPTCHA and/or invitation)
        state_data = json.loads(base64.urlsafe_b64decode(state.encode()).decode())
        return (
            state_data.get('redirect_url', ''),
            state_data.get('recaptcha_token'),
            state_data.get('invitation_token'),
        )
    except Exception:
        # Old format - state is just the redirect URL
        return state, None, None


# Keep alias for backward compatibility
def _extract_recaptcha_state(state: str | None) -> tuple[str, str | None]:
    """Extract redirect URL and reCAPTCHA token from OAuth state.

    Deprecated: Use _extract_oauth_state instead.

    Returns:
        Tuple of (redirect_url, recaptcha_token). Token may be None.
    """
    redirect_url, recaptcha_token, _ = _extract_oauth_state(state)
    return redirect_url, recaptcha_token


@oauth_router.get('/keycloak/callback')
async def keycloak_callback(
    request: Request,
    code: Optional[str] = None,
    state: Optional[str] = None,
    error: Optional[str] = None,
    error_description: Optional[str] = None,
):
    # Extract redirect URL, reCAPTCHA token, and invitation token from state
    redirect_url, recaptcha_token, invitation_token = _extract_oauth_state(state)
    if not redirect_url:
        redirect_url = str(request.base_url)

    if not code:
        # check if this is a forward from the account linking page
        if (
            error == 'temporarily_unavailable'
            and error_description == 'authentication_expired'
        ):
            return RedirectResponse(redirect_url, status_code=302)
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={'error': 'Missing code in request params'},
        )
    scheme = 'http' if request.url.hostname == 'localhost' else 'https'
    redirect_uri = f'{scheme}://{request.url.netloc}{request.url.path}'
    logger.debug(f'code: {code}, redirect_uri: {redirect_uri}')

    (
        keycloak_access_token,
        keycloak_refresh_token,
    ) = await token_manager.get_keycloak_tokens(code, redirect_uri)
    if not keycloak_access_token or not keycloak_refresh_token:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={'error': 'Problem retrieving Keycloak tokens'},
        )

    user_info = await token_manager.get_user_info(keycloak_access_token)
    logger.debug(f'user_info: {user_info}')
    if ROLE_CHECK_ENABLED and 'roles' not in user_info:
        return JSONResponse(
            status_code=status.HTTP_401_UNAUTHORIZED,
            content={'error': 'Missing required role'},
        )

    if 'sub' not in user_info or 'preferred_username' not in user_info:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={'error': 'Missing user ID or username in response'},
        )

    email = user_info.get('email')
    user_id = user_info['sub']
    user = await UserStore.get_user_by_id_async(user_id)
    if not user:
        user = await UserStore.create_user(user_id, user_info)
    else:
        # Existing user — gradually backfill contact_name if it still has a username-style value
        await UserStore.backfill_contact_name(user_id, user_info)
        await UserStore.backfill_user_email(user_id, user_info)

    if not user:
        logger.error(f'Failed to authenticate user {user_info["preferred_username"]}')
        return JSONResponse(
            status_code=status.HTTP_401_UNAUTHORIZED,
            content={
                'error': f'Failed to authenticate user {user_info["preferred_username"]}'
            },
        )

    logger.info(f'Logging in user {str(user.id)} in org {user.current_org_id}')

    # reCAPTCHA verification with Account Defender
    if RECAPTCHA_SITE_KEY:
        if not recaptcha_token:
            logger.warning(
                'recaptcha_token_missing',
                extra={
                    'user_id': user_id,
                    'email': email,
                },
            )
            error_url = f'{request.base_url}login?recaptcha_blocked=true'
            return RedirectResponse(error_url, status_code=302)

        user_ip = request.client.host if request.client else 'unknown'
        user_agent = request.headers.get('User-Agent', '')

        # Handle X-Forwarded-For for proxied requests
        forwarded_for = request.headers.get('X-Forwarded-For')
        if forwarded_for:
            user_ip = forwarded_for.split(',')[0].strip()

        try:
            result = recaptcha_service.create_assessment(
                token=recaptcha_token,
                action='LOGIN',
                user_ip=user_ip,
                user_agent=user_agent,
                email=email,
                user_id=user_id,
            )

            if not result.allowed:
                logger.warning(
                    'recaptcha_blocked_at_callback',
                    extra={
                        'user_ip': user_ip,
                        'score': result.score,
                        'user_id': user_id,
                    },
                )
                # Redirect to home with error parameter
                error_url = f'{request.base_url}login?recaptcha_blocked=true'
                return RedirectResponse(error_url, status_code=302)

        except Exception as e:
            logger.exception(f'reCAPTCHA verification error at callback: {e}')
            # Fail open - continue with login if reCAPTCHA service unavailable

    # Check if email domain is blocked
    if email and await domain_blocker.is_domain_blocked(email):
        logger.warning(
            f'Blocked authentication attempt for email: {email}, user_id: {user_id}'
        )

        # Disable the Keycloak account
        await token_manager.disable_keycloak_user(user_id, email)

        return JSONResponse(
            status_code=status.HTTP_401_UNAUTHORIZED,
            content={
                'error': 'Access denied: Your email domain is not allowed to access this service'
            },
        )

    # Check for duplicate email with + modifier
    if email:
        try:
            has_duplicate = await token_manager.check_duplicate_base_email(
                email, user_id
            )
            if has_duplicate:
                logger.warning(
                    f'Blocked signup attempt for email {email} - duplicate base email found',
                    extra={'user_id': user_id, 'email': email},
                )

                # Delete the Keycloak user that was automatically created during OAuth
                # This prevents orphaned accounts in Keycloak
                # The delete_keycloak_user method already handles all errors internally
                deletion_success = await token_manager.delete_keycloak_user(user_id)
                if deletion_success:
                    logger.info(
                        f'Deleted Keycloak user {user_id} after detecting duplicate email {email}'
                    )
                else:
                    logger.warning(
                        f'Failed to delete Keycloak user {user_id} after detecting duplicate email {email}. '
                        f'User may need to be manually cleaned up.'
                    )

                # Redirect to home page with query parameter indicating the issue
                home_url = f'{request.base_url}/login?duplicated_email=true'
                return RedirectResponse(home_url, status_code=302)
        except Exception as e:
            # Log error but allow signup to proceed (fail open)
            logger.error(
                f'Error checking duplicate email for {email}: {e}',
                extra={'user_id': user_id, 'email': email},
            )

    # Check email verification status
    email_verified = user_info.get('email_verified', False)
    if not email_verified:
        # Send verification email
        # Import locally to avoid circular import with email.py
        from server.routes.email import verify_email

        await verify_email(request=request, user_id=user_id, is_auth_flow=True)
        verification_redirect_url = f'{request.base_url}login?email_verification_required=true&user_id={user_id}'
        # Preserve invitation token so it can be included in OAuth state after verification
        if invitation_token:
            verification_redirect_url = (
                f'{verification_redirect_url}&invitation_token={invitation_token}'
            )
        response = RedirectResponse(verification_redirect_url, status_code=302)
        return response

    # default to github IDP for now.
    # TODO: remove default once Keycloak is updated universally with the new attribute.
    idp: str = user_info.get('identity_provider', ProviderType.GITHUB.value)
    logger.info(f'Full IDP is {idp}')
    idp_type = 'oidc'
    if ':' in idp:
        idp, idp_type = idp.rsplit(':', 1)
        idp_type = idp_type.lower()

    await token_manager.store_idp_tokens(
        ProviderType(idp), user_id, keycloak_access_token
    )

    username = user_info['preferred_username']
    if user_verifier.is_active() and not user_verifier.is_user_allowed(username):
        return JSONResponse(
            status_code=status.HTTP_401_UNAUTHORIZED,
            content={'error': 'Not authorized via waitlist'},
        )

    valid_offline_token = (
        await token_manager.validate_offline_token(user_id=user_info['sub'])
        if idp_type != 'saml'
        else True
    )

    logger.debug(
        f'keycloakAccessToken: {keycloak_access_token}, keycloakUserId: {user_id}'
    )

    # adding in posthog tracking

    # If this is a feature environment, add "FEATURE_" prefix to user_id for PostHog
    posthog_user_id = f'FEATURE_{user_id}' if IS_FEATURE_ENV else user_id

    try:
        posthog.set(
            distinct_id=posthog_user_id,
            properties={
                'user_id': posthog_user_id,
                'original_user_id': user_id,
                'is_feature_env': IS_FEATURE_ENV,
            },
        )
    except Exception as e:
        logger.error(
            'auth:posthog_set:failed',
            extra={
                'user_id': user_id,
                'error': str(e),
            },
        )
        # Continue execution as this is not critical

    logger.info(
        'user_logged_in',
        extra={
            'idp': idp,
            'idp_type': idp_type,
            'posthog_user_id': posthog_user_id,
            'is_feature_env': IS_FEATURE_ENV,
        },
    )

    if not valid_offline_token:
        redirect_url = (
            f'{KEYCLOAK_SERVER_URL_EXT}/realms/{KEYCLOAK_REALM_NAME}/protocol/openid-connect/auth'
            f'?client_id={KEYCLOAK_CLIENT_ID}&response_type=code'
            f'&kc_idp_hint={idp}'
            f'&redirect_uri={scheme}%3A%2F%2F{request.url.netloc}%2Foauth%2Fkeycloak%2Foffline%2Fcallback'
            f'&scope=openid%20email%20profile%20offline_access'
            f'&state={state}'
        )

    has_accepted_tos = user.accepted_tos is not None

    # Process invitation token if present (after email verification but before TOS)
    if invitation_token:
        try:
            logger.info(
                'Processing invitation token during auth callback',
                extra={
                    'user_id': user_id,
                    'invitation_token_prefix': invitation_token[:10] + '...',
                },
            )

            await OrgInvitationService.accept_invitation(
                invitation_token, parse_uuid(user_id)
            )
            logger.info(
                'Invitation accepted during auth callback',
                extra={'user_id': user_id},
            )

        except InvitationExpiredError:
            logger.warning(
                'Invitation expired during auth callback',
                extra={'user_id': user_id},
            )
            # Add query param to redirect URL
            if '?' in redirect_url:
                redirect_url = f'{redirect_url}&invitation_expired=true'
            else:
                redirect_url = f'{redirect_url}?invitation_expired=true'

        except InvitationInvalidError as e:
            logger.warning(
                'Invalid invitation during auth callback',
                extra={'user_id': user_id, 'error': str(e)},
            )
            if '?' in redirect_url:
                redirect_url = f'{redirect_url}&invitation_invalid=true'
            else:
                redirect_url = f'{redirect_url}?invitation_invalid=true'

        except UserAlreadyMemberError:
            logger.info(
                'User already member during invitation acceptance',
                extra={'user_id': user_id},
            )
            if '?' in redirect_url:
                redirect_url = f'{redirect_url}&already_member=true'
            else:
                redirect_url = f'{redirect_url}?already_member=true'

        except EmailMismatchError as e:
            logger.warning(
                'Email mismatch during auth callback invitation acceptance',
                extra={'user_id': user_id, 'error': str(e)},
            )
            if '?' in redirect_url:
                redirect_url = f'{redirect_url}&email_mismatch=true'
            else:
                redirect_url = f'{redirect_url}?email_mismatch=true'

        except Exception as e:
            logger.exception(
                'Unexpected error processing invitation during auth callback',
                extra={'user_id': user_id, 'error': str(e)},
            )
            # Don't fail the login if invitation processing fails
            if '?' in redirect_url:
                redirect_url = f'{redirect_url}&invitation_error=true'
            else:
                redirect_url = f'{redirect_url}?invitation_error=true'

    # If the user hasn't accepted the TOS, redirect to the TOS page
    if not has_accepted_tos:
        encoded_redirect_url = quote(redirect_url, safe='')
        tos_redirect_url = (
            f'{request.base_url}accept-tos?redirect_url={encoded_redirect_url}'
        )
        if invitation_token:
            tos_redirect_url = f'{tos_redirect_url}&invitation_success=true'
        response = RedirectResponse(tos_redirect_url, status_code=302)
    else:
        if invitation_token:
            redirect_url = f'{redirect_url}&invitation_success=true'
        response = RedirectResponse(redirect_url, status_code=302)

    set_response_cookie(
        request=request,
        response=response,
        keycloak_access_token=keycloak_access_token,
        keycloak_refresh_token=keycloak_refresh_token,
        secure=True if scheme == 'https' else False,
        accepted_tos=has_accepted_tos,
    )

    # Sync GitLab repos & set up webhooks
    # Use Keycloak access token (first-time users lack offline token at this stage)
    # Normally, offline token is used to fetch GitLab token via user_id
    schedule_gitlab_repo_sync(user_id, SecretStr(keycloak_access_token))
    return response


@oauth_router.get('/keycloak/offline/callback')
async def keycloak_offline_callback(code: str, state: str, request: Request):
    if not code:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={'error': 'Missing code in request params'},
        )
    scheme = 'https'
    if request.url.hostname == 'localhost':
        scheme = 'http'
    redirect_uri = f'{scheme}://{request.url.netloc}{request.url.path}'
    logger.debug(f'code: {code}, redirect_uri: {redirect_uri}')

    (
        keycloak_access_token,
        keycloak_refresh_token,
    ) = await token_manager.get_keycloak_tokens(code, redirect_uri)
    if not keycloak_access_token or not keycloak_refresh_token:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={'error': 'Problem retrieving Keycloak tokens'},
        )

    user_info = await token_manager.get_user_info(keycloak_access_token)
    logger.debug(f'user_info: {user_info}')
    if 'sub' not in user_info:
        return JSONResponse(
            status_code=status.HTTP_400_BAD_REQUEST,
            content={'error': 'Missing Keycloak ID in response'},
        )

    await token_manager.store_offline_token(
        user_id=user_info['sub'], offline_token=keycloak_refresh_token
    )

    redirect_url, _, _ = _extract_oauth_state(state)
    return RedirectResponse(
        redirect_url if redirect_url else request.base_url, status_code=302
    )


@oauth_router.get('/github/callback')
async def github_dummy_callback(request: Request):
    """Callback for GitHub that just forwards the user to the app base URL."""
    return RedirectResponse(request.base_url, status_code=302)


@api_router.post('/authenticate')
async def authenticate(request: Request):
    try:
        await get_access_token(request)
        return JSONResponse(
            status_code=status.HTTP_200_OK, content={'message': 'User authenticated'}
        )
    except Exception:
        # For any error during authentication, clear the auth cookie and return 401
        response = JSONResponse(
            status_code=status.HTTP_401_UNAUTHORIZED,
            content={'error': 'User is not authenticated'},
        )

        # Delete the auth cookie if it exists
        keycloak_auth_cookie = request.cookies.get('keycloak_auth')
        if keycloak_auth_cookie:
            response.delete_cookie(
                key='keycloak_auth',
                domain=get_cookie_domain(request),
                samesite=get_cookie_samesite(request),
            )

        return response


@api_router.post('/accept_tos')
async def accept_tos(request: Request):
    user_auth: SaasUserAuth = await get_user_auth(request)
    access_token = await user_auth.get_access_token()
    refresh_token = user_auth.refresh_token
    user_id = await user_auth.get_user_id()

    if not access_token or not refresh_token or not user_id:
        logger.warning(
            f'accept_tos: One or more is None: access_token {access_token}, refresh_token {refresh_token}, user_id {user_id}'
        )
        return JSONResponse(
            status_code=status.HTTP_401_UNAUTHORIZED,
            content={'error': 'User is not authenticated'},
        )

    # Get redirect URL from request body
    body = await request.json()
    redirect_url = body.get('redirect_url', str(request.base_url))

    # Update user settings with TOS acceptance
    accepted_tos: datetime = datetime.now(timezone.utc)
    with session_maker() as session:
        user = session.query(User).filter(User.id == uuid.UUID(user_id)).first()
        if not user:
            session.rollback()
            logger.error('User for {user_id} not found.')
            return JSONResponse(
                status_code=status.HTTP_401_UNAUTHORIZED,
                content={'error': 'User does not exist'},
            )
        user.accepted_tos = accepted_tos
        session.commit()

        logger.info(f'User {user_id} accepted TOS')

    response = JSONResponse(
        status_code=status.HTTP_200_OK, content={'redirect_url': redirect_url}
    )

    set_response_cookie(
        request=request,
        response=response,
        keycloak_access_token=access_token.get_secret_value(),
        keycloak_refresh_token=refresh_token.get_secret_value(),
        secure=False if request.url.hostname == 'localhost' else True,
        accepted_tos=True,
    )
    return response


@api_router.post('/logout')
async def logout(request: Request):
    # Always create the response object first to ensure we can return it even if errors occur
    response = JSONResponse(
        status_code=status.HTTP_200_OK,
        content={'message': 'User logged out'},
    )

    # Always delete the cookie regardless of what happens
    response.delete_cookie(
        key='keycloak_auth',
        domain=get_cookie_domain(request),
        samesite=get_cookie_samesite(request),
    )

    # Try to properly logout from Keycloak, but don't fail if it doesn't work
    try:
        user_auth: SaasUserAuth = await get_user_auth(request)
        if user_auth and user_auth.refresh_token:
            refresh_token = user_auth.refresh_token.get_secret_value()
            await token_manager.logout(refresh_token)
    except Exception as e:
        # Log any errors but don't fail the request
        logger.debug(f'Error during logout: {str(e)}')
        # We still want to clear the cookie and return success

    return response


@api_router.get('/refresh-tokens', response_model=TokenResponse)
async def refresh_tokens(
    request: Request,
    provider: ProviderType,
    sid: str,
    x_session_api_key: Annotated[str | None, Header(alias='X-Session-API-Key')],
) -> TokenResponse:
    """Return the latest token for a given provider."""
    user_id = _get_user_id(sid)
    session_api_key = await _get_session_api_key(user_id, sid)
    if session_api_key != x_session_api_key:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail='Forbidden')

    logger.info(f'Refreshing token for conversation {sid}')
    provider_handler = ProviderHandler(
        create_provider_tokens_object([provider]), external_auth_id=user_id
    )
    service = provider_handler.get_service(provider)
    token = await service.get_latest_token()
    if not token:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"No token found for provider '{provider}'",
        )

    return TokenResponse(token=token.get_secret_value())