"""Deep Research workflow implementation adapted for the GSM evaluation system."""

import asyncio
import logging
from typing import Optional, List, Dict, Any

from langchain.chat_models import init_chat_model
from langchain_core.messages import (
    AIMessage, 
    HumanMessage, 
    SystemMessage, 
    ToolMessage, 
    filter_messages,
    get_buffer_string,
)
from langchain_core.runnables import RunnableConfig
from langgraph.graph import END, START, StateGraph
from langgraph.types import Command

from .deep_research_config import DeepResearchConfig
from .deep_research_state import (
    AgentInputState,
    AgentState, 
    SupervisorState,
    ResearcherState,
    ResearcherOutputState,
    ConductResearch,
    ResearchComplete,
    ClarifyWithUser,
    ResearchQuestion,
    BacktrackResearch,
)
from .deep_research_prompts import (
    clarify_with_user_instructions,
    transform_messages_into_research_topic_prompt,
    lead_researcher_prompt,
    research_system_prompt,
    backtracking_research_system_prompt,
    compress_research_system_prompt,
    compress_research_simple_human_message,
    final_report_generation_prompt,
)
from .deep_research_utils import (
    think_tool,
    get_today_str,
    get_api_key_for_model,
    get_notes_from_tool_calls,
)
from ..search import SearchEngine


class DeepResearchWorkflow:
    """Deep Research workflow implementation for GSM agent evaluation."""
    
    def __init__(
        self,
        search_engine: SearchEngine,
        config: DeepResearchConfig,
        search_tools: List[Any]|List[List[Any]],
        search_checker: Any,
        search_engines: Optional[List[SearchEngine]] = None,
        logger: Optional[logging.Logger] = None
    ):
        """Initialize the Deep Research workflow.
        
        Args:
            search_engine: SearchEngine instance for research (legacy, may not be used)
            config: Deep research configuration
            search_tools: List of search tools OR List of search tool groups for multi-agent scenarios
            search_checker: Search checker instance
            search_engines: List of search engines for multi-agent scenarios (for cache clearing)
            logger: Logger instance
        """
        self.search_engine = search_engine
        self.search_engines = search_engines or [search_engine] if search_engine else []
        self.config = config
        self.search_checker = search_checker
        self.logger = logger or logging.getLogger(__name__)
        

        # Handle both single and multi-agent search tools
        if not search_tools:
            raise ValueError("search_tools cannot be empty")
        
        self.search_tool_groups = search_tools
        
        # Initialize configurable model
        self.configurable_model = init_chat_model(
            configurable_fields=("model", "model_provider", "max_tokens", "api_key"),
        )
        
        # Build workflow graphs
        self._build_workflow()
    
    def _build_workflow(self):
        """Build the complete deep research workflow."""
        # Build researcher subgraph first
        self._build_researcher_subgraph()
        
        # Build supervisor subgraph
        self._build_supervisor_subgraph()
        
        # Build main workflow
        self._build_main_workflow()
    
    def _build_researcher_subgraph(self):
        """Build the researcher subgraph for focused research tasks."""
        researcher_builder = StateGraph(
            ResearcherState,
            output=ResearcherOutputState,
            config_schema=DeepResearchConfig
        )
        
        # Add researcher nodes
        researcher_builder.add_node("researcher", self._researcher_node)
        researcher_builder.add_node("researcher_tools", self._researcher_tools_node) 
        researcher_builder.add_node("compress_research", self._compress_research_node)
        
        # Add researcher edges
        researcher_builder.add_edge(START, "researcher")
        researcher_builder.add_edge("compress_research", END)
        
        self.researcher_subgraph = researcher_builder.compile()
    
    def _build_supervisor_subgraph(self):
        """Build the supervisor subgraph for research management."""
        supervisor_builder = StateGraph(SupervisorState, config_schema=DeepResearchConfig)
        
        # Add supervisor nodes
        supervisor_builder.add_node("supervisor", self._supervisor_node)
        supervisor_builder.add_node("supervisor_tools", self._supervisor_tools_node)
        
        # Add supervisor edges
        supervisor_builder.add_edge(START, "supervisor")
        
        self.supervisor_subgraph = supervisor_builder.compile()
    
    def _build_main_workflow(self):
        """Build the main deep research workflow."""
        main_builder = StateGraph(
            AgentState,
            input=AgentInputState,
            config_schema=DeepResearchConfig
        )
        
        # Add main workflow nodes
        main_builder.add_node("clarify_with_user", self._clarify_with_user_node)
        main_builder.add_node("write_research_brief", self._write_research_brief_node)
        main_builder.add_node("research_supervisor", self.supervisor_subgraph)
        main_builder.add_node("final_report_generation", self._final_report_generation_node)
        
        # Add main workflow edges
        main_builder.add_edge(START, "clarify_with_user")
        main_builder.add_edge("research_supervisor", "final_report_generation") 
        main_builder.add_edge("final_report_generation", END)
        
        self.workflow = main_builder.compile()
    
    async def _clarify_with_user_node(self, state: AgentState, config: RunnableConfig) -> Command:
        """Node for user clarification (usually skipped for math problems)."""
        # For math evaluation, usually skip clarification
        if not self.config.allow_clarification:
            return Command(goto="write_research_brief")
        
        # Prepare clarification model
        model_config = {
            "model": self.config.research_model,
            "model_provider": self.config.research_model_provider,
            "max_tokens": self.config.research_model_max_tokens,
            "temperature": self.config.research_model_temperature,
            "api_key": get_api_key_for_model(self.config.research_model, config),
            "tags": ["langsmith:nostream"]
        }
        clarification_model = (
            self.configurable_model
            .with_structured_output(ClarifyWithUser)
            .with_retry(stop_after_attempt=self.config.max_structured_output_retries)
            .with_config(model_config)
        )
        
        # Analyze whether clarification is needed
        messages = state["messages"]
        prompt_content = clarify_with_user_instructions.format(
            messages=get_buffer_string(messages),
            date=get_today_str()
        )
        response = await clarification_model.ainvoke([HumanMessage(content=prompt_content)])
        
        if response.need_clarification:
            return Command(
                goto=END,
                update={"messages": [AIMessage(content=response.question)]}
            )
        else:
            return Command(
                goto="write_research_brief",
                update={"messages": [AIMessage(content=response.verification)]}
            )
    
    async def _write_research_brief_node(self, state: AgentState, config: RunnableConfig) -> Command:
        """Node for generating research brief from user messages."""
        # Set up research model
        model_config = {
            "model": self.config.research_model,
            "model_provider": self.config.research_model_provider,
            "max_tokens": self.config.research_model_max_tokens,
            "temperature": self.config.research_model_temperature,
            "api_key": get_api_key_for_model(self.config.research_model, config),
            "tags": ["langsmith:nostream"]
        }
        research_model = (
            self.configurable_model
            .with_structured_output(ResearchQuestion)
            .with_retry(stop_after_attempt=self.config.max_structured_output_retries)
            .with_config(model_config)
        )
        
        # Generate research brief
        prompt_content = transform_messages_into_research_topic_prompt.format(
            messages=get_buffer_string(state["messages"]),
            date=get_today_str()
        )
        response = await research_model.ainvoke([HumanMessage(content=prompt_content)])
        
        # Initialize supervisor
        supervisor_system_prompt = lead_researcher_prompt.format(
            date=get_today_str(),
            max_concurrent_research_units=self.config.max_concurrent_research_units,
            max_researcher_iterations=self.config.max_researcher_iterations
        )
        
        return Command(
            goto="research_supervisor",
            update={
                "research_brief": response.research_brief,
                "supervisor_messages": {
                    "type": "override",
                    "value": [
                        SystemMessage(content=supervisor_system_prompt),
                        HumanMessage(content=response.research_brief)
                    ]
                }
            }
        )
    
    async def _supervisor_node(self, state: SupervisorState, config: RunnableConfig) -> Command:
        """Supervisor node for managing research tasks."""
        # Configure supervisor model
        model_config = {
            "model": self.config.research_model,
            "model_provider": self.config.research_model_provider,
            "max_tokens": self.config.research_model_max_tokens,
            "temperature": self.config.research_model_temperature,
            "api_key": get_api_key_for_model(self.config.research_model, config),
            "tags": ["langsmith:nostream"]
        }
        # Available tools for supervisor
        supervisor_tools = [ConductResearch, BacktrackResearch, ResearchComplete, think_tool]
        
        research_model = (
            self.configurable_model
            .bind_tools(supervisor_tools)
            .with_retry(stop_after_attempt=self.config.max_structured_output_retries)
            .with_config(model_config)
        )
        
        # Generate supervisor response
        supervisor_messages = state["supervisor_messages"]
        response = await research_model.ainvoke(supervisor_messages)
        
        return Command(
            goto="supervisor_tools",
            update={
                "supervisor_messages": [response],
                "research_iterations": state.get("research_iterations", 0) + 1
            }
        )
    
    async def _supervisor_tools_node(self, state: SupervisorState, config: RunnableConfig) -> Command:
        """Execute supervisor tools including research delegation."""
        supervisor_messages = state["supervisor_messages"]
        research_iterations = state.get("research_iterations", 0)
        most_recent_message = supervisor_messages[-1]
        
        # Check exit conditions
        exceeded_iterations = research_iterations > self.config.max_researcher_iterations
        no_tool_calls = not most_recent_message.tool_calls
        research_complete_called = any(
            tool_call["name"] == "ResearchComplete"
            for tool_call in most_recent_message.tool_calls
        )
        
        if exceeded_iterations or no_tool_calls or research_complete_called:
            return Command(
                goto=END,
                update={
                    "notes": get_notes_from_tool_calls(supervisor_messages),
                    "research_brief": state["research_brief"],
                    "research_topics_by_round": state.get("research_topics_by_round", [])
                }
            )
        
        # Process tool calls - MUST respond to every tool call to avoid OpenAI API errors
        all_tool_messages = []
        update_payload = {"supervisor_messages": []}
        
        # Track which tool calls we've handled
        handled_tool_call_ids = set()
        
        # Handle think_tool calls
        think_tool_calls = [
            tc for tc in most_recent_message.tool_calls
            if tc["name"] == "think_tool"
        ]
        
        for tool_call in think_tool_calls:
            reflection = tool_call["args"]["reflection"]
            all_tool_messages.append(ToolMessage(
                content=f"Reflection recorded: {reflection}",
                name="think_tool",
                tool_call_id=tool_call["id"]
            ))
            handled_tool_call_ids.add(tool_call["id"])
        
        # Handle ConductResearch and BacktrackResearch calls (combined concurrency control)
        conduct_research_calls = [tc for tc in most_recent_message.tool_calls if tc["name"] == "ConductResearch"]
        backtrack_calls = [tc for tc in most_recent_message.tool_calls if tc["name"] == "BacktrackResearch"]
        # import pdb; pdb.set_trace()

        if conduct_research_calls or backtrack_calls:
            try:
                # Start new research round
                self.search_checker.start_new_research_round()
                
                # Clear search engine caches at the beginning of research round
                from ..search import SearchTools
                SearchTools.clear_search_engine_caches(self.search_engines)
                
                # Combine and limit concurrent research across both call types
                combined_calls = conduct_research_calls + backtrack_calls
                allowed_calls = combined_calls[:self.config.max_concurrent_research_units]
                overflow_calls = combined_calls[self.config.max_concurrent_research_units:]
                
                # Record topics for this round
                topics_this_round = []
                for tc in allowed_calls:
                    try:
                        if tc["name"] == "ConductResearch":
                            topics_this_round.append(tc["args"].get("research_topic", ""))
                        elif tc["name"] == "BacktrackResearch":
                            target = tc["args"].get("target_topic", "")
                            topics_this_round.append(f"BACKTRACK: {target}")
                    except Exception:
                        continue
                update_payload.setdefault("research_topics_by_round", [])
                update_payload["research_topics_by_round"] = [
                    {"round": research_iterations, "topics": topics_this_round}
                ]


                # Execute research tasks with distributed search tools
                research_tasks = []
                for idx, tool_call in enumerate(allowed_calls):
                    if tool_call["name"] == "ConductResearch":
                        research_tasks.append(
                            self._execute_research_task(
                                tool_call["args"]["research_topic"],
                                config,
                                task_index=idx,
                                supervisor_messages=state["supervisor_messages"]
                            )
                        )
                    elif tool_call["name"] == "BacktrackResearch":
                        back_ctx = {
                            "target_topic": tool_call["args"].get("target_topic", ""),
                            "reason": tool_call["args"].get("reason", ""),
                            "previous_related_information": tool_call["args"].get("previous_related_information", ""),
                            "updated_focus": tool_call["args"].get("updated_focus", ""),
                        }
                        print(f"back_ctx: {back_ctx}")
                        research_tasks.append(
                            self._execute_research_task(
                                back_ctx["target_topic"] or "Refine previous topic",
                                config,
                                task_index=idx,
                                supervisor_messages=state["supervisor_messages"],
                                backtracking_context=back_ctx,
                            )
                        )
                
                tool_results = await asyncio.gather(*research_tasks)
                
                # Create tool messages
                for observation, tool_call in zip(tool_results, allowed_calls):
                    all_tool_messages.append(ToolMessage(
                        content=observation["compressed_research"],
                        name=tool_call["name"],
                        tool_call_id=tool_call["id"]
                    ))
                    handled_tool_call_ids.add(tool_call["id"])
                
                # Handle overflow calls
                for overflow_call in overflow_calls:
                    all_tool_messages.append(ToolMessage(
                        content=f"Error: Maximum concurrent research units ({self.config.max_concurrent_research_units}) exceeded.",
                        name=overflow_call.get("name", "ConductResearch"),
                        tool_call_id=overflow_call["id"]
                    ))
                    handled_tool_call_ids.add(overflow_call["id"])
                
                # Aggregate raw notes
                raw_notes = "\n".join([
                    "\n".join(observation["raw_notes"])
                    for observation in tool_results
                ])
                
                if raw_notes:
                    update_payload["raw_notes"] = [raw_notes]

                    
            except Exception as e:
                self.logger.error(f"Error in research execution: {e}")
                # No fallback - let error propagate for debugging
                raise
        
        # Handle any remaining unhandled tool calls to avoid OpenAI API errors
        for tool_call in most_recent_message.tool_calls:
            if tool_call["id"] not in handled_tool_call_ids:
                self.logger.warning(f"Unhandled tool call: {tool_call['name']} (id: {tool_call['id']})")
                raise ValueError(f"Unhandled tool call: {tool_call['name']} (id: {tool_call['id']})")
        
        update_payload["supervisor_messages"] = all_tool_messages
        return Command(goto="supervisor", update=update_payload)
    
    async def _execute_research_task(self, research_topic: str, config: RunnableConfig, task_index: int = 0, supervisor_messages: list = None, backtracking_context: Dict[str, Any] | None = None) -> Dict[str, Any]:
        """Execute a single research task using the researcher subgraph."""
        # Select the appropriate search tool group for this task
        tool_group_index = task_index % len(self.search_tool_groups)
        assigned_tools = self.search_tool_groups[tool_group_index]
        self.logger.info(f"Research task {task_index} assigned to search tool group {tool_group_index}")

        
        # Create a modified config that includes the assigned tools
        if config is None:
            raise ValueError("Config cannot be None for research task execution")
        
        # Use the configurable field of RunnableConfig to pass custom data
        task_config = config.copy() if hasattr(config, 'copy') else dict(config) if config else {}
        if 'configurable' not in task_config:
            task_config['configurable'] = {}
        task_config['configurable']['assigned_search_tools'] = assigned_tools

        self.logger.info(f"Task index: {task_index}")
        
        # No fallback - let errors propagate for debugging
        result = await self.researcher_subgraph.ainvoke({
            "researcher_messages": [HumanMessage(content=research_topic)],
            "research_topic": research_topic,
            "tool_call_iterations": 0,
            "compressed_research": "",
            "raw_notes": [],
            "supervisor_messages": supervisor_messages or [],
            "backtracking_context": backtracking_context,
        }, task_config)
        return result
    
    async def _researcher_node(self, state: ResearcherState, config: RunnableConfig) -> Command:
        """Individual researcher node for focused research."""
        researcher_messages = state["researcher_messages"]
        
        # Use assigned search tools - should be in config["configurable"]["assigned_search_tools"]
        assigned_search_tools = config["configurable"]["assigned_search_tools"]
        tools = [*assigned_search_tools, think_tool]


        # Configure researcher model
        model_config = {
            "model": self.config.research_model,
            "model_provider": self.config.research_model_provider,
            "max_tokens": self.config.research_model_max_tokens,
            "temperature": self.config.research_model_temperature,
            "api_key": get_api_key_for_model(self.config.research_model, config),
            "tags": ["langsmith:nostream"]
        }
        
        # Prepare system prompt (supports backtracking mode)
        back_ctx = state.get("backtracking_context")
        if back_ctx:
            print(f"back_ctx: {back_ctx}")
            researcher_prompt = backtracking_research_system_prompt.format(
                date=get_today_str(),
                max_researcher_iterations=self.config.max_researcher_iterations,
                target_topic=back_ctx.get("target_topic", state.get("research_topic", "")),
                reason=back_ctx.get("reason", ""),
                previous_related_information=back_ctx.get("previous_related_information", ""),
                updated_focus=back_ctx.get("updated_focus", "None provided"),
            )
        else:
            researcher_prompt = research_system_prompt.format(
                mcp_prompt="",  # No MCP for this integration
                date=get_today_str(),
                max_researcher_iterations=self.config.max_researcher_iterations
            )
        
        research_model = (
            self.configurable_model
            .bind_tools(tools)
            .with_retry(stop_after_attempt=self.config.max_structured_output_retries)
            .with_config(model_config)
        )
        
        # Generate researcher response
        messages = [SystemMessage(content=researcher_prompt)] + researcher_messages
        response = await research_model.ainvoke(messages)
        
        return Command(
            goto="researcher_tools",
            update={
                "researcher_messages": [response],
                "tool_call_iterations": state.get("tool_call_iterations", 0) + 1
            }
        )
    
    async def _researcher_tools_node(self, state: ResearcherState, config: RunnableConfig) -> Command:
        """Execute researcher tools including search operations."""
        researcher_messages = state["researcher_messages"]
        most_recent_message = researcher_messages[-1]
        
        # Check if any tool calls were made
        if not most_recent_message.tool_calls:
            return Command(goto="compress_research")
        
        # Use assigned search tools - should be in config["configurable"]["assigned_search_tools"]
        assigned_search_tools = config["configurable"]["assigned_search_tools"]
        tools = [*assigned_search_tools, think_tool]
        tools_by_name = {
            tool.name if hasattr(tool, "name") else tool.__name__ if hasattr(tool, "__name__") else None: tool
            for tool in tools
        }
        
        # Execute tool calls
        tool_outputs = []
        for tool_call in most_recent_message.tool_calls:
            tool_name = tool_call["name"]
            tool_call_id = tool_call["id"]
            
            if tool_name in tools_by_name:
                try:
                    result = await tools_by_name[tool_name].ainvoke(tool_call["args"], config)
                    tool_outputs.append(ToolMessage(
                        content=str(result),
                        name=tool_name,
                        tool_call_id=tool_call_id
                    ))
                except Exception as e:
                    self.logger.error(f"Tool execution failed for {tool_name}: {e}")
                    raise ValueError(f"Tool execution failed for {tool_name}: {e}")
            else:
                # Tool not found - there's error in the tool call
                raise ValueError(f"Tool '{tool_name}' is not available. Available tools: {', '.join(tools_by_name.keys())}")
        
        # Check exit conditions
        exceeded_iterations = state.get("tool_call_iterations", 0) >= self.config.max_react_tool_calls
        
        if exceeded_iterations:
            return Command(
                goto="compress_research",
                update={"researcher_messages": tool_outputs}
            )
        
        return Command(
            goto="researcher",
            update={"researcher_messages": tool_outputs}
        )
    
    async def _compress_research_node(self, state: ResearcherState, config: RunnableConfig) -> Dict[str, Any]:
        """Compress research findings into structured summary."""
        # Configure compression model
        # We use gpt-4o for compression because it is the most reliable model for this task
        model_config = {
            "model": "gpt-4o",
            "model_provider": "openai",
            "max_tokens": self.config.compression_model_max_tokens,
            "temperature": self.config.compression_model_temperature,
            "api_key": get_api_key_for_model(self.config.compression_model, config),
            "tags": ["langsmith:nostream"]
        }
        
        compression_model = self.configurable_model.with_config(model_config)
        
        # Prepare messages for compression
        researcher_messages = state["researcher_messages"]
        researcher_messages.append(HumanMessage(content=compress_research_simple_human_message))
        
        # Get original question for context from supervisor messages
        supervisor_messages = state["supervisor_messages"]
        original_question = supervisor_messages[1].content if len(supervisor_messages) > 1 else "No original question available"
        
        # No retry logic or fallback - let errors propagate for debugging
        compression_prompt = compress_research_system_prompt.format(
            original_question=original_question,
            date=get_today_str()
        )
        messages = [SystemMessage(content=compression_prompt)] + researcher_messages
        
        response = await compression_model.ainvoke(messages)
        
        # Extract raw notes
        raw_notes_content = "\n".join([
            str(msg.content)
            for msg in filter_messages(researcher_messages, include_types=["tool", "ai"])
        ])

        # Handle empty response (e.g., from reasoning models that hit token limits)
        response_content = str(response.content).strip()
        if not response_content:
            # Fallback: use raw notes as compressed research
            response_content = f"## Research Findings Summary\n\n{raw_notes_content}"
        
        return {
            "compressed_research": response_content,
            "raw_notes": [raw_notes_content]
        }
    
    async def _final_report_generation_node(self, state: AgentState, config: RunnableConfig) -> Dict[str, Any]:
        """Generate the final comprehensive report."""
        notes = state["notes"]
        findings = "\n".join(notes)
        
        # Configure report generation model
        model_config = {
            "model": self.config.final_report_model,
            "model_provider": self.config.final_report_model_provider,
            "max_tokens": self.config.final_report_model_max_tokens,
            "temperature": self.config.final_report_model_temperature,
            "api_key": get_api_key_for_model(self.config.final_report_model, config),
            "tags": ["langsmith:nostream"]
        }
        
        prompt = final_report_generation_prompt.format(
            research_brief=state["research_brief"],
            messages=get_buffer_string(state["messages"]),
            findings=findings,
            date=get_today_str()
        )
        
        final_report = await self.configurable_model.with_config(model_config).ainvoke([
            HumanMessage(content=prompt)
        ])
        
        return {
            "final_report": final_report.content,
            "final_report_prompt": prompt,
            "messages": [final_report],
            "notes": {"type": "override", "value": []},
            "research_topics_by_round": state.get("research_topics_by_round", [])
        }
    
    async def run(self, problem: str, **kwargs) -> Dict[str, Any]:
        """Run the complete deep research workflow.
        
        Args:
            problem: The problem/question to research
            **kwargs: Additional configuration options
            
        Returns:
            Final workflow result with report
        """
        # No try/catch - let errors propagate for debugging
        # Prepare input
        input_data = {
            "messages": [HumanMessage(content=problem)]
        }
        
        # Create runtime config
        config = {
            "configurable": {
                "research_model": self.config.research_model,
                "compression_model": self.config.compression_model,
                "final_report_model": self.config.final_report_model,
            }
        }
        
        # Execute workflow
        result = await self.workflow.ainvoke(input_data, config)
        return result


if __name__ == "__main__":
    """Display the deep research workflow graph structure."""
    import sys
    from unittest.mock import Mock
    
    # Create mock dependencies for demonstration
    mock_search_engine = Mock()
    mock_config = Mock()
    mock_config.allow_clarification = False
    mock_config.research_model = "gpt-4"
    mock_config.research_model_max_tokens = 4000
    mock_config.compression_model = "gpt-4"
    mock_config.compression_model_max_tokens = 4000
    mock_config.final_report_model = "gpt-4"
    mock_config.final_report_model_max_tokens = 4000
    mock_config.max_concurrent_research_units = 3
    mock_config.max_researcher_iterations = 10
    mock_config.max_react_tool_calls = 15
    mock_config.max_structured_output_retries = 3
    mock_search_tools = []
    mock_search_checker = Mock()
    
    try:
        # Create workflow instance
        workflow = DeepResearchWorkflow(
            search_engine=mock_search_engine,
            config=mock_config,
            search_tools=mock_search_tools,
            search_checker=mock_search_checker
        )
        
        print("=== Deep Research Workflow Graph Structure ===\n")
        
        # Display main workflow graph
        print("Main Workflow:")
        main_graph = workflow.workflow.get_graph()
        print(f"Nodes: {list(main_graph.nodes.keys())}")
        print(f"Edges: {[(edge.source, edge.target) for edge in main_graph.edges]}")
        
        print("\nSupervisor Subgraph:")
        supervisor_graph = workflow.supervisor_subgraph.get_graph()
        print(f"Nodes: {list(supervisor_graph.nodes.keys())}")
        print(f"Edges: {[(edge.source, edge.target) for edge in supervisor_graph.edges]}")
        
        print("\nResearcher Subgraph:")
        researcher_graph = workflow.researcher_subgraph.get_graph()
        print(f"Nodes: {list(researcher_graph.nodes.keys())}")
        print(f"Edges: {[(edge.source, edge.target) for edge in researcher_graph.edges]}")
        
        # Try to generate a visual representation if mermaid is available
        try:
            print("\n=== Main Workflow Mermaid Diagram ===")
            print(main_graph.draw_mermaid())
        except Exception as e:
            print(f"\nMermaid visualization not available: {e}")
            
    except Exception as e:
        print(f"Error creating workflow: {e}")
        sys.exit(1)
