"""
MCP Server for browser automation with Console Tool Support.

This server provides browser automation capabilities through MCP protocol.
It supports console viewing capabilities for debugging web applications.
Uses stdio transport for communication with MCP clients.
"""
import asyncio
import logging
import os
import sys
import io
from pathlib import Path

# --- Set CWD and Path to Package Root ---
# This script now runs as part of an installed package
try:
    # Get the package directory
    package_root = Path(__file__).parent.resolve()
    # For auth_state.json, we'll use the current working directory
    working_dir = Path.cwd()
except Exception as e:
    print(f"Error setting up script environment: {e}", file=sys.stderr)
    sys.exit(1)

# --- Fix Windows encoding and disable telemetry ---
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
os.environ['PYTHONIOENCODING'] = 'utf-8'
os.environ['ANONYMIZED_TELEMETRY'] = 'false'
os.environ['BROWSER_USE_CLOUD_SYNC'] = 'false'

# --- Configure Logging ---
logging.basicConfig(level=logging.INFO, stream=sys.stderr, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- Main Application Imports ---
try:
    from mcp.server.fastmcp import FastMCP
    from browser_use import Agent, BrowserProfile
    from browser_use.browser.session import BrowserSession
    from browser_use.controller.service import Controller
    
    # Import console tool registration functions
    from .console_tool import register_init_console_capture, register_view_console_logs
    
    # Check if playwright chromium is installed
    import subprocess
    try:
        subprocess.run(['playwright', 'show-trace'], capture_output=True, check=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        logger.warning("Playwright chromium might not be installed. If you encounter browser errors, run:")
        logger.warning("playwright install chromium --with-deps --no-shell")
        
except ImportError as e:
    logger.error(f"A required package is not installed. Please run 'pip install browser-use-mcp-console'. Details: {e}")
    sys.exit(1)

# --- MCP Server Definition ---
mcp = FastMCP(
    name="BrowserUseAgent",
    instructions="A server that runs the browser-use agent with console viewing capabilities for debugging and monitoring web applications."
)

@mcp.tool(name="run_browser_tasks", description="Run browser automation tasks with optional console debugging")
async def run_browser_tasks(
    tasks: list[str],
    model: str | None = None,
    headless: bool = False,
    max_steps: int = 100,
) -> str:
    """
    Runs browser tasks with optional console viewing capabilities.
    Can handle single or multiple tasks efficiently.
    For single task, runs normally. For multiple tasks, runs in parallel.
    
    Args:
        tasks: List of tasks to execute (can be a single task)
        model: LLM model to use
        headless: Whether to run browsers in headless mode
        max_steps: Maximum steps per task
    
    Returns:
        Summary of task results
    """
    
    if not tasks:
        return "No tasks provided"
    
    # Configure LLM
    # Support multiple LLM providers
    if os.getenv("OPENROUTER_API_KEY"):
        from browser_use.llm.openrouter.chat import ChatOpenRouter
        if model is None:
            model = "google/gemini-2.5-pro"  # OpenRouter default model
        llm = ChatOpenRouter(model=model, api_key=os.getenv("OPENROUTER_API_KEY"))
    elif os.getenv("OPENAI_API_KEY"):
        from browser_use.llm.openai.chat import ChatOpenAI
        if model is None:
            model = "gpt-4o-mini"  # OpenAI default model
        llm = ChatOpenAI(model=model, api_key=os.getenv("OPENAI_API_KEY"))
    elif os.getenv("ANTHROPIC_API_KEY"):
        from browser_use.llm.anthropic.chat import ChatAnthropic
        if model is None:
            model = "claude-3-5-sonnet-20241022"  # Anthropic default model
        llm = ChatAnthropic(model=model, api_key=os.getenv("ANTHROPIC_API_KEY"))
    else:
        return "Error: No API key found. Please set one of: OPENROUTER_API_KEY, OPENAI_API_KEY, or ANTHROPIC_API_KEY"
    
    storage_state_path = working_dir / 'auth_state.json'
    storage_state = str(storage_state_path) if storage_state_path.exists() else None
    
    if storage_state:
        logger.info(f"Using saved authentication state from: {storage_state}")
    
    profile_template = BrowserProfile(
        headless=headless,
        storage_state=storage_state,  # Use saved login state if exists
        user_data_dir=None,  # Independent temp directory for each session
        keep_alive=False,
        enable_default_extensions=False,
        browser_kwargs={'timeout': 60000}
    )
    
    agents = []
    for task in tasks:
        browser_session = BrowserSession(browser_profile=profile_template)
        
        controller = Controller()
        register_init_console_capture(controller)
        register_view_console_logs(controller)
        
        agent = Agent(
            task=task, 
            llm=llm, 
            browser_session=browser_session,
            controller=controller,
            use_vision=True  # Enable vision capabilities for better page understanding
        )
        
        agents.append(agent)
    
    # Single task or multi-task execution mode
    execution_mode = "single" if len(agents) == 1 else "parallel"
    logger.info(f"Execution mode: {execution_mode} ({len(agents)} agent(s))")
    
    try:
        logger.info("Starting browser session(s)...")
        await asyncio.gather(*[agent.browser_session.start() for agent in agents])
        logger.info(f"{len(agents)} browser session(s) started")

        if execution_mode == "single":
            logger.info("Running single task...")
        else:
            logger.info(f"Running {len(tasks)} tasks in parallel...")
            
        results = await asyncio.gather(
            *[agent.run(max_steps=max_steps) for agent in agents],
            return_exceptions=True  # Continue other tasks even if one fails
        )
        logger.info("Task execution completed")

        task_results = []
        for i, (task, result) in enumerate(zip(tasks, results)):
            if isinstance(result, Exception):
                task_results.append(f"Task {i+1} '{task[:50]}...': Failed - {str(result)}")
            else:
                success = result.is_successful() if hasattr(result, 'is_successful') else False
                final = result.final_result() if hasattr(result, 'final_result') else "No result"
                task_results.append(f"Task {i+1} '{task[:50]}...': {'Success' if success else 'Completed'} - {final}")

        if len(tasks) == 1:
            return task_results[0] if task_results else "No result"
        else:
            return "\n".join([
                f"Executed {len(tasks)} tasks in parallel:",
                "=" * 50,
                *task_results
            ])
        
    except Exception as e:
        logger.error(f"Parallel execution failed: {e}", exc_info=True)
        return f"Parallel execution error: {e}"
    
    finally:
        logger.info("Closing all browser sessions...")
        await asyncio.gather(
            *[agent.browser_session.kill() for agent in agents],
            return_exceptions=True
        )
        logger.info("All sessions closed")

def main():
    """Runs the server using stdio transport for MCP clients."""
    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
    
    logger.info("Starting Browser-Use MCP Server with Console Support...")
    logger.info("Environment optimizations applied:")
    logger.info("  - UTF-8 encoding fixed")
    logger.info("  - Telemetry disabled")
    logger.info("  - Cloud sync disabled")
    logger.info("  - Browser profile optimized")
    logger.info("  - Console viewing tools integrated")
    
    # Run server (stdio mode for uvx/Claude)
    mcp.run()

if __name__ == "__main__":
    main()