roa2web-service-auto/backend/modules/telegram/auth/linking.py

"""
Authentication and User Linking Logic

This module handles the linking process between Telegram users and Oracle ERP accounts.
It manages authentication codes, verifies users through the backend API, and maintains
user sessions with JWT tokens.
"""

import logging
from typing import Optional, Dict, Any
from datetime import datetime, timedelta

from telegram import User as TelegramUser

from backend.modules.telegram.db.operations import (
    get_user,
    create_or_update_user,
    link_user_to_oracle,
    update_user_tokens,
    verify_and_use_auth_code,
    is_user_linked
)
from backend.modules.telegram.api.client import get_backend_client

logger = logging.getLogger(__name__)


async def link_telegram_account(
    telegram_user: TelegramUser,
    auth_code: str
) -> Optional[Dict[str, Any]]:
    """
    Link a Telegram account to an Oracle ERP account using an authentication code.

    Flow:
    1. Verify auth code in database (check exists, not used, not expired)
    2. Extract oracle_username from code
    3. Call backend API to verify user in Oracle and get JWT token
    4. Create/update Telegram user record
    5. Link user to Oracle account with JWT tokens
    6. Return success with user data

    Args:
        telegram_user: Telegram User object from python-telegram-bot
        auth_code: 8-character authentication code from web frontend

    Returns:
        Dict with:
        - success: True if linking succeeded
        - username: Oracle username
        - jwt_token: JWT access token
        - companies: List of companies user has access to
        OR None if linking failed

    Example:
        result = await link_telegram_account(telegram_user, "ABC12345")
        if result:
            print(f"Linked to {result['username']}")
        else:
            print("Linking failed")
    """
    try:
        telegram_user_id = telegram_user.id
        telegram_username = telegram_user.username
        first_name = telegram_user.first_name
        last_name = telegram_user.last_name

        logger.info(
            f"Attempting to link Telegram user {telegram_user_id} "
            f"(@{telegram_username}) with code {auth_code}"
        )

        # Step 1: Verify auth code
        code_data = await verify_and_use_auth_code(auth_code)

        if not code_data:
            logger.warning(f"Invalid, expired, or already used auth code: {auth_code}")
            return None

        oracle_username = code_data.get('oracle_username')
        server_id = code_data.get('server_id')  # Extract server_id from the stored code
        logger.info(f"Auth code valid for Oracle user: {oracle_username} (server_id={server_id})")

        # Step 2: Create/update Telegram user record (basic info)
        user_created = await create_or_update_user(
            telegram_user_id=telegram_user_id,
            username=telegram_username,
            first_name=first_name,
            last_name=last_name
        )

        if not user_created:
            logger.error(f"Failed to create/update Telegram user {telegram_user_id}")
            return None

        # Step 3: Verify user in Oracle and get JWT token via backend API (auto-linking flow)
        backend_client = get_backend_client()
        async with backend_client:
            user_data = await backend_client.verify_user(
                oracle_username=oracle_username,
                linking_code=auth_code,
                server_id=server_id
            )

        if not user_data or not user_data.get('success'):
            logger.error(f"Failed to verify Oracle user {oracle_username} via backend")
            return None

        # Extract tokens and user info from response
        jwt_token = user_data.get('access_token')
        jwt_refresh_token = user_data.get('refresh_token', jwt_token)
        user_info = user_data.get('user', {})
        companies = user_info.get('companies', [])
        permissions = user_info.get('permissions', [])

        # Token expiration (typically 30 minutes for access token)
        token_expires_at = datetime.now() + timedelta(minutes=30)

        # Step 4: Link Telegram user to Oracle account
        linked = await link_user_to_oracle(
            telegram_user_id=telegram_user_id,
            oracle_username=oracle_username,
            jwt_token=jwt_token,
            jwt_refresh_token=jwt_refresh_token,
            token_expires_at=token_expires_at
        )

        if not linked:
            logger.error(f"Failed to link user {telegram_user_id} to Oracle account")
            return None

        logger.info(
            f"Successfully linked Telegram user {telegram_user_id} "
            f"to Oracle user {oracle_username}"
        )

        return {
            "success": True,
            "telegram_user_id": telegram_user_id,
            "username": oracle_username,
            "jwt_token": jwt_token,
            "jwt_refresh_token": jwt_refresh_token,
            "companies": companies,
            "permissions": permissions,
            "linked_at": datetime.now().isoformat()
        }

    except Exception as e:
        logger.error(f"Error linking Telegram account: {e}", exc_info=True)
        return None


async def get_user_auth_data(telegram_user_id: int) -> Optional[Dict[str, Any]]:
    """
    Get authentication data for a linked Telegram user.

    This function retrieves the user's Oracle account information and JWT tokens.
    If the token is expired, it automatically refreshes it.

    Args:
        telegram_user_id: Telegram user ID

    Returns:
        Dict with:
        - telegram_user_id: Telegram user ID
        - username: Oracle username
        - jwt_token: Valid JWT access token (refreshed if needed)
        - jwt_refresh_token: JWT refresh token
        - companies: List of companies (fetched if not cached)
        OR None if user is not linked or error occurred

    Example:
        auth_data = await get_user_auth_data(12345)
        if auth_data:
            jwt = auth_data['jwt_token']
            # Use JWT for API calls
    """
    try:
        # Get user from database
        user_data = await get_user(telegram_user_id)

        if not user_data:
            logger.warning(f"User {telegram_user_id} not found in database")
            return None

        if not user_data.get('oracle_username'):
            logger.warning(f"User {telegram_user_id} is not linked to Oracle account")
            return None

        oracle_username = user_data['oracle_username']
        jwt_token = user_data['jwt_token']
        jwt_refresh_token = user_data['jwt_refresh_token']
        token_expires_at_str = user_data['token_expires_at']

        # Parse token expiration
        token_expires_at = datetime.fromisoformat(token_expires_at_str) if token_expires_at_str else None

        # Check if token is expired or about to expire (< 5 minutes remaining)
        token_expired = (
            token_expires_at is None or
            datetime.now() >= token_expires_at - timedelta(minutes=5)
        )

        if token_expired:
            logger.info(f"Token expired for user {telegram_user_id}, refreshing...")

            # Refresh token via backend API
            backend_client = get_backend_client()
            async with backend_client:
                new_token = await backend_client.refresh_token(jwt_refresh_token)

            if new_token:
                # Update token in database
                new_expires_at = datetime.now() + timedelta(minutes=30)
                await update_user_tokens(
                    telegram_user_id=telegram_user_id,
                    jwt_token=new_token,
                    jwt_refresh_token=jwt_refresh_token,  # Keep same refresh token
                    token_expires_at=new_expires_at
                )

                jwt_token = new_token
                logger.info(f"Token refreshed for user {telegram_user_id}")
            else:
                logger.error(f"Failed to refresh token for user {telegram_user_id}")
                return None

        # Fetch user companies (fresh from backend)
        backend_client = get_backend_client()
        async with backend_client:
            companies = await backend_client.get_user_companies(jwt_token)

        return {
            "telegram_user_id": telegram_user_id,
            "username": oracle_username,
            "jwt_token": jwt_token,
            "jwt_refresh_token": jwt_refresh_token,
            "companies": companies
        }

    except Exception as e:
        logger.error(f"Error getting user auth data: {e}", exc_info=True)
        return None


async def check_user_linked(telegram_user_id: int) -> bool:
    """
    Check if a Telegram user is linked to an Oracle account.

    Args:
        telegram_user_id: Telegram user ID

    Returns:
        bool: True if user is linked, False otherwise

    Example:
        if await check_user_linked(12345):
            print("User is linked")
        else:
            print("User needs to link account")
    """
    try:
        return await is_user_linked(telegram_user_id)
    except Exception as e:
        logger.error(f"Error checking if user is linked: {e}")
        return False


async def get_user_companies(telegram_user_id: int) -> Optional[list]:
    """
    Get list of companies a user has access to.

    This is a convenience function that fetches user auth data and returns
    just the companies list.

    Args:
        telegram_user_id: Telegram user ID

    Returns:
        List of company dicts, or None if user not linked

    Example:
        companies = await get_user_companies(12345)
        if companies:
            for company in companies:
                print(f"{company['id']}: {company['nume_firma']}")
    """
    try:
        auth_data = await get_user_auth_data(telegram_user_id)

        if auth_data:
            return auth_data.get('companies', [])

        return None

    except Exception as e:
        logger.error(f"Error getting user companies: {e}")
        return None


async def unlink_user(telegram_user_id: int) -> bool:
    """
    Unlink a Telegram user from their Oracle account.

    This removes the linking but keeps the Telegram user record.
    Used for account disconnection or security purposes.

    Args:
        telegram_user_id: Telegram user ID

    Returns:
        bool: True if successfully unlinked

    Example:
        if await unlink_user(12345):
            print("Account unlinked")
    """
    try:
        # Set Oracle username and tokens to NULL
        from backend.modules.telegram.db.database import DB_PATH
        import aiosqlite

        async with aiosqlite.connect(DB_PATH) as db:
            db.row_factory = aiosqlite.Row
            await db.execute("""
                UPDATE telegram_users
                SET oracle_username = NULL,
                    jwt_token = NULL,
                    jwt_refresh_token = NULL,
                    token_expires_at = NULL,
                    linked_at = NULL
                WHERE telegram_user_id = ?
            """, (telegram_user_id,))

            await db.commit()

        logger.info(f"User {telegram_user_id} unlinked from Oracle account")
        return True

    except Exception as e:
        logger.error(f"Error unlinking user {telegram_user_id}: {e}")
        return False


# Export main functions
__all__ = [
    'link_telegram_account',
    'get_user_auth_data',
    'check_user_linked',
    'get_user_companies',
    'unlink_user'
]