"""
Deep Research Multi-Agent System with Browser Agent Pools

This example demonstrates using agent pools to enable true parallel execution
of multiple browser agents without state conflicts. Each browser agent instance
has its own isolated Playwright browser/page, memory, and temporary directories.
"""

import asyncio
import json
import logging
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict

# Add project root to Python path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

# Imports
from dotenv import load_dotenv

from marsys.agents import Agent, create_browser_agent_pool
from marsys.agents.browser_agent import BrowserAgent
from marsys.agents.registry import AgentRegistry
from marsys.coordination import Orchestra
from marsys.coordination.communication import CommunicationManager
from marsys.coordination.communication.channels import TerminalChannel
from marsys.environment.tools import tool_google_search_api
from marsys.models import ModelConfig

load_dotenv()


class LoggingAgent(Agent):
    """Agent with built-in logging of all interactions."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.call_count = 0

    async def run_step(self, request, context):
        self.call_count += 1
        start_time = datetime.now()

        log_entry = {
            "agent_call_number": self.call_count,
            "agent_name": self.name,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "caller": context.get("caller_agent", "System"),
            "request": request,
            "response": None,
            "duration_seconds": None,
        }

        result = await super().run_step(request, context)
        duration = (datetime.now() - start_time).total_seconds()
        log_entry["duration_seconds"] = round(duration, 2)
        log_entry["response"] = result

        AGENT_RESPONSE_FILE = (
            Path("examples/real_world/output") / "deep-research-pooled-responses.txt"
        )
        with open(AGENT_RESPONSE_FILE, "a") as f:
            f.write("\n" + "=" * 80 + "\n")
            f.write(
                f"AGENT: {self.name} | CALL #{self.call_count} | CALLED BY: {log_entry['caller']}\n"
            )
            f.write("=" * 80 + "\n\n")
            f.write(f"Timestamp: {log_entry['timestamp']}\n")
            f.write(f"Duration: {duration:.2f} seconds\n\n")
            try:
                f.write(json.dumps(log_entry["request"], indent=2, default=str) + "\n")
            except:
                f.write(str(log_entry["request"]) + "\n")
            f.write("\nRESPONSE:\n")
            try:
                f.write(json.dumps(log_entry["response"], indent=2, default=str) + "\n")
            except:
                f.write(str(log_entry["response"]) + "\n")

        return result


logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def create_model_configs() -> Dict[str, ModelConfig]:
    return {
        "gemini_pro": ModelConfig(
            type="api",
            name="google/gemini-2.5-flash",
            provider="openrouter",
            api_key=os.getenv("OPENROUTER_API_KEY"),
            temperature=0.7,
            max_tokens=10000,
        )
    }


async def ask_user_question(question: str, context: str = "") -> Dict[str, Any]:
    # Auto-response for testing without user interaction
    print("\n" + "=" * 80)
    print("🤖 CLARIFYING QUESTION FROM SCOPER AGENT")
    print("=" * 80)
    if context:
        print(f"📋 Context: {context}")
        print("-" * 80)
    print(f"❓ Question: {question}")
    print("-" * 80)
    
    # Auto-response for testing
    user_response = "Find information about the latest advancements in quantum computing and artificial intelligence integration"
    print(f"💬 Auto-response (test mode): {user_response}")
    print("=" * 80)
    print("✅ Processing response...\n")
    
    return {
        "response": user_response,
        "metadata": {"question": question, "context": context},
    }


async def create_agents(configs: Dict[str, ModelConfig]) -> Dict[str, Agent]:
    """Create agents including a pool of browser agents."""
    agents = {}

    agents["scoper_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="scoper_agent",
        tools={"ask_user_question": ask_user_question},
        description="""You are a Research Scoper Agent. Your role is to understand the user's 
        research requirements and break them down into specific, actionable tasks.
        
        You should:
        1. Clarify any ambiguous requirements
        2. Identify key research areas
        3. Delegate specific research tasks to the retriever_agent
        
        You can ask the user clarifying questions using the ask_user_question tool.""",
        max_tokens=2000,
    )

    agents["retriever_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="retriever_agent",
        tools={"google_search_api": tool_google_search_api},
        description="""You are a Research Retriever Agent specialized in finding and gathering
        information from various sources.
        
        Your capabilities:
        1. Use google_search_api to find relevant web sources
        2. Coordinate multiple browser_agent instances to extract content in parallel
        3. Aggregate findings from multiple sources
        
        IMPORTANT: You have access to a pool of browser_agent instances. You can invoke
        multiple browser agents in parallel by providing an array of invocation requests.
        This allows you to extract content from multiple web pages simultaneously.
        
        Example parallel invocation:
        {
            "next_action": "invoke_agent",
            "action_input": [
                {"agent_name": "browser_agent", "request": {"url": "https://example1.com", "task": "extract main content"}},
                {"agent_name": "browser_agent", "request": {"url": "https://example2.com", "task": "extract key points"}},
                {"agent_name": "browser_agent", "request": {"url": "https://example3.com", "task": "find statistics"}}
            ]
        }
        
        The browser pool will automatically allocate available instances to handle your requests.""",
        max_tokens=2000,
    )

    # Create a POOL of browser agents instead of a single instance
    # This enables true parallel execution without state conflicts
    browser_pool = await create_browser_agent_pool(
        num_instances=3,  # Create 3 independent browser instances
        model_config=configs["gemini_pro"],
        name="browser_agent",
        description="""You are a Browser Agent capable of web navigation and content extraction.
        
        You are part of a pool of browser agents, each with its own isolated browser instance.
        This means you can work in parallel with other browser instances without conflicts.
        
        Your tasks:
        1. Navigate to specified URLs
        2. Extract relevant content based on the request
        3. Handle any popups or cookies dialogs
        4. Return structured information
        
        Each instance in the pool has its own:
        - Playwright browser and page instance
        - Memory and conversation history
        - Temporary directory for downloads/screenshots
        
        This isolation ensures reliable parallel operation.""",
        headless=True,
        viewport_width=1440,
        viewport_height=960,
        auto_screenshot=True,
        timeout=10000,
    )
    
    # Register the pool with the AgentRegistry
    AgentRegistry.register_pool(browser_pool)
    
    # Note: We register the pool under the name "browser_agent"
    # The coordination system will automatically handle instance allocation
    agents["browser_agent"] = browser_pool

    agents["analyzer_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="analyzer_agent",
        description="""You are a Research Analyzer Agent. Your role is to analyze and synthesize
        information gathered by the retriever and browser agents.
        
        You should:
        1. Identify key themes and patterns
        2. Evaluate the credibility of sources
        3. Extract the most important insights
        4. Prepare structured analysis for the reporter""",
        max_tokens=2000,
    )

    agents["reporter_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="reporter_agent",
        description="""You are a Research Reporter Agent. Your role is to create comprehensive
        research reports based on the analysis provided.
        
        You should:
        1. Organize findings into a clear structure
        2. Highlight key discoveries
        3. Provide executive summary
        4. Include relevant citations and sources
        5. Return the final report to the user""",
        max_tokens=4000,
    )

    return agents


async def main():
    """Main execution function demonstrating agent pools."""
    
    # Setup
    configs = create_model_configs()
    
    # Create output directory
    output_dir = Path("examples/real_world/output")
    output_dir.mkdir(exist_ok=True)
    
    # Clear previous log
    log_file = output_dir / "deep-research-pooled-responses.txt"
    if log_file.exists():
        log_file.unlink()
    
    print("\n" + "=" * 80)
    print("🚀 DEEP RESEARCH SYSTEM WITH BROWSER AGENT POOLS")
    print("=" * 80)
    print("\n📋 Creating agents with browser pool...")
    
    # Create agents including browser pool
    agents = await create_agents(configs)
    
    # Display pool information
    browser_pool = agents["browser_agent"]
    print(f"\n✅ Created browser agent pool with {browser_pool.num_instances} instances")
    print(f"   Each instance has isolated browser, memory, and temp directories")
    print(f"   Instances can work in parallel without conflicts")
    
    # Define topology - note that browser_agent is now a pool
    topology = {
        "agents": [
            "User",
            "scoper_agent",
            "retriever_agent",
            "browser_agent",  # This is now a pool!
            "analyzer_agent",
            "reporter_agent",
        ],
        "flows": [
            "User -> scoper_agent",
            "scoper_agent -> retriever_agent",
            "retriever_agent -> browser_agent",
            "retriever_agent -> analyzer_agent",
            "browser_agent -> analyzer_agent",
            "analyzer_agent -> reporter_agent",
            "reporter_agent -> User",
        ],
        "rules": [
            "parallel(browser_agent)",  # Enable parallel execution of browser pool
            "max_turns(scoper_agent <-> User, 3)",
        ],
    }
    
    print("\n📊 Topology configured with parallel browser execution")
    print("   The retriever can invoke multiple browser instances simultaneously")
    
    # Setup communication
    terminal_channel = TerminalChannel()
    comm_manager = CommunicationManager({"terminal": terminal_channel})
    comm_manager.set_default_channel("terminal")
    
    # Initial task
    initial_task = """
    Research the latest developments in quantum computing and AI integration.
    Focus on:
    1. Recent breakthroughs (2024-2025)
    2. Key research institutions and companies
    3. Practical applications being developed
    4. Technical challenges and limitations
    
    Please provide a comprehensive report with sources.
    """
    
    print("\n📝 Starting research task...")
    print("-" * 80)
    
    # Run the orchestra
    try:
        result = await Orchestra.run(
            task=initial_task,
            topology=topology,
            communication_manager=comm_manager,
            context={"session_id": f"research_{datetime.now().strftime('%Y%m%d_%H%M%S')}"},
            max_steps=50,
        )
        
        # Display results
        print("\n" + "=" * 80)
        print("📊 RESEARCH COMPLETE")
        print("=" * 80)
        
        if result.success:
            print("\n✅ Research completed successfully!")
            print(f"📈 Total steps: {result.total_steps}")
            print(f"⏱️ Duration: {result.total_duration:.2f} seconds")
            
            # Show pool statistics
            pool_stats = browser_pool.get_statistics()
            print(f"\n📊 Browser Pool Statistics:")
            print(f"   Total allocations: {pool_stats['total_allocations']}")
            print(f"   Peak concurrent usage: {pool_stats['peak_concurrent_usage']}")
            print(f"   Average wait time: {pool_stats['average_wait_time']:.2f}s")
            
            print("\n📄 Final Report:")
            print("-" * 80)
            if result.final_response:
                print(result.final_response)
            else:
                print("No final response generated")
        else:
            print(f"\n❌ Research failed: {result.error}")
    
    except Exception as e:
        print(f"\n❌ Error during execution: {e}")
        logger.error(f"Execution error: {e}", exc_info=True)
    
    finally:
        # Clean up browser pool
        print("\n🧹 Cleaning up browser pool...")
        await browser_pool.cleanup()
        print("✅ Cleanup complete")


if __name__ == "__main__":
    print("\n🎯 Starting Deep Research System with Agent Pools")
    print("   This example demonstrates parallel browser execution")
    print("   without state conflicts using isolated instances\n")
    
    asyncio.run(main())