"""
Gen8i MCP Server - OAuth 2.0 Authenticated MCP Server

This server provides three tools:
1. login - Initiates browser-based login flow
2. get_use_case - Retrieves available use cases for authenticated user
3. query - Executes queries within a specific use case

Authentication Flow:
- User configures CLIENT_ID in claude_config.json
- Tools check for valid token
- If token missing/expired, tools prompt for login
- login tool opens browser for secure authentication
- Tokens stored securely: access token in memory, refresh token in OS keyring
"""

import asyncio
import json
import keyring
import sys
import webbrowser
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional

from collections import defaultdict

import os

import httpx
from fastapi.responses import HTMLResponse
from mcp.server.fastmcp import FastMCP
from pydantic import Field
from pydantic import BaseModel as _BaseModel
from starlette.applications import Starlette
from starlette.routing import Route
import uvicorn

def _mask(value: str, visible: int = 4) -> str:
    """Mask sensitive string, showing only the first `visible` characters."""
    if not value:
        return "***"
    if len(value) <= visible:
        return "*" * len(value)
    return value[:visible] + "*" * (len(value) - visible)


_SECURITY_HEADERS = {
    "Content-Security-Policy": "default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'",
    "X-Content-Type-Options": "nosniff",
    "X-Frame-Options": "DENY",
    "X-XSS-Protection": "1; mode=block",
}


def _secure_html(content: str, status_code: int = 200) -> HTMLResponse:
    """Return an HTMLResponse with security headers."""
    return HTMLResponse(content, status_code=status_code, headers=_SECURITY_HEADERS)


class RateLimiter:
    """Simple in-memory rate limiter for login attempts."""

    def __init__(self, max_attempts: int = 5, window_seconds: int = 300):
        self.max_attempts = max_attempts
        self.window_seconds = window_seconds
        self._attempts: dict[str, list[datetime]] = defaultdict(list)

    def is_allowed(self, identifier: str) -> bool:
        now = datetime.now()
        cutoff = now - timedelta(seconds=self.window_seconds)
        self._attempts[identifier] = [
            t for t in self._attempts[identifier] if t > cutoff
        ]
        if len(self._attempts[identifier]) >= self.max_attempts:
            return False
        self._attempts[identifier].append(now)
        return True


# Set UTF-8 encoding for Windows console
if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stderr.reconfigure(encoding="utf-8")
    except Exception as e:
        print(f"[ERROR] Failed to set console encoding: {e}", file=sys.stderr)
        pass


# ============================================================================
# Configuration
# ============================================================================


def _require_env(name: str) -> str:
    """Read a required value from environment variables."""
    value = os.environ.get(name, "")
    if not value:
        print(f"[ERROR] {name} environment variable is required.", file=sys.stderr)
        sys.exit(1)
    return value


class ServerConfig(_BaseModel):
    """Server configuration from environment variables"""

    CLIENT_ID: str = Field(
        default_factory=lambda: _require_env("CLIENT_ID"),
        description="Client ID for authentication",
    )

    # API endpoints
    API_BASE_URL: str = Field(
        default="https://mcp.gen8i.com", description="Gen8i MCP API base URL"
    )
    # OAuth endpoints (PKCE flow)
    AUTHORIZE_ENDPOINT: str = Field(
        default="/api/v1/oauth/authorize", description="OAuth authorization endpoint"
    )
    TOKEN_ENDPOINT: str = Field(
        default="/api/v1/oauth/token", description="OAuth token exchange endpoint"
    )
    LOGIN_ENDPOINT: str = Field(
        default="/api/v1/auth/login", description="Login endpoint"
    )
    # Legacy endpoints (for refresh)
    REFRESH_ENDPOINT: str = Field(
        default="/api/v1/auth/refresh", description="Token refresh endpoint"
    )
    USE_CASES_ENDPOINT: str = Field(
        default="/api/v1/usecases", description="Use cases endpoint"
    )
    QUERY_ENDPOINT: str = Field(default="/api/v1/query", description="Query endpoint")

    # Local server settings
    CALLBACK_PORT: int = Field(
        default=8899, description="Port for OAuth callback server"
    )

    # Token storage
    TOKEN_FILE: Path = Field(
        default=Path.home() / ".g8_mcp_auth_token.json",
        description="Token storage file",
    )


# ============================================================================
# PKCE Manager
# ============================================================================


class PKCEManager:
    """Manages PKCE (Proof Key for Code Exchange) parameters for OAuth flow"""

    @staticmethod
    def generate_code_verifier(length: int = 64) -> str:
        """
        Generate a cryptographically random code verifier.

        Args:
            length: Length of verifier (43-128 characters, default 64)

        Returns:
            URL-safe base64 encoded random string
        """
        import secrets
        import base64

        if not 43 <= length <= 128:
            raise ValueError("Code verifier length must be between 43 and 128")

        # Generate random bytes
        random_bytes = secrets.token_bytes(length)

        # Encode as URL-safe base64 (no padding)
        verifier = base64.urlsafe_b64encode(random_bytes).decode("utf-8").rstrip("=")

        # Truncate to desired length
        return verifier[:length]

    @staticmethod
    def generate_code_challenge(verifier: str) -> str:
        """
        Generate PKCE code challenge from verifier using SHA256.

        Args:
            verifier: Code verifier string

        Returns:
            URL-safe base64 encoded SHA256 hash (no padding)
        """
        import hashlib
        import base64

        # Create SHA256 hash of verifier
        digest = hashlib.sha256(verifier.encode("utf-8")).digest()

        # Encode as URL-safe base64 (no padding)
        challenge = base64.urlsafe_b64encode(digest).decode("utf-8").rstrip("=")

        return challenge

    @staticmethod
    def generate_state() -> str:
        """
        Generate random state parameter for CSRF protection.

        Returns:
            Random URL-safe string
        """
        import secrets

        return secrets.token_urlsafe(32)

    @staticmethod
    def validate_state(received: str, expected: str) -> bool:
        """
        Validate state parameter to prevent CSRF attacks.

        Args:
            received: State from callback
            expected: Original state we generated

        Returns:
            True if states match, False otherwise
        """
        import secrets

        return secrets.compare_digest(received, expected)


# ============================================================================
# Token Manager
# ============================================================================


class TokenManager:
    """Manages token storage, retrieval, and validation"""

    SERVICE_NAME = "g8-mcp-server"

    def __init__(self, config: ServerConfig):
        self.config = config
        self.token_file = config.TOKEN_FILE
        self._current_token: Optional[str] = None
        self._token_expiry: Optional[datetime] = None
        self._refresh_token: Optional[str] = None
        self._refresh_expiry: Optional[datetime] = None

    def save_token(
        self,
        access_token: str,
        refresh_token: Optional[str] = None,
        expires_in: int = 3600,
        refresh_expires_in: int = 2592000,
        token_type: str = "bearer",
    ):
        """
        Save tokens securely - ONLY refresh token in keyring, access token in memory.

        Security improvements:
        - Access token stored ONLY in memory (not persisted)
        - Refresh token stored in OS keyring
        - Metadata file has restrictive permissions (0600)
        """
        # Store access token in MEMORY ONLY
        self._current_token = access_token
        self._token_expiry = datetime.now() + timedelta(seconds=expires_in)
        self._refresh_token = refresh_token
        self._refresh_expiry = (
            datetime.now() + timedelta(seconds=refresh_expires_in)
            if refresh_token
            else None
        )

        print("[OK] Access token stored in memory (not persisted)", file=sys.stderr)

        # Store ONLY refresh token in OS keyring (secure vault)
        if refresh_token:
            try:
                keyring.set_password(self.SERVICE_NAME, "refresh_token", refresh_token)
                print("[OK] Refresh token saved to OS Secure Vault", file=sys.stderr)
            except Exception as e:
                print(
                    f"[ERROR] Failed to save refresh token to keyring: {e}",
                    file=sys.stderr,
                )

        # Store only metadata in JSON file (NO TOKENS)
        data = {
            "token_type": token_type,
            "access_expires_at": self._token_expiry.isoformat(),
            "refresh_expires_at": self._refresh_expiry.isoformat()
            if self._refresh_expiry
            else None,
            "client_id": self.config.CLIENT_ID,
        }

        self.token_file.parent.mkdir(parents=True, exist_ok=True)
        self.token_file.write_text(json.dumps(data, indent=2))

        # Set restrictive file permissions
        self._set_secure_file_permissions(self.token_file)

        print(f"[OK] Token metadata saved to {self.token_file}", file=sys.stderr)

    def load_token(self) -> Optional[str]:
        """
        Load token - refresh token from keyring, access token from memory if valid.

        Returns:
            Access token if valid and in memory, None otherwise
        """
        # If access token in memory and still valid, return it
        if self._current_token and self._token_expiry:
            if datetime.now() < self._token_expiry:
                print(
                    "[OK] Using valid access token from memory",
                    file=sys.stderr,
                )
                return self._current_token
            else:
                print("[WARN] Access token in memory expired", file=sys.stderr)
                self._current_token = None  # Clear expired token

        # Check metadata file
        if not self.token_file.exists():
            return None

        try:
            data = json.loads(self.token_file.read_text())

            # Check if token is for the same client
            if data.get("client_id") != self.config.CLIENT_ID:
                print("[WARN] Token is for different client, ignoring", file=sys.stderr)
                return None

            # Parse expiry times
            access_expires_at = datetime.fromisoformat(data["access_expires_at"])
            self._token_expiry = access_expires_at

            if data.get("refresh_expires_at"):
                self._refresh_expiry = datetime.fromisoformat(
                    data["refresh_expires_at"]
                )

            # Load ONLY refresh token from keyring
            try:
                self._refresh_token = keyring.get_password(
                    self.SERVICE_NAME, "refresh_token"
                )
            except Exception as e:
                print(
                    f"[ERROR] Failed to retrieve refresh token from keyring: {e}",
                    file=sys.stderr,
                )
                return None

            if not self._refresh_token:
                print(
                    "[WARN] Metadata exists but refresh token missing from keyring",
                    file=sys.stderr,
                )
                return None

            # Access token not in memory and expired - need to refresh
            print(
                "[WARN] Access token not in memory, will attempt refresh",
                file=sys.stderr,
            )
            return None
        except Exception as e:
            print(f"[WARN] Error loading token: {e}", file=sys.stderr)
            return None

    def get_refresh_token(self) -> Optional[str]:
        """Get refresh token if available"""
        return self._refresh_token

    def _set_secure_file_permissions(self, file_path: Path):
        """
        Set restrictive file permissions on token metadata file.

        - Unix/Mac: 0600 (owner read/write only)
        - Windows: Remove all permissions except owner
        """
        import platform

        try:
            system = platform.system()

            if system in ["Linux", "Darwin"]:  # Unix/Mac
                import os
                import stat

                # Set to 0600 (rw-------)
                os.chmod(file_path, stat.S_IRUSR | stat.S_IWUSR)
                print(f"[OK] Set file permissions to 0600 (Unix)", file=sys.stderr)

            elif system == "Windows":
                import subprocess

                # Use icacls to set restrictive permissions
                # Remove inheritance and grant only current user full control
                try:
                    # Get current username
                    username = os.environ.get("USERNAME")
                    if not username:
                        print(
                            "[WARN] Could not get username for Windows permissions",
                            file=sys.stderr,
                        )
                        return

                    # Disable inheritance and remove all permissions
                    subprocess.run(
                        ["icacls", str(file_path), "/inheritance:r"],
                        check=True,
                        capture_output=True,
                    )

                    # Grant only current user full control
                    subprocess.run(
                        ["icacls", str(file_path), f"/grant:r", f"{username}:F"],
                        check=True,
                        capture_output=True,
                    )
                    print(
                        f"[OK] Set restrictive permissions (Windows)", file=sys.stderr
                    )

                except subprocess.CalledProcessError as e:
                    print(
                        f"[WARN] Failed to set Windows permissions: {e}",
                        file=sys.stderr,
                    )
            else:
                print(
                    f"[WARN] Unknown OS: {system}, skipping permission setting",
                    file=sys.stderr,
                )

        except Exception as e:
            print(f"[WARN] Failed to set file permissions: {e}", file=sys.stderr)

    def get_token(self) -> Optional[str]:
        """Get current valid token (no auto-refresh - use get_valid_token for that)"""
        # Check memory first
        if self._current_token and self._token_expiry:
            if datetime.now() < self._token_expiry:
                return self._current_token

        # Try loading from file
        return self.load_token()

    async def get_valid_token(self) -> Optional[str]:
        """Get valid token, automatically refresh if expired"""
        # Check if access token is valid
        current = self.get_token()
        if current:
            return current

        # Access token expired or missing, try refresh
        if self._refresh_token and self._refresh_expiry:
            if datetime.now() < self._refresh_expiry:
                print(
                    "[INFO] Access token expired, attempting refresh...",
                    file=sys.stderr,
                )
                try:
                    refreshed = await self.refresh_access_token()
                    if refreshed:
                        return self._current_token
                except Exception as e:
                    print(f"[WARN] Token refresh failed: {e}", file=sys.stderr)

        # Both tokens invalid or refresh failed
        print("[INFO] No valid tokens available", file=sys.stderr)
        return None

    async def refresh_access_token(self) -> bool:
        """
        Refresh access token using refresh token.

        Handles refresh token rotation:
        - Saves new refresh token to keyring
        - Stores new access token in memory only
        - Updates metadata file
        """
        if not self._refresh_token:
            print("[ERROR] No refresh token available", file=sys.stderr)
            return False

        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(
                    f"{self.config.API_BASE_URL}{self.config.REFRESH_ENDPOINT}",
                    json={
                        "refresh_token": self._refresh_token,
                        "client_id": self.config.CLIENT_ID,
                    },
                    timeout=30.0,
                )

                if response.status_code == 200:
                    result = response.json()
                    access_token = result.get("access_token")
                    new_refresh_token = result.get(
                        "refresh_token"
                    )  # NEW: Handle rotation
                    token_type = result.get("token_type", "bearer")
                    expires_in = result.get("expires_in", 3600)
                    refresh_expires_in = result.get("refresh_expires_in", 2592000)

                    if access_token:
                        # Update access token IN MEMORY ONLY (not keyring)
                        self._current_token = access_token
                        self._token_expiry = datetime.now() + timedelta(
                            seconds=expires_in
                        )
                        print(
                            "[OK] Access token refreshed (in memory)", file=sys.stderr
                        )

                        # Handle refresh token rotation
                        if new_refresh_token:
                            print(
                                "[OK] Refresh token rotated, saving new token",
                                file=sys.stderr,
                            )
                            self._refresh_token = new_refresh_token
                            self._refresh_expiry = datetime.now() + timedelta(
                                seconds=refresh_expires_in
                            )

                            # Save NEW refresh token to keyring
                            try:
                                keyring.set_password(
                                    self.SERVICE_NAME,
                                    "refresh_token",
                                    new_refresh_token,
                                )
                                print(
                                    "[OK] New refresh token saved to keyring",
                                    file=sys.stderr,
                                )
                            except Exception as e:
                                print(
                                    f"[ERROR] Failed to save new refresh token: {e}",
                                    file=sys.stderr,
                                )

                        # Update metadata file
                        data = {
                            "token_type": token_type,
                            "access_expires_at": self._token_expiry.isoformat(),
                            "refresh_expires_at": self._refresh_expiry.isoformat()
                            if self._refresh_expiry
                            else None,
                            "client_id": self.config.CLIENT_ID,
                        }
                        self.token_file.write_text(json.dumps(data, indent=2))

                        # Set secure permissions
                        self._set_secure_file_permissions(self.token_file)

                        print(
                            "✅ Token refresh complete with rotation!", file=sys.stderr
                        )
                        return True
                    else:
                        print(
                            "[ERROR] No access_token in refresh response",
                            file=sys.stderr,
                        )
                        return False
                else:
                    print(
                        f"[ERROR] Token refresh failed: {response.status_code}",
                        file=sys.stderr,
                    )
                    return False

        except Exception as e:
            print(f"[ERROR] Error refreshing token: {e}", file=sys.stderr)
            return False

    def clear_token(self):
        """Clear all tokens from memory and file"""
        self._current_token = None
        self._token_expiry = None
        self._refresh_token = None
        self._refresh_expiry = None
        if self.token_file.exists():
            self.token_file.unlink()

        try:
            keyring.delete_password(self.SERVICE_NAME, "access_token")
            keyring.delete_password(self.SERVICE_NAME, "refresh_token")
        except Exception:
            pass  # Ignore if they don't exist

        print("[CLEAR] All tokens cleared from Vault and file", file=sys.stderr)


# ============================================================================
# Login Flow Server
# ============================================================================


class LoginFlowServer:
    """Handles browser-based login flow with Gen8i API integration"""

    def __init__(self, config: ServerConfig, token_manager: TokenManager):
        self.config = config
        self.token_manager = token_manager
        self.received_token: Optional[dict] = (
            None  # Store full response with access and refresh tokens
        )
        self.login_complete = asyncio.Event()

        # PKCE parameters (generated during login flow)
        self.code_verifier: Optional[str] = None
        self.state: Optional[str] = None
        self._rate_limiter = RateLimiter(max_attempts=5, window_seconds=300)

    async def handle_login_page(self, request):
        """Show login form for user to enter credentials"""

        # Validate client credentials are configured
        if not self.config.CLIENT_ID:
            return _secure_html(
                """
                <!DOCTYPE html>
                <html>
                <head>
                    <title>Configuration Error</title>
                    <style>
                        body {
                            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
                            background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
                            display: flex;
                            justify-content: center;
                            align-items: center;
                            min-height: 100vh;
                        }
                        .container {
                            background: white;
                            padding: 3rem;
                            border-radius: 1rem;
                            box-shadow: 0 20px 60px rgba(0,0,0,0.3);
                            text-align: center;
                        }
                        h1 { color: #e53e3e; }
                    </style>
                </head>
                <body>
                    <div class="container">
                        <h1>❌ Configuration Error</h1>
                        <p>CLIENT_ID must be configured in Claude Desktop config.</p>
                        <p style="font-size: 0.9rem; color: #666;">Configure CLIENT_ID in claude_config.json.</p>
                    </div>
                </body>
                </html>
            """,
                status_code=400,
            )

        # Return HTML login form
        return _secure_html(f"""
        <!DOCTYPE html>
        <html>
        <head>
            <title>Gen8i Login</title>
            <style>
                * {{
                    margin: 0;
                    padding: 0;
                    box-sizing: border-box;
                }}
                body {{
                    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                    display: flex;
                    justify-content: center;
                    align-items: center;
                    min-height: 100vh;
                    padding: 20px;
                }}
                .container {{
                    background: white;
                    padding: 3rem;
                    border-radius: 1rem;
                    box-shadow: 0 20px 60px rgba(0,0,0,0.3);
                    max-width: 400px;
                    width: 100%;
                }}
                h1 {{
                    color: #2d3748;
                    margin-bottom: 0.5rem;
                    font-size: 2rem;
                }}
                .subtitle {{
                    color: #718096;
                    margin-bottom: 2rem;
                    font-size: 0.9rem;
                }}
                .form-group {{
                    margin-bottom: 1.5rem;
                }}
                label {{
                    display: block;
                    color: #4a5568;
                    font-weight: 600;
                    margin-bottom: 0.5rem;
                    font-size: 0.9rem;
                }}
                input {{
                    width: 100%;
                    padding: 0.75rem;
                    border: 2px solid #e2e8f0;
                    border-radius: 0.5rem;
                    font-size: 1rem;
                    transition: border-color 0.2s;
                }}
                input:focus {{
                    outline: none;
                    border-color: #667eea;
                }}
                button {{
                    width: 100%;
                    padding: 0.75rem;
                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                    color: white;
                    border: none;
                    border-radius: 0.5rem;
                    font-size: 1rem;
                    font-weight: 600;
                    cursor: pointer;
                    transition: transform 0.2s, box-shadow 0.2s;
                }}
                button:hover {{
                    transform: translateY(-2px);
                    box-shadow: 0 10px 20px rgba(102, 126, 234, 0.4);
                }}
                button:active {{
                    transform: translateY(0);
                }}
                button:disabled {{
                    opacity: 0.6;
                    cursor: not-allowed;
                }}
                .client-info {{
                    margin-top: 1rem;
                    padding: 0.75rem;
                    background: #edf2f7;
                    border-radius: 0.5rem;
                    font-size: 0.85rem;
                    color: #4a5568;
                }}
                .error {{
                    background: #fee;
                    color: #c53030;
                    padding: 0.75rem;
                    border-radius: 0.5rem;
                    margin-bottom: 1rem;
                    display: none;
                }}
                .loading {{
                    display: none;
                    text-align: center;
                    margin-top: 1rem;
                    color: #4a5568;
                }}
            </style>
        </head>
        <body>
            <div class="container">
                <h1>🔐 Gen8i Login</h1>
                <p class="subtitle">Enter your credentials to continue</p>
                
                <div class="client-info">
                    <strong>Client ID:</strong> {self.config.CLIENT_ID}
                </div>
                
                <div id="error" class="error"></div>
                
                <form id="loginForm">
                    <div class="form-group">
                        <label for="username">Email</label>
                        <input type="email" id="username" name="username" required autofocus placeholder="jane@example.com">
                    </div>
                    
                    <div class="form-group">
                        <label for="password">Password</label>
                        <input type="password" id="password" name="password" required>
                    </div>
                    
                    <button type="submit" id="submitBtn">Login</button>
                </form>
                
                <div id="loading" class="loading">
                    <p>🔄 Authenticating...</p>
                </div>
            </div>
            
            <script>
                document.getElementById('loginForm').addEventListener('submit', async (e) => {{
                    e.preventDefault();
                    
                    const submitBtn = document.getElementById('submitBtn');
                    const errorDiv = document.getElementById('error');
                    const loadingDiv = document.getElementById('loading');
                    const username = document.getElementById('username').value;
                    const password = document.getElementById('password').value;
                    
                    // Show loading
                    submitBtn.disabled = true;
                    errorDiv.style.display = 'none';
                    loadingDiv.style.display = 'block';
                    
                    try {{
                        // Call authenticate endpoint
                        const formData = new FormData();
                        formData.append('username', username);
                        formData.append('password', password);
                        
                        const response = await fetch('/authenticate', {{
                            method: 'POST',
                            body: formData
                        }});
                        
                        if (response.ok) {{
                            // Success - will be redirected
                            const result = await response.json();
                            if (result.redirect) {{
                                window.location.href = result.redirect;
                            }}
                        }} else {{
                            const error = await response.json();
                            errorDiv.textContent = error.detail || 'Login failed';
                            errorDiv.style.display = 'block';
                            submitBtn.disabled = false;
                            loadingDiv.style.display = 'none';
                        }}
                    }} catch (error) {{
                        errorDiv.textContent = 'Network error: ' + error.message;
                        errorDiv.style.display = 'block';
                        submitBtn.disabled = false;
                        loadingDiv.style.display = 'none';
                    }}
                }});
            </script>
        </body>
        </html>
        """)

    async def handle_authenticate(self, request):
        """Process login form submission and call Gen8i API"""
        try:
            form_data = await request.form()
            username = form_data.get("username")
            password = form_data.get("password")

            if username and not self._rate_limiter.is_allowed(username):
                return HTMLResponse(
                    json.dumps({"detail": "Too many login attempts. Please try again later."}),
                    status_code=429,
                    media_type="application/json",
                )

            if not username or not password:
                return HTMLResponse(
                    json.dumps({"detail": "Username and password are required"}),
                    status_code=400,
                    media_type="application/json",
                )

            # Call Gen8i API
            async with httpx.AsyncClient() as client:
                # Prepare JSON payload for new API format
                payload = {
                    "email": username,
                    "password": password,
                    "client_id": self.config.CLIENT_ID,
                }

                print(f"[AUTH] Calling Gen8i API...", file=sys.stderr)
                print(
                    f"[INFO] URL: {self.config.API_BASE_URL}{self.config.LOGIN_ENDPOINT}",
                    file=sys.stderr,
                )
                print(f"[INFO] Email: {_mask(username)}", file=sys.stderr)
                print(f"[INFO] Client ID: {self.config.CLIENT_ID}", file=sys.stderr)

                response = await client.post(
                    f"{self.config.API_BASE_URL}{self.config.LOGIN_ENDPOINT}",
                    json=payload,
                    timeout=30.0,
                )

                if response.status_code == 200:
                    result = response.json()
                    access_token = result.get("access_token")
                    refresh_token = result.get("refresh_token")
                    token_type = result.get("token_type", "bearer")
                    expires_in = result.get("expires_in", 3600)
                    refresh_expires_in = result.get("refresh_expires_in", 2592000)

                    if access_token:
                        # Store complete token response
                        self.received_token = result
                        self.login_complete.set()

                        # Redirect to callback success page without token in URL
                        return HTMLResponse(
                            json.dumps(
                                {
                                    "redirect": f"http://localhost:{self.config.CALLBACK_PORT}/callback?status=success"
                                }
                            ),
                            status_code=200,
                            media_type="application/json",
                        )
                    else:
                        return HTMLResponse(
                            json.dumps({"detail": "No access_token in response"}),
                            status_code=400,
                            media_type="application/json",
                        )
                else:
                    # Log detailed error server-side only
                    print(f"[ERROR] Auth failed with status {response.status_code}", file=sys.stderr)
                    # Return generic message to client
                    return HTMLResponse(
                        json.dumps({"detail": "Authentication failed. Please check your credentials and try again."}),
                        status_code=401,
                        media_type="application/json",
                    )

        except Exception as e:
            print(f"[ERROR] Authentication error: {e}", file=sys.stderr)
            return HTMLResponse(
                json.dumps({"detail": "An internal error occurred. Please try again."}),
                status_code=500,
                media_type="application/json",
            )

    async def handle_callback(self, request):
        """Handle OAuth callback - exchange authorization code for tokens"""
        # Get code and state from query parameters
        code = request.query_params.get("code")
        received_state = request.query_params.get("state")

        # Validate state parameter (CSRF protection)
        if not received_state or not PKCEManager.validate_state(
            received_state, self.state
        ):
            print(
                "[ERROR] State validation failed - possible CSRF attack!",
                file=sys.stderr,
            )
            return _secure_html(
                """
                <!DOCTYPE html>
                <html>
                <head><title>Login Failed</title></head>
                <body style="font-family: Arial; text-align: center; padding: 50px; background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);">
                    <div style="background: white; padding: 3rem; border-radius: 1rem; max-width: 400px; margin: 0 auto;">
                        <h1 style="color: #e53e3e;">Security Error</h1>
                        <p>State validation failed. Possible CSRF attack detected.</p>
                    </div>
                </body>
                </html>
                """,
                status_code=400,
            )

        if not code:
            print("[ERROR] No authorization code received", file=sys.stderr)
            return _secure_html(
                """
                <!DOCTYPE html>
                <html>
                <head><title>Login Failed</title></head>
                <body style="font-family: Arial; text-align: center; padding: 50px; background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);">
                    <div style="background: white; padding: 3rem; border-radius: 1rem; max-width: 400px; margin: 0 auto;">
                        <h1 style="color: #e53e3e;">Login Failed</h1>
                        <p>No authorization code received.</p>
                    </div>
                </body>
                </html>
                """,
                status_code=400,
            )

        try:
            # Exchange authorization code for tokens
            print(
                f"[AUTH] Exchanging authorization code for tokens...", file=sys.stderr
            )
            print(f"[DEBUG] Code: {_mask(code)}",  file=sys.stderr)
            print(
                f"[DEBUG] Code verifier: {_mask(self.code_verifier)}", file=sys.stderr
            )

            token_url = f"{self.config.API_BASE_URL}{self.config.TOKEN_ENDPOINT}"
            redirect_uri = f"http://localhost:{self.config.CALLBACK_PORT}/callback"

            async with httpx.AsyncClient() as client:
                response = await client.post(
                    token_url,
                    json={
                        "code": code,
                        "client_id": self.config.CLIENT_ID,
                        "code_verifier": self.code_verifier,
                        "redirect_uri": redirect_uri,
                        "grant_type": "authorization_code",
                    },
                    headers={"Content-Type": "application/json"},
                    timeout=30.0,
                )

            if response.status_code != 200:
                print(f"[ERROR] Token exchange failed: {response.status_code}", file=sys.stderr)
                return _secure_html(
                    """
                    <!DOCTYPE html>
                    <html>
                    <head><title>Login Failed</title></head>
                    <body style="font-family: Arial; text-align: center; padding: 50px; background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);">
                        <div style="background: white; padding: 3rem; border-radius: 1rem; max-width: 400px; margin: 0 auto;">
                            <h1 style="color: #e53e3e;">Token Exchange Failed</h1>
                            <p>Could not exchange authorization code for tokens. Please try again.</p>
                        </div>
                    </body>
                    </html>
                    """,
                    status_code=500,
                )

            # Parse token response
            token_data = response.json()
            self.received_token = token_data

            print(f"[SUCCESS] Tokens received successfully!", file=sys.stderr)
            print(
                f"[DEBUG] Access token: {_mask(token_data.get('access_token', ''))}",
                file=sys.stderr,
            )
            print(
                f"[DEBUG] Refresh token: {_mask(token_data.get('refresh_token', ''))}",
                file=sys.stderr,
            )

            # Signal completion
            self.login_complete.set()

            # Return success page
            return _secure_html("""
                <!DOCTYPE html>
                <html>
                <head>
                    <title>Login Successful</title>
                    <style>
                        body {
                            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
                            display: flex;
                            justify-content: center;
                            align-items: center;
                            height: 100vh;
                            margin: 0;
                            background: linear-gradient(135deg, #F27842 0%, #ff9a6c 100%);
                        }
                        .container {
                            background: white;
                            padding: 3rem;
                            border-radius: 1rem;
                            box-shadow: 0 20px 60px rgba(242, 120, 66, 0.3);
                            text-align: center;
                            max-width: 400px;
                        }
                        h1 { color: #F27842; margin-bottom: 1rem; }
                        p { color: #4a5568; font-size: 1.1rem; }
                        .checkmark {
                            width: 80px;
                            height: 80px;
                            margin: 0 auto 1.5rem;
                            border-radius: 50%;
                            display: block;
                            stroke-width: 2;
                            stroke: #F27842;
                            stroke-miterlimit: 10;
                            box-shadow: inset 0px 0px 0px #F27842;
                            animation: fill .4s ease-in-out .4s forwards, scale .3s ease-in-out .9s both;
                        }
                        .checkmark__circle {
                            stroke-dasharray: 166;
                            stroke-dashoffset: 166;
                            stroke-width: 2;
                            stroke-miterlimit: 10;
                            stroke: #F27842;
                            fill: none;
                            animation: stroke 0.6s cubic-bezier(0.65, 0, 0.45, 1) forwards;
                        }
                        .checkmark__check {
                            transform-origin: 50% 50%;
                            stroke-dasharray: 48;
                            stroke-dashoffset: 48;
                            animation: stroke 0.3s cubic-bezier(0.65, 0, 0.45, 1) 0.8s forwards;
                        }
                        @keyframes stroke {
                            100% { stroke-dashoffset: 0; }
                        }
                        @keyframes scale {
                            0%, 100% { transform: none; }
                            50% { transform: scale3d(1.1, 1.1, 1); }
                        }
                        @keyframes fill {
                            100% { box-shadow: inset 0px 0px 0px 30px #F27842; }
                        }
                        #status { font-size: 0.9rem; color: #718096; margin-top: 2rem; }
                    </style>
                </head>
                <body>
                    <div class="container">
                        <svg class="checkmark" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 52 52">
                            <circle class="checkmark__circle" cx="26" cy="26" r="25" fill="none"/>
                            <path class="checkmark__check" fill="none" d="M14.1 27.2l7.1 7.2 16.7-16.8"/>
                        </svg>
                        <h1>✅ Login Successful!</h1>
                        <p>Authentication completed successfully.</p>
                        <p id="status">Attempting to close this window in <span id="countdown">3</span> seconds...</p>
                    </div>
                    <script>
                        let count = 3;
                        const countdownEl = document.getElementById('countdown');
                        const statusEl = document.getElementById('status');

                        const interval = setInterval(() => {
                            count--;
                            if (count > 0) {
                                countdownEl.textContent = count;
                            } else {
                                clearInterval(interval);
                                // Try to close the window
                                window.close();
                                // If window.close() fails (which it often does), show a message
                                setTimeout(() => {
                                    statusEl.textContent = 'You can now close this window and return to Claude.';
                                }, 100);
                            }
                        }, 1000);
                    </script>
                </body>
                </html>
            """)

        except Exception as e:
            print(f"[ERROR] Token exchange error: {e}", file=sys.stderr)
            return _secure_html(
                """
                <!DOCTYPE html>
                <html>
                <head><title>Error</title></head>
                <body style="font-family: Arial; text-align: center; padding: 50px; background: linear-gradient(135deg, #F27842 0%, #ff9a6c 100%);">
                    <div style="background: white; padding: 3rem; border-radius: 1rem; max-width: 400px; margin: 0 auto;">
                        <h1 style="color: #e53e3e;">Error</h1>
                        <p>An error occurred during authentication. Please try again.</p>
                    </div>
                </body>
                </html>
                """,
                status_code=500,
            )

    async def handle_login_callback(self, request):
        """Handle callback after direct login - just show success page"""
        return _secure_html("""
            <!DOCTYPE html>
            <html>
            <head><title>Login Successful</title></head>
            <body style="font-family: Arial; text-align: center; padding: 50px; background: linear-gradient(135deg, #F27842 0%, #ff9a6c 100%);">
                <div style="background: white; padding: 3rem; border-radius: 1rem; max-width: 400px; margin: 0 auto;">
                    <h1 style="color: #F27842;">Login Successful!</h1>
                    <p>Authentication completed. You can close this window.</p>
                </div>
            </body>
            </html>
        """)

    def create_app(self):
        """Create Starlette app for OAuth callback handling"""
        return Starlette(
            routes=[
                Route("/callback", self.handle_callback, methods=["GET"]),
            ]
        )

    async def start_login_flow(self, timeout: int = 120) -> str:
        """Start PKCE OAuth login flow - open browser to Gen8i's login page"""
        from urllib.parse import urlencode

        # Generate PKCE parameters
        print("[AUTH] Generating PKCE parameters...", file=sys.stderr)
        self.code_verifier = PKCEManager.generate_code_verifier()
        code_challenge = PKCEManager.generate_code_challenge(self.code_verifier)
        self.state = PKCEManager.generate_state()

        print(f"[DEBUG] Code verifier: {_mask(self.code_verifier)}", file=sys.stderr)
        print(f"[DEBUG] Code challenge: {_mask(code_challenge)}", file=sys.stderr)
        print(f"[DEBUG] State: {_mask(self.state)}", file=sys.stderr)

        # Create callback server
        app = self.create_app()

        # Configure server
        config = uvicorn.Config(
            app, host="localhost", port=self.config.CALLBACK_PORT, log_level="warning"
        )
        server = uvicorn.Server(config)

        # Run server in background
        server_task = asyncio.create_task(server.serve())

        try:
            # Give server time to start
            await asyncio.sleep(1)

            # Build OAuth authorization URL
            redirect_uri = f"http://localhost:{self.config.CALLBACK_PORT}/callback"

            params = {
                "client_id": self.config.CLIENT_ID,
                "redirect_uri": redirect_uri,
                "state": self.state,
                "code_challenge": code_challenge,
                "code_challenge_method": "S256",
                "response_type": "code",
            }

            authorize_url = f"{self.config.API_BASE_URL}{self.config.AUTHORIZE_ENDPOINT}?{urlencode(params)}"

            print(f"\n[AUTH] Opening browser for OAuth login...", file=sys.stderr)
            print(f"[INFO] Authorization URL: {authorize_url}\n", file=sys.stderr)

            # Open browser to Gen8i's login page
            webbrowser.open(authorize_url)

            # Wait for login completion or timeout
            try:
                await asyncio.wait_for(self.login_complete.wait(), timeout=timeout)
            except asyncio.TimeoutError:
                raise TimeoutError(
                    f"Login timeout - no response within {timeout} seconds"
                )

            if not self.received_token:
                raise ValueError("Login completed but no token received")

            print("✅ Login successful!\n", file=sys.stderr)
            return self.received_token  # Return full token response dict

        finally:
            # Cleanup server
            server.should_exit = True
            try:
                await asyncio.wait_for(server_task, timeout=5)
            except asyncio.TimeoutError:
                print("[WARN] Server shutdown timeout", file=sys.stderr)


# ============================================================================
# API Client
# ============================================================================


class APIClient:
    """Client for making authenticated API requests"""

    def __init__(self, config: ServerConfig, token_manager: TokenManager):
        self.config = config
        self.token_manager = token_manager

    async def get_use_cases(self) -> dict:
        """Fetch available use cases for authenticated user"""
        token = await self.token_manager.get_valid_token()
        if not token:
            raise ValueError("No authentication token available. Please login first.")

        # Gen8i API uses GET request with Authorization header
        headers = {
            "Authorization": f"Bearer {token}",
        }

        async with httpx.AsyncClient() as client:
            response = await client.get(
                f"{self.config.API_BASE_URL}{self.config.USE_CASES_ENDPOINT}",
                headers=headers,
                timeout=30.0,
            )
            response.raise_for_status()
            return response.json()

    async def query_use_case(
        self, usecase_id: str, usecase_name: str, query_text: str
    ) -> dict:
        """Execute query in specific use case"""
        token = await self.token_manager.get_valid_token()
        if not token:
            raise ValueError("No authentication token available. Please login first.")

        # Gen8i API format
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
            "accept": "application/json",
        }

        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.config.API_BASE_URL}{self.config.QUERY_ENDPOINT}",
                headers=headers,
                json={
                    "usecase_id": usecase_id,
                    "usecase_name": usecase_name,
                    "query_text": query_text,
                },
                timeout=60.0,
            )
            response.raise_for_status()
            return response.json()


# ============================================================================
# MCP Server
# ============================================================================

# Initialize configuration and global instances
print("[START] Initializing Gen8i MCP Server...\n", file=sys.stderr)

# Load configuration
config = ServerConfig()
print(f"[CONFIG] Configuration:", file=sys.stderr)
print(f"   Client ID: {config.CLIENT_ID}", file=sys.stderr)
print(f"   API Base URL: {config.API_BASE_URL}", file=sys.stderr)
print(f"   Callback Port: {config.CALLBACK_PORT}\n", file=sys.stderr)

# Initialize components
token_manager = TokenManager(config)
api_client = APIClient(config, token_manager)
login_server = LoginFlowServer(config, token_manager)


# Try to load existing token on startup
existing_token = token_manager.load_token()
if existing_token:
    print("[OK] Found existing token\n", file=sys.stderr)
else:
    print("[INFO] No existing token found\n", file=sys.stderr)

# Create FastMCP server
mcp = FastMCP("Gen8i MCP Server")


# ============================================================================
# MCP Tools
# ============================================================================


@mcp.tool()
async def login() -> str:
    """
    Authenticate with Gen8i server to access resources.

    Opens a browser window for secure authentication on the Gen8i server.
    Uses OAuth 2.0 Authorization Code flow with PKCE for security.

    The system uses CLIENT_ID from claude_config.json.
    Tokens are stored securely in the OS keyring.

    Returns:
        dict: Success message with authentication status
    """
    global login_server, token_manager

    print("\n" + "=" * 60, file=sys.stderr)
    print("[AUTH] STARTING GEN8I BROWSER LOGIN", file=sys.stderr)
    print("=" * 60 + "\n", file=sys.stderr)

    try:
        # Start browser-based login flow - returns full token response dict
        token_response = await login_server.start_login_flow(timeout=120)

        # Extract token data
        access_token = token_response.get("access_token")
        refresh_token = token_response.get("refresh_token")
        token_type = token_response.get("token_type", "bearer")
        expires_in = token_response.get("expires_in", 3600)
        refresh_expires_in = token_response.get("refresh_expires_in", 2592000)

        # Save complete token data
        token_manager.save_token(
            access_token=access_token,
            refresh_token=refresh_token,
            expires_in=expires_in,
            refresh_expires_in=refresh_expires_in,
            token_type=token_type,
        )

        return f"""[SUCCESS] Login successful! 

Client ID: {config.CLIENT_ID}

Your authentication tokens have been saved:
- Access token valid for {expires_in // 60} minutes
- Refresh token valid for {refresh_expires_in // 86400} days

Tokens will be automatically refreshed when needed.
You can now use the other tools:
- get_use_case: View your available use cases
- query: Execute queries in your use cases
"""

    except TimeoutError:
        return "[TIMEOUT] Login timeout: No response received within 2 minutes. Please try again."

    except Exception as e:
        return f"[ERROR] Login failed: {str(e)}"


@mcp.tool()
async def get_use_case() -> str:
    """
    Get available use cases (data sources) for the authenticated user.

    **IMPORTANT**: Call this tool FIRST before answering any business/analytical questions
    to see which use cases are available and get their IDs.

    Use cases represent different data sources or analytical contexts (e.g., sales data,
    financial reports, product analytics). You'll need the use case ID to query data.

    Requires authentication. If not logged in, returns an error with instructions to login.

    Returns:
        JSON string containing available use cases with their IDs and descriptions

    Example workflow:
        1. User asks: "What were GSK's sales in 2024?"
        2. Call get_use_case() to see available use cases
        3. Call query() with the appropriate usecase_id and the question
    """
    global api_client, token_manager

    # Check if token exists or can be refreshed
    token = await token_manager.get_valid_token()
    if not token:
        return """[ERROR] Authentication required!

You need to login first. Please use the 'login' tool to authenticate:

Example: Use the login tool to open browser and authenticate.
"""

    try:
        # Fetch use cases
        use_cases = await api_client.get_use_cases()

        return f"""[SUCCESS] Available Use Cases:

{json.dumps(use_cases, indent=2)}

You can now query these use cases using the 'query' tool.
"""

    except httpx.HTTPStatusError as e:
        if e.response.status_code == 401:
            token_manager.clear_token()
            return """[ERROR] Authentication failed!

Your token is invalid. Please login again using the 'login' tool.
"""
        elif e.response.status_code == 403:
            return """[ERROR] Access denied!

You don't have permission to view use cases. Please contact your administrator.
"""
        else:
            return f"[ERROR] API Error: {e.response.status_code} - {e.response.text}"

    except Exception as e:
        return f"[ERROR] Error fetching use cases: {str(e)}"


@mcp.tool()
async def query(usecase_id: str, usecase_name: str, query_text: str) -> str:
    """
    Execute analytical/business queries using Gen8i's AI-powered data analysis.

    **USE THIS TOOL** when the user asks questions about:
    - Sales, revenue, or financial data (e.g., "How much sales did GSK do in 2024?")
    - Business metrics and KPIs (e.g., "What was the revenue growth?")
    - Product performance (e.g., "Top performing products last quarter")
    - Company analytics (e.g., "Compare Roche and Pfizer sales")
    - Trends and patterns (e.g., "Sales trends for FY 2021")
    - Any data-driven or analytical questions

    **WORKFLOW**:
    1. If you haven't called get_use_case() yet, call it first to get available use cases
    2. Use the appropriate usecase_id and usecase_name for the data domain
    3. Pass the user's question as the 'query_text' parameter

    Args:
        usecase_id: The ID of the use case (string, get from get_use_case tool)
        usecase_name: The name of the use case (e.g., "EC Transcript")
        query_text: The business/analytical question to answer (e.g., "Types of explosive")

    Requires authentication. If not logged in, returns an error with instructions to login.

    Returns:
        AI-generated analysis and data results as JSON string

    Examples:
        - query(usecase_id="1", usecase_name="EC Transcript", query_text="Types of explosive")
        - query(usecase_id="298", usecase_name="Sales Data", query_text="how much sales done by GSK in 2024")
    """
    global api_client, token_manager

    # Check if token exists or can be refreshed
    token = await token_manager.get_valid_token()
    if not token:
        return """[ERROR] Authentication required!

You need to login first. Please use the 'login' tool to authenticate.
"""

    try:
        # Execute query
        result = await api_client.query_use_case(usecase_id, usecase_name, query_text)

        return f"""[SUCCESS] Query Results:

Use Case ID: {usecase_id}
Use Case Name: {usecase_name}
Query: {query_text}

{json.dumps(result, indent=2)}
"""

    except httpx.HTTPStatusError as e:
        if e.response.status_code == 401:
            token_manager.clear_token()
            return """[ERROR] Authentication failed!

Your token is invalid. Please login again using the 'login' tool.
"""
        elif e.response.status_code == 403:
            return f"""[ERROR] Access denied!

You don't have permission to query use case '{usecase_id}'. 
Please check your available use cases using the 'get_use_case' tool.
"""
        elif e.response.status_code == 404:
            return f"""[ERROR] Use case not found!

The use case '{usecase_id}' doesn't exist. 
Please check your available use cases using the 'get_use_case' tool.
"""
        else:
            return f"[ERROR] API Error: {e.response.status_code} - {e.response.text}"

    except Exception as e:
        return f"[ERROR] Error executing query: {str(e)}"


# ============================================================================
# Main
# ============================================================================

def main():
    """Entry point for the MCP server - called by uvx"""
    try:
        mcp.run(transport="stdio")
    except KeyboardInterrupt:
        print("\n[EXIT] Server stopped by user", file=sys.stderr)
    except Exception as e:
        print(f"\n[ERROR] Server error: {e}", file=sys.stderr)
        raise


if __name__ == "__main__":
    main()
