"""Simplified Multi-Agent Deep Research System

This example demonstrates a streamlined multi-agent research workflow with:
- Query analyzer creates research brief directly (no user intent or research statement agents)
- Web retrieval searches for research papers once
- Two parallel orchestrators process top 2 papers
- Each orchestrator manages browser agents and synthesizer
- Final synthesizer combines both reports

Workflow:
1. QueryAnalyzer → Research brief
2. WebRetrieval → Search and select top 2 papers
3. Parallel: Orchestrator1 & Orchestrator2 (each with same research statement + URLs)
4. Each Orchestrator: Browser agents (parallel) → Synthesizer
5. FinalSynthesizer → Final report
"""

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

import requests
from dotenv import load_dotenv

from marsys.agents import Agent
from marsys.agents.agent_pool import AgentPool
from marsys.agents.browser_agent import BrowserAgent
from marsys.agents.registry import AgentRegistry
from marsys.agents.utils import init_agent_logging
from marsys.coordination import Orchestra
from marsys.environment.tools import tool_google_search_api
from marsys.models.models import ModelConfig

# init_agent_logging(level=logging.INFO, clear_existing_handlers=True)


# ============================================================================
# TOOL DEFINITIONS
# ============================================================================


def tool_semantic_scholar_search(query: str, num_results: int = 10, year_filter: str = None) -> str:
    """
    Search academic papers using Semantic Scholar API.

    Args:
        query: Search query for academic papers
        num_results: Number of results to return (default: 10)
        year_filter: Optional year filter (e.g., "2020-" for 2020 onwards)

    Returns:
        JSON string with search results
    """
    url = "https://api.semanticscholar.org/graph/v1/paper/search/bulk"

    params = {
        "query": query,
        "fields": "title,url,abstract,year,authors,citationCount,publicationDate,venue",
        "limit": min(num_results, 100),
    }

    if year_filter:
        params["year"] = year_filter

    # API key optional but recommended for higher rate limits
    headers = {}
    api_key = os.getenv("SEMANTIC_SCHOLAR_API_KEY")
    if api_key:
        headers["x-api-key"] = api_key

    try:
        response = requests.get(url, params=params, headers=headers, timeout=10)
        response.raise_for_status()
        data = response.json()

        results = []
        for paper in data.get("data", []):
            # Handle None values for abstract and authors
            abstract = paper.get("abstract") or ""
            authors = paper.get("authors") or []

            results.append(
                {
                    "title": paper.get("title", "N/A"),
                    "url": paper.get("url", ""),
                    "abstract": abstract[:500],  # Truncate abstract
                    "year": paper.get("year"),
                    "authors": [a.get("name", "") for a in authors[:3]],  # First 3 authors
                    "citations": paper.get("citationCount", 0),
                    "venue": paper.get("venue", ""),
                }
            )

        return json.dumps({"results": results, "total": len(results)})

    except Exception as e:
        return json.dumps({"error": f"Semantic Scholar search failed: {str(e)}"})


def write_to_scratch_pad(url: str, title: str, content: str, scratch_pad_file: str):
    """Write extracted content to scratch pad file in JSONL format."""
    try:
        source_id = 1
        if os.path.exists(scratch_pad_file):
            with open(scratch_pad_file, "r", encoding="utf-8") as f:
                source_id = sum(1 for line in f if line.strip()) + 1

        data = {
            "source_id": source_id,
            "url": url,
            "title": title,
            "content": content,
            "timestamp": datetime.now().isoformat(),
        }

        # Ensure directory exists
        Path(scratch_pad_file).parent.mkdir(parents=True, exist_ok=True)

        with open(scratch_pad_file, "a", encoding="utf-8") as f:
            f.write(json.dumps(data, ensure_ascii=False) + "\n")

        return {"success": True, "source_id": source_id, "file": scratch_pad_file}
    except Exception as e:
        return {"success": False, "error": str(e)}


def read_scratch_pad_content(scratch_pad_file: str):
    """Read all content from scratch pad file."""
    sources = []
    if os.path.exists(scratch_pad_file):
        with open(scratch_pad_file, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    sources.append(json.loads(line.strip()))
    return {"sources": sources, "total": len(sources)}


def write_file(file_path: str, content: str):
    """Write content to a file."""
    try:
        Path(file_path).parent.mkdir(parents=True, exist_ok=True)
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(content)
        return {"success": True, "file_path": file_path}
    except Exception as e:
        return {"success": False, "error": str(e)}


def read_file(file_path: str):
    """Read content from a file."""
    try:
        if not os.path.exists(file_path):
            return {"success": False, "error": f"File not found: {file_path}"}

        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()
        return {"success": True, "content": content, "file_path": file_path}
    except Exception as e:
        return {"success": False, "error": str(e)}


# ============================================================================
# AGENT DESCRIPTIONS
# ============================================================================

QUERY_ANALYZER_DESC = """You create a research brief directly from user input.

Your role:
1. Ask the user about their research topic
2. Based on their response, create a clear research brief with:
   - Research topic/question
   - Key areas of focus
   - What information is needed
3. Send the research brief to WebRetrievalAgent

Be concise and direct. After one question to the user, create the research brief."""

WEB_RETRIEVAL_DESC = """You find research papers using Semantic Scholar.

Your role:
1. Receive research brief from QueryAnalyzer
2. Create one focused search query for Semantic Scholar
3. Execute search using tool_semantic_scholar_search
4. Select the top 2 most relevant papers based on citations, title, and abstract
5. Invoke BOTH Orchestrator1 and Orchestrator2 in PARALLEL using parallel_invoke
6. Provide BOTH orchestrators with:
   - The research brief
   - Top 3 paper URLs (url_1, url_2, url_3)
   - Their respective scratch_pad_path and report_path

Available tools: tool_semantic_scholar_search
Use parallel_invoke to send to both orchestrators simultaneously."""

ORCHESTRATOR_DESC = """You manage one research workflow stream.

Your role:
1. Receive research brief and 3 URLs (url_1, url_2, url_3) with your file paths
2. Invoke BrowserAgent for BOTH URLs in parallel
3. After browser agents complete, invoke SynthesizerAgent with your paths
4. After synthesizer completes, invoke FinalSynthesizer with your report_path
5. Return completion signal

You coordinate: Browser agents (parallel) → Synthesizer → FinalSynthesizer
Pass the correct file paths to each agent.
If BrowserAgent failed, do not try again. Make sure you pass both `url` and `scratch_pad_path` to BrowserAgent."""

BROWSER_DESC = """You extract content from a single web page.

Your role:
1. Receive URL and scratch_pad_path
2. Extract content from the URL
3. Clean and format content (markdown format, relevant text only)
4. Save to scratch_pad using write_to_scratch_pad tool
5. Return success/failure status

Extract only content relevant to the research topic."""

SYNTHESIZER_DESC = """You create a research report from collected sources.

Your role:
1. Receive scratch_pad_path and report_path
2. Read all sources using read_scratch_pad_content tool
3. Analyze and organize the information
4. Create a well-structured markdown report with:
   - Executive summary
   - Main findings
   - Citations to sources (use [source_id] format)
   - References section
5. Save report using write_file tool to report_path

Return the report_path when done."""

FINAL_SYNTHESIZER_DESC = """You create the final research report from multiple sources.

Your role:
1. Receive two report paths (report_1.md, report_2.md)
2. Read both reports using read_file tool
3. Synthesize the BEST final report by:
   - Combining insights from both reports
   - Creating a coherent narrative
   - Removing duplicates
4. Structure with:
   - Executive Summary
   - Key Findings
   - Detailed Analysis
   - References
5. Save to final_report_path using write_file tool

Create a comprehensive research report that synthesizes both streams."""


# ============================================================================
# MAIN EXECUTION
# ============================================================================


async def run_simple_research():
    """Run the simplified multi-agent research workflow."""

    load_dotenv()

    OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
    if not OPENROUTER_API_KEY:
        raise ValueError("OPENROUTER_API_KEY required")

    # Single model configuration - Gemini 2.5 Flash
    model_config = ModelConfig(
        type="api",
        provider="openrouter",
        name="google/gemini-2.5-flash",
        temperature=0.7,
        thinking_budget=6000,
        max_tokens=16000,
        api_key=OPENROUTER_API_KEY,
    )

    # Create output directory
    research_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_directory = f"./tmp/research_{research_id}"
    Path(output_directory).mkdir(parents=True, exist_ok=True)
    print(f"\n📁 Research directory: {output_directory}")

    # File paths for 2 parallel research streams
    scratch_paths = [
        f"{output_directory}/scratch_pad_1.jsonl",
        f"{output_directory}/scratch_pad_2.jsonl",
    ]
    report_paths = [
        f"{output_directory}/report_1.md",
        f"{output_directory}/report_2.md",
    ]
    final_report_path = f"{output_directory}/final_report.md"

    # Create agents
    print("\n🤖 Creating agents...")

    query_analyzer = Agent(
        model_config=model_config,
        goal="Create research briefs from user input",
        instruction=QUERY_ANALYZER_DESC,
        name="QueryAnalyzer",
    )

    web_retrieval = Agent(
        model_config=model_config,
        goal="Find research papers using Semantic Scholar",
        instruction=WEB_RETRIEVAL_DESC,
        name="WebRetrievalAgent",
        tools={
            "tool_semantic_scholar_search": tool_semantic_scholar_search,
        },
    )

    # Create 2 orchestrator instances
    orchestrators = []
    for i in range(1, 3):
        orchestrators.append(
            Agent(
                model_config=model_config,
                goal="Manage one research workflow stream",
                instruction=ORCHESTRATOR_DESC,
                name=f"Orchestrator{i}",
                memory_retention="session",
            )
        )

    # Browser agent pools (one per orchestrator)
    browser_pools = []
    for i in range(1, 3):
        browser_pool = await AgentPool.create_async(
            agent_class=BrowserAgent,
            num_instances=3,
            model_config=model_config,
            goal="Extract content from web pages",
            instruction=BROWSER_DESC,
            name=f"BrowserAgent{i}",
            headless=True,
            memory_retention="single_run",
            tools={"write_to_scratch_pad": write_to_scratch_pad},
        )
        AgentRegistry.register_pool(browser_pool)
        browser_pools.append(browser_pool)

    # Synthesizer agents
    synthesizers = []
    for i in range(1, 3):
        synthesizers.append(
            Agent(
                model_config=model_config,
                goal="Create research reports from collected sources",
                instruction=SYNTHESIZER_DESC,
                name=f"Synthesizer{i}",
                tools={
                    "read_scratch_pad_content": read_scratch_pad_content,
                    "write_file": write_file,
                },
            )
        )

    # Final synthesizer
    final_synthesizer = Agent(
        model_config=model_config,
        goal="Create the final research report from multiple sources",
        instruction=FINAL_SYNTHESIZER_DESC,
        name="FinalSynthesizer",
        tools={"read_file": read_file, "write_file": write_file},
    )

    # Helper: surface filesystem context to each orchestrator once
    def _attach_workspace_context():
        for idx, orchestrator in enumerate(orchestrators, start=1):
            stream_note = f"You run stream #{idx}\n" f"scratch_pad_path (to be sent to BrowserAgent): {scratch_paths[idx - 1]}\n" f"report_path (to be sent to Synthesizer): {report_paths[idx - 1]}\n"
            message_id = orchestrator.memory.add(role="system", content=stream_note)
            orchestrator._context_selector.save_selection(
                {"message_ids": [message_id]},
                key="workspace_paths",
            )

        # Attach context to final_synthesizer
        synthesizer_note = f"You will create the final report at path:\n" f"{final_report_path}\n" f"Source reports (you need to read these resport to generate the final report):\n" f"1: {report_paths[0]}\n" f"2: {report_paths[1]}"
        message_id = final_synthesizer.memory.add(role="system", content=synthesizer_note)
        final_synthesizer._context_selector.save_selection({"message_ids": [message_id]}, key="workspace_paths")

    _attach_workspace_context()

    # Define topology
    print("\n🔗 Defining topology...")

    topology = {
        "agents": [
            {"name": "User", "type": "user"},
            "QueryAnalyzer",
            "WebRetrievalAgent",
            "Orchestrator1",
            "Orchestrator2",
            "BrowserAgent1",
            "BrowserAgent2",
            "Synthesizer1",
            "Synthesizer2",
            "FinalSynthesizer",
        ],
        "flows": [
            # User interaction
            "User -> QueryAnalyzer",
            "QueryAnalyzer -> User",
            # Query analyzer to web retrieval
            "QueryAnalyzer -> WebRetrievalAgent",
            # Web retrieval to parallel orchestrators
            "WebRetrievalAgent -> Orchestrator1",
            "WebRetrievalAgent -> Orchestrator2",
            # Orchestrator 1 workflow
            "Orchestrator1 -> BrowserAgent1",
            "BrowserAgent1 -> Orchestrator1",
            "Orchestrator1 -> Synthesizer1",
            "Synthesizer1 -> Orchestrator1",
            # Orchestrator 2 workflow
            "Orchestrator2 -> BrowserAgent2",
            "BrowserAgent2 -> Orchestrator2",
            "Orchestrator2 -> Synthesizer2",
            "Synthesizer2 -> Orchestrator2",
            # Both orchestrators to final synthesizer
            "Orchestrator1 -> FinalSynthesizer",
            "Orchestrator2 -> FinalSynthesizer",
            # Final synthesizer to user
            "FinalSynthesizer -> User",
        ],
        "rules": [
            "timeout(1800)",  # 30 min timeout
            "max_steps(150)",
        ],
    }

    # Run research workflow
    print("\n🚀 Starting research workflow...\n")

    try:
        # Prepare context with all file paths
        context = {
            "output_directory": output_directory,
            "scratch_paths": scratch_paths,
            "report_paths": report_paths,
            "final_report_path": final_report_path,
            "research_id": research_id,
        }

        # Run with Orchestra (auto-creates CommunicationManager for user interaction)
        from marsys.coordination.config import ExecutionConfig, StatusConfig

        execution_config = ExecutionConfig(
            user_interaction="terminal",  # Auto-creates CommunicationManager
            user_first=True,
            initial_user_msg="Hello! I'll help you research academic papers. What topic would you like to research?",
            status=StatusConfig.from_verbosity(1),
            convergence_timeout=600.0,  # 10 min for parallel workflows
        )

        result = await Orchestra.run(
            task="Begin research workflow",
            topology=topology,
            agent_registry=AgentRegistry,
            context={"request_context": context},
            execution_config=execution_config,
            max_steps=150,
            verbosity=2,
        )

        # Display results
        print("\n" + "=" * 80)
        print("📊 RESEARCH COMPLETE")
        print("=" * 80)
        print(f"Success: {result.success}")
        print(f"Total steps: {result.total_steps}")
        print(f"Duration: {result.total_duration:.2f}s")
        print(f"Branches executed: {len(result.branch_results)}")

        if result.success:
            print(f"\n✅ Final report saved to: {final_report_path}")
            print(f"\n📄 Individual reports:")
            for i, path in enumerate(report_paths, 1):
                if os.path.exists(path):
                    print(f"   - Report {i}: {path}")

        # Browser pool statistics
        for i, browser_pool in enumerate(browser_pools, 1):
            stats = browser_pool.get_statistics()
            print(f"\n🌐 Browser Pool {i} Statistics:")
            print(f"   - Total allocations: {stats['total_allocations']}")
            print(f"   - Peak concurrent usage: {stats['peak_concurrent_usage']}")
            print(f"   - Average wait time: {stats['average_wait_time']:.2f}s")

    except KeyboardInterrupt:
        print("\n\n⚠️ Research interrupted by user")
    except Exception as e:
        print(f"\n❌ Error during research: {e}")
        import traceback

        traceback.print_exc()
    finally:
        # Cleanup
        print("\n🧹 Cleaning up...")
        cleanup_tasks = []
        for browser_pool in browser_pools:
            cleanup_tasks.append(browser_pool.cleanup())

        for agent in (
            [
                query_analyzer,
                web_retrieval,
                final_synthesizer,
            ]
            + orchestrators
            + synthesizers
        ):
            if hasattr(agent, "cleanup"):
                cleanup_tasks.append(agent.cleanup())

        await asyncio.gather(*cleanup_tasks, return_exceptions=True)
        print("✅ Cleanup complete\n")


def main():
    """Main entry point."""
    print("\n🔬 Simplified Multi-Agent Deep Research System")
    print("=" * 80)
    print("This system searches for academic papers and synthesizes them into a report.")
    print("=" * 80)

    asyncio.run(run_simple_research())


if __name__ == "__main__":
    main()
