"""
Uses broker-based authentication handler with platform-specific support.

This module provides functionality to acquire access tokens using MSAL Python
with a platform-appropriate broker:
  - Windows 10+: Windows Account Manager (WAM)
  - macOS: Company Portal broker
  - Linux: Intune broker

For more information, see:
https://learn.microsoft.com/en-us/entra/msal/python/advanced/wam
https://github.com/AzureAD/microsoft-authentication-extensions-for-python
"""

import os
import platform
from typing import Callable, Optional
from pathlib import Path
import jwt
from msal import PublicClientApplication
from msal_extensions import PersistedTokenCache, build_encrypted_persistence

from cli_logging.cli_logger import emit_structured_log
from cli_logging.logging_utils import Operation

APP_DIR_NAME = ".m365-copilot-agent-evals"
DEFAULT_AUTH_TIMEOUT_SECONDS = 120


class AuthHandler:
    """Handler for platform-specific broker-based authentication (WAM on Windows, Company Portal on macOS, Intune on Linux)."""

    def __init__(
        self,
        client_id: str,
        tenant_id: str,
        scopes_str: str,
        cache_dir: Optional[str] = None,
        auth_timeout: int = DEFAULT_AUTH_TIMEOUT_SECONDS,
        account_hint: Optional[str] = None,
    ):
        """
        Initialize the auth handler.

        Args:
            client_id: App registration client ID (required).
            tenant_id: Directory/tenant ID (required).
            scopes_str: Comma-separated scopes (required).
            cache_dir: Optional directory for token cache file. Defaults to ~/.m365-copilot-agent-evals.
            auth_timeout: Seconds to wait for interactive auth before giving up. Defaults to 120.
            account_hint: Optional account (email/UPN) to scope authentication to a
                specific cached MSAL account. When set, silent acquisition is limited
                to matching accounts and interactive auth pre-fills the account picker
                via ``login_hint``. When None, all cached accounts are tried (default).
        """
        if not client_id:
            raise ValueError("client_id is required")

        if not tenant_id:
            raise ValueError("tenant_id is required")

        if not scopes_str:
            raise ValueError("scopes_str is required")

        scopes = [s.strip() for s in scopes_str.split(",") if s.strip()]

        self.client_id = client_id
        self.authority = f"https://login.microsoftonline.com/{tenant_id}"
        self.scopes = scopes
        self._auth_timeout = auth_timeout
        self.account_hint = (account_hint or "").strip() or None

        # Initialize the public client application with platform-appropriate broker.
        # See: https://learn.microsoft.com/en-us/entra/msal/python/advanced/linux-broker-py?tabs=ubuntudep#how-to-opt-in-to-use-broker
        self._current_os = platform.system()
        broker_kwargs = {}
        if self._current_os == "Windows":
            broker_kwargs["enable_broker_on_windows"] = True
        elif self._current_os == "Darwin":
            broker_kwargs["enable_broker_on_mac"] = True
        elif self._current_os == "Linux":
            broker_kwargs["enable_broker_on_linux"] = True

        try:
            self.app = PublicClientApplication(
                client_id=self.client_id,
                authority=self.authority,
                token_cache=self._setup_token_cache(cache_dir),
                **broker_kwargs
            )
            emit_structured_log(
                "info",
                f"Auth handler initialized successfully (OS: {self._current_os})",
                operation=Operation.AUTHENTICATE,
            )
        except OSError as e:
            if self._current_os == "Linux":
                emit_structured_log(
                    "error",
                    f"Linux broker failed to load a required system library: {e}. "
                    "Install the missing dependencies and retry:\n"
                    "  sudo apt install libwebkit2gtk-4.1-0 libdbus-1-dev "
                    "python3-gi gir1.2-secret-1 libubsan1",
                    operation=Operation.AUTHENTICATE,
                )
            else:
                emit_structured_log(
                    "error",
                    f"Failed to initialize auth handler: {e}",
                    operation=Operation.AUTHENTICATE,
                )
            raise
        except ImportError as e:
            emit_structured_log(
                "error",
                f"Broker dependencies not installed: {e}. "
                "Install with: pip install 'msal[broker]'",
                operation=Operation.AUTHENTICATE,
            )
            raise
        except Exception as e:
            emit_structured_log(
                "error",
                f"Failed to initialize auth handler: {e}",
                operation=Operation.AUTHENTICATE,
            )
            raise

    def _setup_token_cache(
        self, cache_dir: Optional[str] = None
    ) -> Optional[PersistedTokenCache]:
        """
        Setup encrypted persistent token cache using MSAL Extensions.

        Creates a platform-dependent encrypted cache (DPAPI on Windows,
        Keychain on macOS, Libsecret on Linux). Falls back to an in-memory
        session cache if the encrypted store is unavailable (e.g. libsecret /
        PyGObject missing on WSL); tokens will not persist across process
        restarts in that case.

        Args:
            cache_dir: Optional directory for token cache file. Defaults to ~/.m365-copilot-agent-evals.

        Returns:
            PersistedTokenCache if available, None for in-memory session cache.
        """
        if cache_dir is None:
            cache_dir = str(Path.home() / APP_DIR_NAME)

        cache_path = os.path.join(cache_dir, "token_cache.bin")

        try:
            # Create cache directory if it doesn't exist
            Path(cache_dir).mkdir(parents=True, exist_ok=True)

            # Build encrypted persistence (DPAPI on Windows, Keychain on Mac, Libsecret on Linux)
            persistence = build_encrypted_persistence(cache_path)
            token_cache = PersistedTokenCache(persistence)
            emit_structured_log(
                "info",
                f"Encrypted token cache initialized at {cache_path} (encrypted: {persistence.is_encrypted})",
                operation=Operation.AUTHENTICATE,
            )
            return token_cache
        except Exception as e:
            emit_structured_log(
                "error",
                f"Could not initialize persistent token cache: {e}. "
                "Falling back to in-memory session cache; "
                "tokens will not persist across sessions.",
                operation=Operation.AUTHENTICATE,
            )
            return None

    def acquire_token_interactive(self):
        """
        Acquire a token interactively using the platform broker and return the full MSAL result dict.

        This method first attempts to retrieve a cached token. If no valid cached
        token is found or if it has expired, it will prompt the user to authenticate
        via the platform broker (WAM on Windows, Company Portal on macOS, Intune on Linux).

        Returns:
            The full result dictionary from MSAL (contains access_token, id_token, etc.)
            or None if acquisition fails.

        Raises:
            Exception: If broker communication fails or user denies the request.
        """
        try:
            # Attempt to acquire a token silently from cache. When an account
            # hint is set, MSAL's get_accounts(username=...) scopes the lookup to
            # matching accounts (case-insensitive UPN match); otherwise all
            # cached accounts are returned.
            silent_accounts = self.get_accounts(username=self.account_hint)

            if silent_accounts:
                emit_structured_log(
                    "info",
                    f"Attempting silent token acquisition from {len(silent_accounts)} cached account(s)",
                    operation=Operation.AUTHENTICATE,
                )
                for account in silent_accounts:
                    silent_result = self.app.acquire_token_silent(
                        scopes=self.scopes, account=account
                    )
                    if silent_result and "access_token" in silent_result:
                        emit_structured_log(
                            "info",
                            "Access token acquired successfully from cache",
                            operation=Operation.AUTHENTICATE,
                        )
                        return silent_result
                emit_structured_log(
                    "info",
                    "No valid cached token found; proceeding with interactive authentication",
                    operation=Operation.AUTHENTICATE,
                )
            elif self.account_hint:
                emit_structured_log(
                    "info",
                    f"No cached account matched '{self.account_hint}'; proceeding with "
                    "interactive authentication (account pre-filled via login_hint)",
                    operation=Operation.AUTHENTICATE,
                )
            else:
                emit_structured_log(
                    "info",
                    "No cached accounts found; proceeding with interactive authentication",
                    operation=Operation.AUTHENTICATE,
                )

            interactive_kwargs = {
                "scopes": self.scopes,
                "parent_window_handle": self.app.CONSOLE_WINDOW_HANDLE,
                "timeout": self._auth_timeout,
            }
            # Pre-fill the broker/WAM account picker with the requested identity.
            if self.account_hint:
                interactive_kwargs["login_hint"] = self.account_hint

            result = self.app.acquire_token_interactive(**interactive_kwargs)

            if result and "access_token" in result:
                emit_structured_log(
                    "info",
                    "Access token acquired successfully via interactive authentication",
                    operation=Operation.AUTHENTICATE,
                )
                return result

            error_msg = None
            if result:
                error_msg = result.get(
                    "error_description", result.get("error", "Unknown error")
                )
            emit_structured_log(
                "error",
                f"Failed to acquire token. {error_msg or 'Unknown error'}",
                operation=Operation.AUTHENTICATE,
            )
            return result

        except Exception as e:
            emit_structured_log(
                "error",
                f"Error during token acquisition: {e}",
                operation=Operation.AUTHENTICATE,
            )
            raise

    def acquire_token_silent(self, account: Optional[dict] = None) -> Optional[str]:
        """
        Acquire an access token silently (without user interaction).

        This is useful for cached tokens or when you have a known account.
        Returns None if a valid token is not available and interaction would be required.

        Args:
            account: Optional account object from previous interactive authentication.
                    If None, attempts to get a cached token for any account.

        Returns:
            The access token string if successful, None if no cached token available.
        """
        try:
            result = self.app.acquire_token_silent(scopes=self.scopes, account=account)

            if result and "access_token" in result:
                emit_structured_log(
                    "info",
                    "Access token acquired silently",
                    operation=Operation.AUTHENTICATE,
                )
                return result["access_token"]

            emit_structured_log(
                "debug",
                "No cached token available; interactive acquisition would be required",
                operation=Operation.AUTHENTICATE,
            )
            return None

        except Exception as e:
            emit_structured_log(
                "debug",
                f"Silent token acquisition failed: {e}",
                operation=Operation.AUTHENTICATE,
            )
            return None

    def get_accounts(self, username: Optional[str] = None) -> list:
        """
        Get the list of accounts available in WAM cache.

        Args:
            username: Optional account (email/UPN) to filter by. MSAL performs a
                case-insensitive match on the account username. When None, all
                cached accounts are returned.

        Returns:
            List of account objects cached by WAM (filtered by username when set).
        """
        try:
            accounts = self.app.get_accounts(username=username)
            emit_structured_log(
                "info",
                f"Found {len(accounts)} cached account(s)",
                operation=Operation.AUTHENTICATE,
            )
            return accounts
        except Exception as e:
            emit_structured_log(
                "error",
                f"Error retrieving cached accounts: {e}",
                operation=Operation.AUTHENTICATE,
            )
            return []

    def clear_cache(self) -> bool:
        """
        Clear all cached tokens and accounts from the token cache.

        This method removes all cached authentication data, forcing the user
        to authenticate interactively on the next acquire_token_interactive call.

        Returns:
            True if cache was successfully cleared, False otherwise.
        """
        try:
            # Get all cached accounts
            accounts = self.app.get_accounts()

            # Remove each account from the cache
            for account in accounts:
                self.app.remove_account(account)
                emit_structured_log(
                    "info",
                    f"Removed account {account.get('username', 'unknown')} from cache",
                    operation=Operation.AUTHENTICATE,
                )

            emit_structured_log(
                "info",
                "Token cache cleared successfully",
                operation=Operation.AUTHENTICATE,
            )
            return True
        except Exception as e:
            emit_structured_log(
                "error",
                f"Error clearing token cache: {e}",
                operation=Operation.AUTHENTICATE,
            )
            return False

    @staticmethod
    def extract_user_oid_from_access_token(access_token: str) -> str:
        """
        Extract the user OID from an access token using MSAL's JWT decoding.

        MSAL includes PyJWT under the hood, so we can use jwt.decode
        directly without adding new dependencies.

        Args:
            access_token: The access token string

        Returns:
            The user OID if found, empty string otherwise

        Raises:
            ValueError: If token format is invalid or OID not found
        """
        try:
            # Decode without verification (we're just reading claims, not validating signature)
            decoded = jwt.decode(access_token, options={"verify_signature": False})
            oid = decoded.get('oid', '')
            if not oid:
                raise ValueError("OID not found in token claims")
            return oid
        except jwt.DecodeError as e:
            raise ValueError(f"Failed to decode token: {e}")


def make_token_refresh_fn(auth_handler: "AuthHandler") -> Callable[[], str]:
    """Return a callable that silently refreshes the A2A access token.

    On a 401 response the caller invokes this function.  It first attempts a
    silent refresh (using the MSAL refresh token) and falls back to interactive
    authentication only when a silent refresh is not possible.  The returned
    string is the new access token; an empty string signals failure.

    Args:
        auth_handler: An initialized AuthHandler instance to use for token acquisition.

    Returns:
        A zero-argument callable that returns a fresh access token string.
    """
    def _refresh() -> str:
        result = auth_handler.acquire_token_interactive() or {}
        return result.get("access_token") or ""
    return _refresh