"""
OAuth User Pool Management for Load Testing.

Manages multiple OAuth-authenticated users for realistic multi-user load testing scenarios.
"""

import logging
import secrets
import string
import time
from dataclasses import dataclass
from typing import Any
from urllib.parse import quote

import anyio
import httpx
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

logger = logging.getLogger(__name__)


@dataclass
class UserConfig:
    """Configuration for a single test user."""

    username: str
    password: str
    display_name: str
    email: str
    groups: list[str]


@dataclass
class UserProfile:
    """Profile for an OAuth-authenticated user."""

    username: str
    password: str
    token: str
    session: ClientSession | None = None
    streamable_context: Any | None = None  # Store for proper cleanup
    operation_count: int = 0
    error_count: int = 0


class OAuthUserPool:
    """
    Manages a pool of OAuth-authenticated users for load testing.

    Handles token acquisition, session management, and user lifecycle.
    """

    def __init__(
        self,
        admin_client: Any,  # NextcloudClient with admin credentials
        client_id: str,
        client_secret: str,
        callback_url: str,
        token_endpoint: str,
        authorization_endpoint: str,
    ):
        self.admin_client = admin_client  # For user management
        self.nextcloud_host = str(admin_client._client.base_url)
        self.client_id = client_id
        self.client_secret = client_secret
        self.callback_url = callback_url
        self.token_endpoint = token_endpoint
        self.authorization_endpoint = authorization_endpoint
        self.users: dict[str, UserProfile] = {}
        self._http_client: httpx.AsyncClient | None = None

    async def __aenter__(self):
        """Initialize HTTP client."""
        self._http_client = httpx.AsyncClient(verify=False, timeout=30.0)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Cleanup HTTP client."""
        if self._http_client:
            await self._http_client.aclose()

    async def acquire_token(self, username: str, password: str, auth_code: str) -> str:
        """
        Exchange authorization code for OAuth access token.

        Args:
            username: Username for logging
            password: Password (for logging/debugging)
            auth_code: Authorization code from OAuth flow

        Returns:
            OAuth access token
        """
        logger.info(f"Exchanging auth code for access token (user: {username})...")

        if not self._http_client:
            raise RuntimeError(
                "HTTP client not initialized - use async context manager"
            )

        # Exchange authorization code for access token
        token_response = await self._http_client.post(
            self.token_endpoint,
            data={
                "grant_type": "authorization_code",
                "code": auth_code,
                "redirect_uri": self.callback_url,
                "client_id": self.client_id,
                "client_secret": self.client_secret,
            },
            headers={"Content-Type": "application/x-www-form-urlencoded"},
        )
        token_response.raise_for_status()
        token_data = token_response.json()

        access_token = token_data.get("access_token")
        if not access_token:
            raise ValueError(f"No access token in response for {username}")

        logger.info(f"Successfully acquired OAuth token for {username}")
        return access_token

    async def add_user(self, username: str, password: str, token: str) -> UserProfile:
        """
        Add a user to the pool with their OAuth token.

        Args:
            username: Username
            password: Password (for future re-auth if needed)
            token: OAuth access token

        Returns:
            UserProfile for the added user
        """
        if username in self.users:
            logger.warning(f"User {username} already in pool, updating token")

        profile = UserProfile(username=username, password=password, token=token)
        self.users[username] = profile
        logger.info(f"Added user {username} to pool (total: {len(self.users)})")
        return profile

    async def create_user_session(
        self, username: str, mcp_url: str = "http://localhost:8001/mcp"
    ) -> ClientSession:
        """
        Create an MCP client session for a user.

        Args:
            username: Username to create session for
            mcp_url: MCP server URL

        Returns:
            Initialized ClientSession

        Raises:
            KeyError: If user not in pool
        """
        if username not in self.users:
            raise KeyError(f"User {username} not in pool")

        profile = self.users[username]

        # Create streamable HTTP connection with OAuth token in Authorization header
        # This matches the pattern from tests/conftest.py create_mcp_client_session()
        headers = {"Authorization": f"Bearer {profile.token}"}
        streamable_context = streamablehttp_client(mcp_url, headers=headers)

        try:
            read_stream, write_stream, _ = await streamable_context.__aenter__()

            session = ClientSession(read_stream, write_stream)
            await session.__aenter__()
            await session.initialize()

            # Store both session and context for proper cleanup
            profile.session = session
            profile.streamable_context = streamable_context
            logger.info(f"Created MCP session for {username}")
            return session

        except Exception as e:
            # Clean up streamable context if session creation failed
            try:
                await streamable_context.__aexit__(None, None, None)
            except Exception as cleanup_error:
                logger.debug(f"Error during cleanup: {cleanup_error}")
            raise e

    async def close_user_session(self, username: str):
        """Close the MCP session for a user."""
        if username not in self.users:
            return

        profile = self.users[username]

        # Close ClientSession
        if profile.session:
            try:
                await profile.session.__aexit__(None, None, None)
            except Exception as e:
                logger.debug(f"Error closing session for {username}: {e}")
            profile.session = None

        # Close streamable context
        if profile.streamable_context:
            try:
                await profile.streamable_context.__aexit__(None, None, None)
            except Exception as e:
                logger.debug(f"Error closing streamable context for {username}: {e}")
            profile.streamable_context = None

    async def close_all_sessions(self):
        """Close all user sessions."""
        for username in list(self.users.keys()):
            await self.close_user_session(username)

    def get_user(self, username: str) -> UserProfile:
        """Get user profile by username."""
        if username not in self.users:
            raise KeyError(f"User {username} not in pool")
        return self.users[username]

    def get_all_users(self) -> list[UserProfile]:
        """Get all user profiles."""
        return list(self.users.values())

    def record_operation(self, username: str, success: bool = True):
        """Record an operation for user stats."""
        if username in self.users:
            self.users[username].operation_count += 1
            if not success:
                self.users[username].error_count += 1

    def get_stats(self) -> dict[str, dict[str, int | float]]:
        """Get per-user operation statistics."""
        return {
            username: {
                "operations": profile.operation_count,
                "errors": profile.error_count,
                "success_rate": (
                    (profile.operation_count - profile.error_count)
                    / max(profile.operation_count, 1)
                    * 100
                ),
            }
            for username, profile in self.users.items()
        }

    async def create_nextcloud_user(
        self,
        username: str,
        password: str,
        display_name: str | None = None,
        email: str | None = None,
    ) -> UserConfig:
        """
        Create a Nextcloud user via the Users API.

        Args:
            username: Username for the new user
            password: Password for the new user
            display_name: Optional display name
            email: Optional email address

        Returns:
            UserConfig for the created user

        Raises:
            HTTPStatusError: If user creation fails
        """
        logger.info(f"Creating Nextcloud user: {username}")

        await self.admin_client.users.create_user(
            userid=username,
            password=password,
            display_name=display_name or username,
            email=email or f"{username}@benchmark.local",
        )

        logger.info(f"Successfully created Nextcloud user: {username}")

        return UserConfig(
            username=username,
            password=password,
            display_name=display_name or username,
            email=email or f"{username}@benchmark.local",
            groups=[],
        )

    async def delete_nextcloud_user(self, username: str):
        """
        Delete a Nextcloud user via the Users API.

        Args:
            username: Username to delete
        """
        logger.info(f"Deleting Nextcloud user: {username}")

        try:
            await self.admin_client.users.delete_user(userid=username)
            logger.info(f"Successfully deleted Nextcloud user: {username}")
        except Exception as e:
            logger.warning(f"Failed to delete user {username}: {e}")

    async def acquire_token_playwright(
        self,
        browser: Any,
        username: str,
        password: str,
        state: str,
        auth_states: dict[str, str],
    ) -> str:
        """
        Acquire OAuth token via Playwright browser automation.

        Based on conftest.py playwright_oauth_token fixture.
        Automates the full OAuth flow:
        1. Navigate to authorization URL
        2. Fill login form
        3. Handle OAuth consent
        4. Wait for callback server to receive auth code
        5. Exchange code for access token

        Args:
            browser: Playwright browser instance
            username: Username to authenticate
            password: Password for the user
            state: Unique state parameter for this OAuth flow
            auth_states: Dict mapping state -> auth_code (shared with callback server)

        Returns:
            OAuth access token

        Raises:
            TimeoutError: If callback not received within timeout
            ValueError: If token exchange fails
        """

        logger.info(f"Starting Playwright OAuth flow for {username}...")
        logger.debug(f"Using state: {state[:16]}...")

        # Construct authorization URL
        auth_url = (
            f"{self.authorization_endpoint}?"
            f"response_type=code&"
            f"client_id={self.client_id}&"
            f"redirect_uri={quote(self.callback_url, safe='')}&"
            f"state={state}&"
            f"scope=openid%20profile%20email"
        )

        # Browser automation
        context = await browser.new_context(ignore_https_errors=True)
        page = await context.new_page()

        try:
            # Navigate to authorization URL
            logger.debug("Navigating to authorization URL...")
            await page.goto(auth_url, wait_until="networkidle", timeout=30000)
            current_url = page.url

            # Login if needed
            if "/login" in current_url or "/index.php/login" in current_url:
                logger.info(f"Logging in as {username}...")
                await page.wait_for_selector('input[name="user"]', timeout=10000)
                await page.fill('input[name="user"]', username)
                await page.fill('input[name="password"]', password)
                await page.click('button[type="submit"]')
                await page.wait_for_load_state("networkidle", timeout=30000)
                current_url = page.url
                logger.info("Login completed")

            # Handle OAuth consent if present
            try:
                authorize_button = await page.query_selector(
                    'button:has-text("Authorize"), button:has-text("Allow"), input[type="submit"][value*="uthoriz"]'
                )
                if authorize_button:
                    logger.info("Authorizing OAuth client...")
                    await authorize_button.click()
                    await page.wait_for_load_state("networkidle", timeout=10000)
            except Exception as e:
                logger.debug(f"No authorization needed: {e}")

            # Wait for callback server to receive auth code
            logger.info("Waiting for OAuth callback...")
            timeout_seconds = 30
            start_time = time.time()
            while state not in auth_states:
                if time.time() - start_time > timeout_seconds:
                    screenshot_path = f"/tmp/oauth_timeout_{username}.png"
                    await page.screenshot(path=screenshot_path)
                    logger.error(f"Screenshot saved to {screenshot_path}")
                    raise TimeoutError(
                        f"Timeout waiting for OAuth callback for {username}"
                    )
                await anyio.sleep(0.5)

            auth_code = auth_states[state]
            logger.info(f"Received auth code for {username}")

        finally:
            await context.close()

        # Exchange code for token
        logger.info(f"Exchanging auth code for access token ({username})...")
        token_response = await self._http_client.post(
            self.token_endpoint,
            data={
                "grant_type": "authorization_code",
                "code": auth_code,
                "redirect_uri": self.callback_url,
                "client_id": self.client_id,
                "client_secret": self.client_secret,
            },
            headers={"Content-Type": "application/x-www-form-urlencoded"},
        )
        token_response.raise_for_status()
        token_data = token_response.json()

        access_token = token_data.get("access_token")
        if not access_token:
            raise ValueError(f"No access token for {username}: {token_data}")

        logger.info(f"Successfully acquired OAuth token for {username}")
        return access_token


class UserSessionWrapper:
    """
    Wrapper for a user-specific MCP session with operation tracking.

    Provides a convenient interface for executing operations as a specific user.
    """

    def __init__(self, username: str, session: ClientSession, pool: OAuthUserPool):
        self.username = username
        self.session = session
        self.pool = pool

    async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any:
        """
        Call an MCP tool and record the operation.

        Args:
            tool_name: Name of the tool to call
            arguments: Tool arguments

        Returns:
            Tool result
        """
        try:
            result = await self.session.call_tool(tool_name, arguments)
            self.pool.record_operation(self.username, success=True)
            return result
        except Exception:
            self.pool.record_operation(self.username, success=False)
            raise

    async def read_resource(self, uri: str) -> Any:
        """
        Read an MCP resource and record the operation.

        Args:
            uri: Resource URI

        Returns:
            Resource data
        """
        try:
            result = await self.session.read_resource(uri)
            self.pool.record_operation(self.username, success=True)
            return result
        except Exception:
            self.pool.record_operation(self.username, success=False)
            raise


def generate_secure_password(length: int = 20) -> str:
    """Generate a secure random password."""

    alphabet = string.ascii_letters + string.digits + "!@#$%^&*()"
    return "".join(secrets.choice(alphabet) for _ in range(length))