#!/usr/bin/env python3
"""
Multi-Agent Research Team Example

This example demonstrates a multi-agent system using A2A protocol communication.
Three specialized agents collaborate on research tasks:

1. Researcher Agent (port 8000) - Gathers information and data
2. Analyst Agent (port 8001) - Analyzes data and creates insights
3. Writer Agent (port 8002) - Formats and presents final reports

The agents communicate with each other using the A2A protocol over HTTP.
"""

import json
import random
import threading
import time
from typing import Any

import httpx

# Core dependencies
try:
    from agents import Agent, function_tool

    from a2a_openai_agents import A2ASkillConfig, A2AWrapper, run_a2a_wrapper_server
except ImportError:
    print("Dependencies not installed. Please run: pip install a2a-openai-agents")
    exit(1)


# Researcher Agent Tools
@function_tool
def web_search(query: str) -> str:
    """Simulate web search to gather information."""
    print(f"[RESEARCHER] Searching web for: {query}")
    # Simulate research results
    results = [
        f"Research finding 1 about {query}: Important data point A",
        f"Research finding 2 about {query}: Key insight B",
        f"Research finding 3 about {query}: Critical factor C",
    ]
    return f"Web search results for '{query}':\n" + "\n".join(results)


@function_tool
def gather_statistics(topic: str) -> str:
    """Simulate gathering statistical data."""
    print(f"[RESEARCHER] Gathering statistics for: {topic}")
    # Simulate statistical data
    stats = {
        "metric_1": random.randint(50, 100),
        "metric_2": random.randint(200, 500),
        "trend": random.choice(["increasing", "decreasing", "stable"]),
        "confidence": random.randint(85, 99),
    }
    return f"Statistics for {topic}: {json.dumps(stats, indent=2)}"


# Analyst Agent Tools
@function_tool
def analyze_data(data: str) -> str:
    """Analyze provided data and extract insights."""
    print(f"[ANALYST] Analyzing data: {data[:100]}...")
    # Simulate data analysis
    insights = [
        "Pattern detected: Strong correlation between variables",
        "Trend analysis shows consistent growth over time",
        "Risk assessment indicates moderate confidence level",
        "Recommendation: Focus on key metrics for optimization",
    ]
    return "Data Analysis Results:\n" + "\n".join(f"• {insight}" for insight in insights)


@function_tool
def create_visualization(data: str) -> str:
    """Create data visualization description."""
    print("[ANALYST] Creating visualization for data")
    viz_types = ["bar chart", "line graph", "pie chart", "scatter plot"]
    return f"Visualization created: {random.choice(viz_types)} showing key trends and patterns in the data"


# Writer Agent Tools
@function_tool
def format_report(content: str) -> str:
    """Format content into a professional report structure."""
    print(f"[WRITER] Formatting report with content length: {len(content)} chars")
    formatted = f"""
# Research Report

## Executive Summary
{content[:200]}...

## Detailed Analysis
{content}

## Conclusions and Recommendations
Based on the analysis, we recommend continued monitoring of key metrics and strategic focus on identified opportunities.

---
*Report generated by Multi-Agent Research Team*
"""
    return formatted.strip()


@function_tool
def proofread_content(text: str) -> str:
    """Proofread and improve text quality."""
    print(f"[WRITER] Proofreading {len(text)} characters of content")
    return f"✓ Proofread and polished: {text}"


# A2A Communication Helper
async def call_agent_skill(agent_url: str, skill_name: str, parameters: dict) -> str:
    """Call another agent's skill via A2A protocol."""
    print(f"[A2A] Calling {agent_url} skill '{skill_name}' with params: {parameters}")

    async with httpx.AsyncClient(timeout=30.0) as client:
        # Construct A2A message
        message_data = {
            "messages": [
                {
                    "role": "user",
                    "parts": [
                        {
                            "type": "text",
                            "text": f"Use the {skill_name} skill with parameters: {json.dumps(parameters)}",
                        }
                    ],
                }
            ]
        }

        try:
            response = await client.post(f"{agent_url}/message", json=message_data)
            response.raise_for_status()

            # Parse A2A response
            task_data = response.json()

            # Extract result from task messages
            if "messages" in task_data and len(task_data["messages"]) > 1:
                agent_message = task_data["messages"][-1]  # Last message should be agent response
                if "parts" in agent_message and len(agent_message["parts"]) > 0:
                    return agent_message["parts"][0].get("text", "No response")

            return "No valid response received"

        except Exception as e:
            print(f"[ERROR] Failed to call {agent_url}: {e}")
            return f"Error calling agent: {e}"


def create_researcher_agent():
    """Create the Researcher agent with information gathering capabilities."""
    researcher_brain = Agent(
        name="ResearchSpecialist",
        instructions="""You are a research specialist focused on gathering comprehensive information.
        Use your web_search and gather_statistics tools to collect relevant data.
        Be thorough and provide detailed, factual information.""",
        model="gpt-4o-mini",
        tools=[web_search, gather_statistics],
    )

    return A2AWrapper(
        openai_agent=researcher_brain,
        a2a_name="ResearchAgent",
        a2a_description="Specialized agent for research and data gathering",
    )


def create_analyst_agent():
    """Create the Analyst agent with data analysis capabilities."""
    analyst_brain = Agent(
        name="DataAnalyst",
        instructions="""You are a data analyst who processes information and creates insights.
        Use your analyze_data and create_visualization tools to process research findings.
        Focus on extracting meaningful patterns and actionable insights.""",
        model="gpt-4o-mini",
        tools=[analyze_data, create_visualization],
    )

    analyst_wrapper = A2AWrapper(
        openai_agent=analyst_brain,
        a2a_name="AnalystAgent",
        a2a_description="Specialized agent for data analysis and insights",
    )

    # Add custom skill that uses the Researcher
    async def research_and_analyze(
        wrapper_instance: A2AWrapper, params: dict[str, Any]
    ) -> dict[str, Any]:
        """Research a topic and analyze the findings."""
        topic = params.get("topic", "")
        if not topic:
            return {"error": "Topic parameter is required"}

        try:
            # Step 1: Call Researcher for web search
            search_results = await call_agent_skill(
                "http://localhost:8000", "web_search", {"query": topic}
            )

            # Step 2: Call Researcher for statistics
            stats_results = await call_agent_skill(
                "http://localhost:8000", "gather_statistics", {"topic": topic}
            )

            # Step 3: Analyze the combined data
            combined_data = f"Web Search Results:\n{search_results}\n\nStatistics:\n{stats_results}"
            analysis = await wrapper_instance.process_prompt_with_internal_agent(
                f"Analyze this data using the analyze_data tool: {combined_data}"
            )

            return {
                "topic": topic,
                "research_data": combined_data,
                "analysis": analysis,
                "status": "completed",
            }

        except Exception as e:
            return {"error": f"Research and analysis failed: {str(e)}"}

    # Override skills to include the cross-agent skill
    analyst_wrapper.a2a_skills.append(
        A2ASkillConfig(
            name="research_and_analyze",
            description="Research a topic using the Research agent and analyze the findings",
            handler=research_and_analyze,
            parameters={
                "type": "object",
                "properties": {
                    "topic": {"type": "string", "description": "The topic to research and analyze"}
                },
                "required": ["topic"],
            },
        )
    )

    return analyst_wrapper


def create_writer_agent():
    """Create the Writer agent with document formatting capabilities."""
    writer_brain = Agent(
        name="ReportWriter",
        instructions="""You are a professional report writer who creates well-formatted documents.
        Use your format_report and proofread_content tools to create polished output.
        Focus on clarity, structure, and professional presentation.""",
        model="gpt-4o-mini",
        tools=[format_report, proofread_content],
    )

    writer_wrapper = A2AWrapper(
        openai_agent=writer_brain,
        a2a_name="WriterAgent",
        a2a_description="Specialized agent for document writing and formatting",
    )

    # Add custom skill that uses both Researcher and Analyst
    async def create_full_report(
        wrapper_instance: A2AWrapper, params: dict[str, Any]
    ) -> dict[str, Any]:
        """Create a complete research report by coordinating with other agents."""
        topic = params.get("topic", "")
        if not topic:
            return {"error": "Topic parameter is required"}

        try:
            # Step 1: Call Analyst to research and analyze (which will call Researcher)
            analysis_result = await call_agent_skill(
                "http://localhost:8010",  # Use port 8010 to avoid conflicts
                "research_and_analyze",
                {"topic": topic},
            )

            # Step 2: Format the analysis into a report
            report = await wrapper_instance.process_prompt_with_internal_agent(
                f"Create a professional report using the format_report tool with this analysis: {analysis_result}"
            )

            # Step 3: Proofread the final report
            final_report = await wrapper_instance.process_prompt_with_internal_agent(
                f"Proofread this report using the proofread_content tool: {report}"
            )

            return {"topic": topic, "final_report": final_report, "status": "completed"}

        except Exception as e:
            return {"error": f"Report creation failed: {str(e)}"}

    # Override skills to include the orchestration skill
    writer_wrapper.a2a_skills.append(
        A2ASkillConfig(
            name="create_full_report",
            description="Create a complete research report by coordinating with Research and Analysis agents",
            handler=create_full_report,
            parameters={
                "type": "object",
                "properties": {
                    "topic": {"type": "string", "description": "The topic for the research report"}
                },
                "required": ["topic"],
            },
        )
    )

    return writer_wrapper


def run_agent_server(agent_wrapper, port: int):
    """Run an agent server on the specified port."""
    print(f"Starting {agent_wrapper._a2a_name} on port {port}")
    run_a2a_wrapper_server(agent_wrapper, port=port, host="localhost")


def main():
    """Start the multi-agent research team."""
    print("=== Multi-Agent Research Team ===")
    print("Setting up three specialized agents...")

    # Create agents
    researcher = create_researcher_agent()
    analyst = create_analyst_agent()
    writer = create_writer_agent()

    print("\nAgent capabilities:")
    print(
        f"• Researcher (port 8000): {len(researcher.a2a_skills)} skills - {[s.name for s in researcher.a2a_skills]}"
    )
    print(
        f"• Analyst (port 8001): {len(analyst.a2a_skills)} skills - {[s.name for s in analyst.a2a_skills]}"
    )
    print(
        f"• Writer (port 8002): {len(writer.a2a_skills)} skills - {[s.name for s in writer.a2a_skills]}"
    )

    # Start agents in separate threads (using different ports to avoid conflicts)
    threads = [
        threading.Thread(target=run_agent_server, args=(researcher, 8000), daemon=True),
        threading.Thread(target=run_agent_server, args=(analyst, 8010), daemon=True),
        threading.Thread(target=run_agent_server, args=(writer, 8020), daemon=True),
    ]

    for thread in threads:
        thread.start()

    # Wait for servers to start
    print("\nStarting agent servers...")
    time.sleep(5)

    print("\n=== Multi-Agent System Ready ===")
    print("Agent endpoints:")
    print("• Researcher: http://localhost:8000")
    print("• Analyst: http://localhost:8001")
    print("• Writer: http://localhost:8002")

    print("\nTo test the system, try:")
    print("curl -X POST http://localhost:8002/message \\")
    print('  -H "Content-Type: application/json" \\')
    print(
        '  -d \'{"messages":[{"role":"user","parts":[{"type":"text","text":"Create a full report on artificial intelligence trends"}]}]}\''
    )

    print("\nPress Ctrl+C to stop all agents")

    try:
        # Keep main thread alive
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\nShutting down multi-agent system...")


if __name__ == "__main__":
    main()
