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
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()


async def save_report_to_file(content: str) -> Dict[str, Any]:
    """
    Save the research report to a markdown file.

    Args:
        content: The markdown content of the report

    Returns:
        Dictionary with success status and filename
    """
    from datetime import datetime
    from pathlib import Path

    # Create output directory if it doesn't exist
    output_dir = Path("examples/real_world/output")
    output_dir.mkdir(parents=True, exist_ok=True)

    # Generate filename with timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"deep_research_report_{timestamp}.md"
    filepath = output_dir / filename

    try:
        # Save the report to file
        with open(filepath, "w", encoding="utf-8") as f:
            f.write(content)

        return {
            "success": True,
            "message": f"Report successfully saved to {filepath}",
            "filename": filename,
            "filepath": str(filepath),
        }
    except Exception as e:
        return {
            "success": False,
            "message": f"Failed to save report: {str(e)}",
            "error": str(e),
        }


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-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


class LoggingBrowserAgent(BrowserAgent):
    """BrowserAgent with built-in logging of all interactions"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.call_count = 0
        self.AGENT_RESPONSE_FILE = (
            Path("examples/real_world/output")
            / f"deep-research-responses={datetime.now().strftime("%Y-%m-%d %H:%M:%S")}.txt"
        )

    async def run_step(self, request, context):
        """Override run_step to add logging"""
        self.call_count += 1

        # Prepare log entry
        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": None,
            "response": None,
            "duration_seconds": None,
        }

        # Log the request
        if isinstance(request, dict):
            log_entry["request"] = request
        else:
            log_entry["request"] = str(request)

        # Call the parent's run_step
        start_time = datetime.now()
        result = await super().run_step(request, context)
        duration = (datetime.now() - start_time).total_seconds()
        log_entry["duration_seconds"] = round(duration, 2)

        # Log the response
        log_entry["response"] = result

        # Write to the response file
        AGENT_RESPONSE_FILE = (
            Path("examples/real_world/output") / "deep-research-responses.txt"
        )
        with open(AGENT_RESPONSE_FILE, "a") as f:
            f.write("\n" + "=" * 100 + "\n")
            f.write(f"AGENT CALL #{self.call_count}: {self.name} (BrowserAgent)\n")
            f.write("=" * 100 + "\n\n")
            f.write(f"Agent: {self.name}\n")
            f.write(f"Called by: {log_entry['caller']}\n")
            f.write(f"Timestamp: {log_entry['timestamp']}\n")
            f.write(f"Duration: {duration:.2f} seconds\n\n")

            f.write("REQUEST:\n")
            f.write("-" * 50 + "\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")
            f.write("-" * 50 + "\n")
            try:
                f.write(json.dumps(log_entry["response"], indent=2, default=str) + "\n")
            except:
                f.write(str(log_entry["response"]) + "\n")

            f.write("\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 create_agents(configs: Dict[str, ModelConfig]) -> Dict[str, Agent]:
    agents = {}

    agents["scoper_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="scoper_agent",
        goal="Refine vague user queries into clear, scoped research objectives",
        instruction="""You are the first agent in the deep research workflow. Your role is to refine vague user queries into clear, scoped research objectives.

RESPONSIBILITIES:
- Analyze the initial user topic for clarity and specificity
- Transform the topic into a well-defined research objective
- Pass the refined topic to planner_agent to continue the workflow
- Ensure the research scope is clear and achievable

IMPORTANT: Create a clear, specific research topic before proceeding to planner_agent.""",
    )

    agents["planner_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="planner_agent",
        goal="Orchestrate the research process and manage task delegation to other agents",
        instruction="""You are the central coordinator in the deep research workflow. Your role is to orchestrate the research process and manage task delegation.

RESPONSIBILITIES:
- Receive the refined research topic from scoper_agent (initial planning)
- OR receive requests for additional resources from evaluator_agent (refinement planning)
- Create comprehensive research plans with specific objectives and criteria
- Direct research_agent to generate targeted search queries
- Ensure the research stays focused and productive
- Respect the resource_limit specified in the context
- Handle both initial research planning and refinement requests

RESEARCH PLANNING STRATEGIES:
- Define clear research objectives and success criteria
- Identify key research questions to answer
- Determine the scope and depth of research needed
- Plan for different types of sources (academic, industry, news, reports)
- Consider timeline and resource constraints
- Plan to find exactly the number of resources specified in resource_limit

IMPORTANT: Be adaptive. Handle both initial planning and refinement requests efficiently. Always respect the resource_limit and focus on quality.""",
    )

    agents["research_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="research_agent",
        goal="Generate effective search queries that translate research objectives into searchable terms",
        instruction="""You are responsible for generating effective search queries based on the research plan. Your role is to translate research objectives into searchable terms.

RESPONSIBILITIES:
- Receive research objectives and criteria from planner_agent
- Handle both initial research planning and refinement requests
- Analyze the research topic to identify key concepts, synonyms, and related terms
- Generate search query variations to ensure comprehensive coverage
- Optimize queries for search engine effectiveness (use relevant keywords, phrases, and operators)
- Consider different search strategies (broad vs. specific, recent vs. historical, academic vs. commercial)
- Pass optimized search queries to retriever_agent for execution
- Respect the resource_limit from context and focus on remaining needs

INITIAL RESEARCH QUERY GENERATION:
- Generate comprehensive queries covering all research objectives
- Create 3-5 search queries to find the full resource_limit
- Cover different aspects and perspectives of the research topic
- Use specific, targeted keywords rather than broad terms

REFINEMENT QUERY GENERATION (when evaluator_agent needs more resources):
- Analyze feedback about what types of sources are still needed
- Generate targeted queries for specific gaps in research coverage
- Focus on missing aspects or underrepresented perspectives
- Create 2-3 focused queries to find the remaining resources needed
- Ensure new queries complement existing findings without duplication

RESOURCE LIMIT CONSIDERATION:
- The context contains a "resource_limit" parameter
- Track how many resources have already been found (context.resources_found)
- Generate enough search queries to find the remaining resources needed
- Focus on quality over quantity - better to have fewer, more targeted queries
- Avoid generating excessive queries that might exceed the resource limit

SEARCH QUERY STRATEGIES:
- Use specific, targeted keywords rather than broad terms
- Include relevant time periods (e.g., "2024", "2025")
- Add industry-specific terminology and jargon
- Use quotation marks for exact phrases
- Include location or region if relevant
- Consider academic vs. commercial sources
- Generate appropriate number of queries based on remaining resource needs

IMPORTANT: Your queries should be specific enough to find relevant results but broad enough to capture comprehensive information. Generate enough queries to find the remaining resources needed, not the total resource_limit.""",
    )

    agents["retriever_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="retriever_agent",
        goal="Execute search queries using Google Search API and extract relevant URLs",
        tools={"tool_google_search_api": tool_google_search_api},
        instruction="""You are responsible for executing search queries and finding relevant URLs. Your role is to use the Google Search API to discover candidate web pages for research.

RESPONSIBILITIES:
- Execute search queries received from research_agent using the Google Search API
- Extract ALL relevant URLs from search results
- Pass ALL URLs to browser_agent(s) for content extraction
- DO NOT save context or use context tools - process URLs immediately

SEARCH EXECUTION:
- Execute each search query from research_agent
- The tool returns a JSON list with objects containing 'url' field like:
  [{"title": "...", "content": "...", "url": "https://real-url.com"}, ...]
- Extract ALL URLs from the 'url' field in search results (up to 5 URLs max)
- NEVER make up URLs - use only the URLs returned by the search tool

CRITICAL WORKFLOW:
1. Execute ONE search query with Google Search API (use num_results=5)
2. Extract ALL URLs from the search results immediately
3. Invoke browser_agent for EACH URL in parallel (all at once)
4. DO NOT save context, DO NOT use preview_saved_context

IMPORTANT RULES:
- Process ALL URLs from search results in ONE batch
- DO NOT save or preview context - that breaks the workflow
- DO NOT process URLs one at a time - process them ALL together
- Browser agents will handle content extraction and pass to summarizer
- Each search should result in multiple browser_agent invocations (one per URL)""",
    )

    agents["browser_agent"] = await LoggingBrowserAgent.create_safe(
        model_config=configs["gemini_pro"],
        name="browser_agent",
        goal="Fetch and extract clean content from web pages for research analysis",
        input_schema={
            "type": "object",
            "properties": {
                "url": {
                    "type": "string",
                    "description": "The URL to extract content from",
                },
                "query": {
                    "type": "string",
                    "description": "What specific information to extract from the URL",
                },
            },
            "required": ["url", "query"],
        },
        instruction="""You are responsible for fetching and extracting content from web pages. Your role is to retrieve, clean, and prepare web content for analysis.

RESPONSIBILITIES:
- Fetch web page content from the URL provided
- Use the extract_content_from_url tool to get the page content
- After extracting content, invoke summarizer_agent with the extracted content

WORKFLOW:
1. Use extract_content_from_url tool to fetch and extract content from the URL
2. Once content is extracted, invoke summarizer_agent with the content

OUTPUT FORMAT (after extracting content):
{
  "next_action": "invoke_agent",
  "action_input": [
    {
      "agent_name": "summarizer_agent",
      "request": {
        "title": "extracted_title",
        "content": "extracted_content",
        "url": "original_url"
      }
    }
  ]
}

IMPORTANT:
- Always invoke summarizer_agent after successfully extracting content
- Include the title, content, and URL in the request to summarizer_agent
- Do NOT use JSON escape sequences in your response""",
        headless=True,
        viewport_width=1440,
        viewport_height=900,
        auto_screenshot=False,  # Don't need screenshots for content extraction
        timeout=10000,  # 10 second timeout for page operations
    )

    agents["evaluator_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="evaluator_agent",
        goal="Assess content quality and relevance, decide when sufficient research resources are gathered",
        instruction="""You are the quality control agent in the research workflow. Your role is to assess content relevance, quality, and value for the research objectives.

RESPONSIBILITIES:
- Evaluate content quality, relevance, and credibility from browser_agent
- Score content based on relevance to research topic, source authority, and content depth
- Filter out low-quality or irrelevant content
- Track the number of high-quality resources found (resources_found)
- Make conditional routing decisions based on resource count:
  * If resources_found < resource_limit: Request more resources from planner_agent
  * If resources_found >= resource_limit: Proceed to summarizer_agent and cross_check_agent
- Ensure only valuable content proceeds to the next stages of analysis

CONDITIONAL ROUTING LOGIC:
- Check context.resources_found vs context.resource_limit
- If resources_found < resource_limit:
  * Route to planner_agent with request for more resources
  * Include feedback on what types of sources are still needed
  * Specify how many more resources are required
- If resources_found >= resource_limit:
  * Route high-quality content to summarizer_agent for summarization
  * Route high-quality content to cross_check_agent for contradiction analysis
  * Proceed to final report generation


CRITICAL TERMINATION RULES:
- Only send feedback to planner_agent if quality score is below 60%
- Maximum 2 feedback iterations to prevent infinite loops
- If quality is above 60%, ALWAYS proceed to summarizer_agent + cross_check_agent
- Track iteration count and force termination after 2 rounds

IMPORTANT: Always check the resource count before making routing decisions. The goal is to find exactly the number of high-quality resources specified in resource_limit.""",
    )

    agents["summarizer_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="summarizer_agent",
        goal="Create concise, structured summaries of research content capturing key insights",
        instruction="""You are responsible for creating concise summaries of high-quality research content. Your role is to distill key information into digestible summaries.

RESPONSIBILITIES:
- Receive content that has been extracted from web pages
- Create comprehensive yet concise summaries that capture key findings, insights, and data points
- Maintain the original context and meaning while condensing information
- Structure summaries in a consistent format for easy integration
- Highlight critical information, statistics, and conclusions
- Ensure summaries are factual and accurately represent the source content

SUMMARY STRUCTURE:
- Key findings and insights
- Important statistics and data points
- Methodology and approach (if applicable)
- Conclusions and recommendations
- Source attribution and credibility notes
- Context and relevance to research topic

QUALITY GUIDELINES: 
- Focus on creating clear, factual summaries
- Include the source URL for each summary
- Process each piece of content thoroughly
- Maintain objectivity and accuracy""",
    )

    # Cross-check agent commented out - not used in current topology
    # agents["cross_check_agent"] = LoggingAgent(
    #     model_config=configs["gemini_pro"],
    #     agent_name="cross_check_agent",
    #     description="""You are responsible for analyzing research content for contradictions, corroborations, and consistency. Your role is to identify conflicting information and validate findings across multiple sources.
    #
    # RESPONSIBILITIES:
    # - Receive high-quality content that has been validated by evaluator_agent
    # - Analyze content for contradictions, inconsistencies, and conflicting information
    # - Identify corroborating evidence and consensus across multiple sources
    # - Detect bias, outdated information, or unreliable claims
    # - Compare findings, statistics, and conclusions across different sources
    # - Flag areas of agreement and disagreement in the research
    # - Provide analysis of source reliability and credibility
    # - Pass contradiction/corroboration analysis to synthesizer_agent for final report compilation
    #
    # PARALLEL EXECUTION:
    # - You work in parallel with summarizer_agent (as per topology rules)
    # - Both agents receive the same content from evaluator_agent simultaneously
    # - Work independently and efficiently to complete your analysis
    # - Focus solely on validation - summarizer_agent handles content summarization
    # - Your analysis will be combined with summarizer_agent's summaries by synthesizer_agent
    #
    # VALIDATION FOCUS AREAS:
    # - Statistical consistency across sources
    # - Factual accuracy and verification
    # - Temporal relevance and currency
    # - Source credibility and authority
    # - Bias detection and assessment
    # - Methodology differences and implications
    # - Geographic or demographic variations
    # - Industry vs. academic perspectives
    #
    # IMPORTANT: Focus on identifying inconsistencies and validating claims. The synthesizer_agent will integrate your analysis with summarizer_agent's summaries.""",
    # )

    agents["synthesizer_agent"] = LoggingAgent(
        model_config=configs["gemini_pro"],
        name="synthesizer_agent",
        goal="Synthesize all research summaries into a comprehensive markdown report and save to file",
        tools={"save_report_to_file": save_report_to_file},
        instruction="""You are the final agent responsible for creating and saving the comprehensive research report. Your role is to synthesize all research summaries into a coherent, well-structured markdown document and save it to a file.

RESPONSIBILITIES:
- You may receive multiple summaries from different sources
- Integrate all research findings into a coherent narrative
- Identify common themes and unique insights across different sources
- Organize information logically with clear sections and subsections
- Highlight key insights, trends, and conclusions
- Create a professional, well-formatted markdown report
- Ensure the report answers the original research question comprehensively
- SAVE the report to a markdown file in the output directory

REPORT STRUCTURE:
- Executive Summary
- Introduction with research objectives
- Main findings organized by topic/theme
- Key statistics and data points
- Methodology notes (including number of sources analyzed)
- Conclusions and recommendations
- References and source URLs

FINAL OUTPUT:
- After creating the report, you MUST save it to a file
- Use the tool: save_report_to_file to save the markdown content
- The file will be automatically saved with timestamp in the output directory
- Return a final_response confirming the report has been saved

QUALITY STANDARDS: 
- Generate the complete markdown report based on all received summaries
- Include source URLs and references from all research content
- Every claim, statistic, or finding should be properly attributed
- Ensure coherent flow and logical organization
- Maintain professional tone and formatting
- Note when findings are corroborated across multiple sources

IMPORTANT: As the final agent, you must save the report and return final_response to complete the workflow.""",
    )

    return agents


async def run_deep_research(
    # topic: str,
    comm_manager: CommunicationManager,
    session_id: str,
    resource_limit: int = 10,
) -> Dict[str, Any]:
    # Create a clean agent response log file in the output directory
    output_dir = Path("./examples/real_world/output")
    output_dir.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    agent_response_file = output_dir / f"deep_research_{timestamp}_execution.json"

    # Initialize the response file with header
    with open(agent_response_file, "w") as f:
        f.write("=" * 100 + "\n")
        f.write(
            f"DEEP RESEARCH EXECUTION LOG - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        )
        # f.write(f"Topic: {topic}\n")
        f.write(f"Session ID: {session_id}\n")
        f.write("Communication Mode: Terminal (Synchronous)\n")
        f.write("=" * 100 + "\n\n")

    topology = {
        "agents": [
            # {"name": "User", "type": "user"},  # Skip User node for automated execution
            "scoper_agent",
            "planner_agent",
            "research_agent",
            "retriever_agent",
            "browser_agent",
            # "evaluator_agent",
            "summarizer_agent",
            # "cross_check_agent",  # Commented out - has no edges defined in topology
            "synthesizer_agent",  # This will be a convergence point and final agent
        ],
        "flows": [
            # "User -> scoper_agent",  # Skip User interaction
            # "scoper_agent -> User",  # Skip User interaction
            "scoper_agent -> planner_agent",
            "planner_agent -> research_agent",
            "research_agent -> retriever_agent",
            "retriever_agent -> browser_agent",  # Retriever spawns multiple browser_agents
            "browser_agent -> summarizer_agent",  # Each browser_agent leads to summarizer
            # "research_agent -> evaluator_agent",
            # "evaluator_agent -> research_agent",  # Bidirectional - evaluator can request more resources
            # "evaluator_agent -> planner_agent",  # Conditional - only if more resources needed
            # "evaluator_agent -> summarizer_agent",  # If enough resources, proceed to analysis
            # "evaluator_agent -> cross_check_agent",  # If enough resources, proceed to analysis
            "summarizer_agent -> synthesizer_agent",  # CONVERGENCE POINT: All summarizers converge here
            # "cross_check_agent -> synthesizer_agent",
            # synthesizer_agent is now the final agent - it saves report to file
        ],
        "rules": [
            "timeout(2000)",
            "max_steps(10)",
            "parallel(planner_agent)",
            "parallel(retriever_agent, browser_agent)",  # Enable parallel execution
            "convergence_point(synthesizer_agent)",  # Mark synthesizer as convergence point
            # This ensures synthesizer waits for ALL parallel summarizer instances to complete
            # before executing, aggregating all their results into a single synthesis
            # "conditional(evaluator_agent -> planner_agent, resources_found < resource_limit)",
            # "conditional(evaluator_agent -> summarizer_agent, resources_found >= resource_limit)",
        ],
        "entry_point": "scoper_agent",  # Start with scoper_agent after User
        "exit_point": "synthesizer_agent",  # synthesizer_agent is now the final agent
    }

    # Define the research task with the specific topic
    task = f"""Conduct comprehensive deep research on the topic: synthetic data robotics academic paper 2025

The research should follow this workflow:
1. Scope and clarify the research topic
2. Plan the research approach to find exactly {resource_limit} high-quality resources  
3. Generate targeted search queries
4. Retrieve relevant URLs (limit: {resource_limit} sources total)
5. Extract and analyze web content
6. Summarize findings
7. Synthesize into a comprehensive report
8. Save the final report

IMPORTANT CONSTRAINTS:
- Find exactly {resource_limit} high-quality, relevant resources
- Focus on quality over quantity
- Complete the workflow efficiently
"""
    # logger.info(f"Starting deep research for topic: {topic}")

    try:
        # Create context
        context = {
            # "topic": topic,
            "research_type": "comprehensive",
            "output_format": "markdown",
            "response_file": str(agent_response_file),
            "caller_agent": "User",  # Initial caller
            "session_id": session_id,  # For user communication
            "communication_mode": "sync",  # Synchronous terminal mode
            "resource_limit": resource_limit,  # Pass resource limit
            "resources_found": 0,  # Track how many resources have been found
            "resources_processed": 0,  # Track how many resources have been processed
            "resource_tracking": {
                "total_found": 0,
                "high_quality_count": 0,
                "current_batch": [],
                "evaluation_status": "pending",
            },
        }

        # Create Orchestra with communication manager
        orchestra = Orchestra(
            agent_registry=AgentRegistry, communication_manager=comm_manager
        )

        # Run with Orchestra and capture detailed execution
        result = await orchestra.execute(
            task=task,
            topology=topology,
            context=context,
            max_steps=80,
        )

        logger.info(f"Deep research completed. Success: {result.success}")

        # Add execution summary to the response file
        with open(agent_response_file, "a") as f:
            f.write("\n" + "=" * 100 + "\n")
            f.write("EXECUTION SUMMARY\n")
            f.write("=" * 100 + "\n\n")
            f.write(f"Total branches executed: {len(result.branch_results)}\n")
            f.write(f"Total steps: {result.total_steps}\n")
            f.write(f"Duration: {result.total_duration:.2f} seconds\n")
            f.write(f"Success: {result.success}\n")

        return {
            "success": result.success,
            "final_response": result.final_response,
            "total_steps": result.total_steps,
            "duration": result.total_duration,
            "parallel_branches": len(
                [br for br in result.branch_results if "child_" in br.branch_id]
            ),
            "metadata": result.metadata,
            "resources_found": result.metadata.get(
                "resources_found", 0
            ),  # Add resources found to result
        }
    except Exception as e:
        logger.error(f"Deep research failed: {e}")
        with open(agent_response_file, "a") as f:
            f.write("\n" + "=" * 100 + "\n")
            f.write("ERROR\n")
            f.write("=" * 100 + "\n\n")
            f.write(f"Deep research failed: {e}\n")
            import traceback

            f.write(f"\nTraceback:\n{traceback.format_exc()}")
        return {"success": False, "error": str(e)}


def save_results(results: Dict[str, Any]):
    """Save the deep research results to files."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    base_filename = f"deep_research_{timestamp}"

    # Ensure output directory exists
    output_dir = Path("examples/real_world/output")
    output_dir.mkdir(parents=True, exist_ok=True)

    # Save the final research report
    if results.get("success") and results.get("final_response"):
        # Extract the actual response content from the JSON structure
        final_response = results["final_response"]

        # If the response is a dict with the agent's structured format, extract the actual content
        if isinstance(final_response, dict):
            if "action_input" in final_response and isinstance(
                final_response["action_input"], dict
            ):
                # Extract from action_input.response
                actual_content = final_response["action_input"].get(
                    "response", str(final_response)
                )
            elif "content" in final_response:
                # Try content field
                actual_content = final_response["content"]
            elif "final_response" in final_response:
                # Try final_response field
                actual_content = final_response["final_response"]
            else:
                # Fallback to string representation
                actual_content = json.dumps(final_response, indent=2)
        else:
            actual_content = str(final_response)

        report_path = output_dir / f"{base_filename}_report.md"
        report_content = f"""# Deep Research Report

**Generated**: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}  
**Session**: Deep Research Session  
**Resources Analyzed**: {results.get('resources_found', 'N/A')}  
**Total Steps**: {results.get('total_steps', 0)}  
**Duration**: {results.get('duration', 0):.2f} seconds  

---

{actual_content}

---

*Report generated by MARS Deep Research System*  
*Timestamp: {timestamp}*
"""
        report_path.write_text(report_content)
        logger.info(f"Research report saved to: {report_path}")
        print(f"\n📄 Markdown report saved: {report_path}")

    # Save execution details
    execution_path = output_dir / f"{base_filename}_execution.json"
    execution_data = {
        "timestamp": timestamp,
        "success": results.get("success", False),
        "total_steps": results.get("total_steps", 0),
        "duration_seconds": results.get("duration", 0),
        "parallel_branches": results.get("parallel_branches", 0),
        "metadata": results.get("metadata", {}),
    }

    with open(execution_path, "w") as f:
        json.dump(execution_data, f, indent=2)
    logger.info(f"Execution details saved to: {execution_path}")


async def main():
    AgentRegistry.clear()
    configs = create_model_configs()
    agents = await create_agents(configs)  # Agents are auto-registered in AgentRegistry

    # Create communication infrastructure
    comm_manager = CommunicationManager()

    # Create and register terminal channel with automated response
    terminal_channel = TerminalChannel("deep_research_terminal")
    # Enable automated mode - no user interaction
    terminal_channel.test_mode = True
    terminal_channel.test_response = "synthetic data robotics academic paper 2025"
    await terminal_channel.start()
    comm_manager.register_channel(terminal_channel)

    # Create session
    session_id = f"deep_research_session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    comm_manager.assign_channel_to_session(session_id, terminal_channel.channel_id)

    # topic = input("Enter the topic for deep research: ")

    # Ask user for resource limit
    # while True:
    #     try:
    #         resource_limit_input = input(
    #             "Enter the maximum number of resources to research (default: 10): "
    #         ).strip()
    #         if resource_limit_input == "":
    #             resource_limit = 10
    #             break
    #         resource_limit = int(resource_limit_input)
    #         if resource_limit > 0:
    #             break
    #         else:
    #             print("Please enter a positive number.")
    #     except ValueError:
    #         print("Please enter a valid number.")

    # print(f"Researching topic: '{topic}' with maximum {resource_limit} resources")
    resource_limit = 5

    try:
        results = await run_deep_research(comm_manager, session_id, resource_limit)

        # Save results to markdown file
        save_results(results)

        # Print brief summary of where the report was saved
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_dir = Path("examples/real_world/output")
        report_path = output_dir / f"deep_research_{timestamp}_report.md"

        print("\n" + "=" * 80)
        print("DEEP RESEARCH COMPLETED")
        print("=" * 80)
        print(f"Success: {results.get('success', False)}")
        print(f"Total Steps: {results.get('total_steps', 0)}")
        print(f"Duration: {results.get('duration', 0):.2f} seconds")
        print(f"Resources Found: {results.get('resources_found', 0)}")

        if results.get("success"):
            print(f"\n✅ Research report saved to: {report_path}")
            print(
                f"✅ Execution details saved to: {output_dir / f'deep_research_{timestamp}_execution.json'}"
            )
        else:
            if results.get("error"):
                print(f"\n❌ Error: {results['error']}")

        print("=" * 80)
        print("\nExiting program...")

        # Exit successfully after saving report
        import sys

        sys.exit(0 if results.get("success") else 1)

    finally:
        # Cleanup
        await terminal_channel.stop()
        await comm_manager.cleanup()


if __name__ == "__main__":
    asyncio.run(main())
