"""
Clean SWE Workflow implementation following LangGraph best practices.
Removes hardcoding and uses tool-based approach with proper prompts.
"""

import operator
import logging
import traceback
from typing import Annotated, Literal, Sequence, TypedDict, Dict, Any
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langgraph.graph import END, START, StateGraph
from langgraph.prebuilt import ToolNode
from pathlib import Path
from tenacity import retry, stop_after_attempt, wait_exponential

from prompts.swe_prompts import SOFTWARE_ENGINEER_PROMPT, CODE_ANALYZER_PROMPT, EDITING_AGENT_PROMPT
from langchain_anthropic import ChatAnthropic
from tools.langraph_tools import LangGraphTools
from tools.advanced_langraph_tools import AdvancedLangGraphTools
from tools.planning_tools import PlanningTools
from agents.planner import PlannerAgent
from state.agent_state import AgentState, create_initial_state

# Optional MCP integration using langchain_mcp_adapters
try:
    from langchain_mcp_adapters.client import MultiServerMCPClient
    from langchain_core.tools import StructuredTool
    import asyncio
    import concurrent.futures
    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False

logger = logging.getLogger(__name__)

class CleanSWEWorkflow:
    """
    Clean SWE Workflow implementation following LangGraph best practices.
    Uses tool-based agents with proper prompts instead of hardcoded logic.
    """

    def __init__(self, repo_path: str, output_dir: str, use_planner: bool = False, enable_mcp: bool = False):
        self.repo_path = Path(repo_path)
        self.output_dir = Path(output_dir)
        self.use_planner = use_planner  # Feature flag for planner logic
        self.enable_mcp = enable_mcp  # Feature flag for MCP integration

        # Initialize Claude client for LangChain
        # The newest Anthropic model is "claude-sonnet-4-20250514", not "claude-3-7-sonnet-20250219", "claude-3-5-sonnet-20241022" nor "claude-3-sonnet-20240229"
        self.anthropic_client = ChatAnthropic(
            model="claude-sonnet-4-20250514",
            temperature=0,
            max_tokens=4096
        )

        # Initialize tools
        self.langraph_tools = LangGraphTools(str(repo_path))
        self.advanced_tools = AdvancedLangGraphTools(str(repo_path))

        # Initialize Planner Agent only if enabled
        if self.use_planner:
            self.planning_tools = PlanningTools()
            self.planner_agent = PlannerAgent(self.planning_tools)
        else:
            self.planning_tools = None
            self.planner_agent = None

        # All agents get access to all tools for maximum flexibility
        all_tools = self.advanced_tools.get_all_tools()

        # Initialize MCP tools if enabled
        mcp_tools = []
        if self.enable_mcp:
            mcp_tools = self._initialize_mcp_integration()

        # Combine built-in tools with optional MCP tools
        self.all_agent_tools = all_tools + mcp_tools

        # Create unified tool node for all agents
        self.unified_tool_node = ToolNode(self.all_agent_tools)

        # Agent names
        self.planner_name = "Planner"
        self.software_engineer_name = "SoftwareEngineer"
        self.code_analyzer_name = "CodeAnalyzer"
        self.editor_name = "Editor"

        # Build the workflow
        self.workflow = self._build_workflow()

        planner_status = "enabled" if self.use_planner else "disabled"
        mcp_status = "enabled" if self.enable_mcp else "disabled"
        total_tools = len(self.all_agent_tools)
        mcp_count = len(mcp_tools) if self.enable_mcp else 0

        print(f"🔧 Clean SWE Workflow initialized:")
        print(f"   📋 Planner: {planner_status}")
        print(f"   🔗 MCP Integration: {mcp_status}")
        print(f"   🛠️  Total Tools: {total_tools} ({len(all_tools)} built-in + {mcp_count} MCP)")

        logger.info(f"🔧 Clean SWE Workflow initialized with Planner {planner_status}, MCP {mcp_status}, and {total_tools} tools")

    def _initialize_mcp_integration(self) -> list:
        """Initialize MCP servers (DeepWiki and Semgrep) safely without breaking existing workflow."""
        if not MCP_AVAILABLE:
            print("⚠️  MCP adapters not available. Install with: pip install langchain-mcp-adapters")
            return []

        try:
            print("🔗 Initializing MCP servers...")

            # Configure MCP servers as per your example
            mcp_servers = {
                "deepwiki": {
                    "url": "https://mcp.deepwiki.com/mcp",
                    "transport": "streamable_http"
                },
                "semgrep": {
                    "url": "https://mcp.semgrep.ai/mcp", 
                    "transport": "streamable_http"
                }
            }

            # Initialize MCP client in a separate async context
            async def get_mcp_tools():
                try:
                    client = MultiServerMCPClient(mcp_servers)
                    async_tools = await client.get_tools()

                    # Convert async MCP tools to sync LangChain tools
                    sync_tools = []
                    for async_tool in async_tools:
                        def create_sync_tool(captured_tool):
                            def sync_wrapper(**kwargs):
                                """Sync wrapper for MCP async tools"""
                                try:
                                    print(f"🔧 Executing MCP tool: {captured_tool.name}")

                                    def run_async():
                                        loop = asyncio.new_event_loop()
                                        asyncio.set_event_loop(loop)
                                        try:
                                            return loop.run_until_complete(captured_tool.ainvoke(kwargs))
                                        finally:
                                            loop.close()

                                    # Run in thread pool to avoid blocking
                                    with concurrent.futures.ThreadPoolExecutor() as executor:
                                        future = executor.submit(run_async)
                                        return future.result(timeout=60)  # 60s timeout

                                except Exception as e:
                                    return f"MCP tool error: {str(e)}"

                            # Create StructuredTool with proper schema
                            return StructuredTool(
                                name=captured_tool.name,
                                description=captured_tool.description or f"MCP tool: {captured_tool.name}",
                                func=sync_wrapper,
                                args_schema=captured_tool.args_schema
                            )

                        sync_tools.append(create_sync_tool(async_tool))

                    return sync_tools

                except Exception as e:
                    print(f"⚠️  MCP server connection failed: {e}")
                    return []

            # Run async initialization
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                mcp_tools = loop.run_until_complete(get_mcp_tools())
                print(f"✅ MCP integration ready: {len(mcp_tools)} tools from {len(mcp_servers)} servers")
                return mcp_tools
            finally:
                loop.close()

        except Exception as e:
            print(f"⚠️  MCP initialization failed: {e}")
            print("🔄 Continuing with built-in tools only")
            return []

    def _create_agent(self, system_prompt: str, tools: list):
        """Create an agent with system prompt and tools."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", system_prompt),
            MessagesPlaceholder(variable_name="messages"),
        ])

        llm = self.anthropic_client
        if tools:
            return prompt | llm.bind_tools(tools)
        else:
            return prompt | llm

    def _create_agent_node(self, agent, name: str):
        """Create an agent node for the workflow following the reference pattern."""
        def agent_node(state):
            @retry(
                stop=stop_after_attempt(3),
                wait=wait_exponential(multiplier=1, min=4, max=10),
            )
            def invoke_with_retry(agent, state):
                return agent.invoke(state)

            try:
                # Track consecutive visits to prevent infinite loops
                consecutive_visits = state.get("consecutive_visits", {})
                visit_count = consecutive_visits.get(name, 0) + 1
                consecutive_visits[name] = visit_count

                # Display agent activity
                print(f"\n💭 {name}: Analyzing task and planning approach...")

                result = invoke_with_retry(agent, state)

                # Display the LLM's thought process
                if hasattr(result, 'content') and result.content:
                    # Handle both string and list content
                    if isinstance(result.content, list):
                        # Extract text from list of content blocks
                        text_parts = []
                        for block in result.content:
                            if isinstance(block, dict) and block.get('type') == 'text':
                                text_parts.append(block.get('text', ''))
                            elif hasattr(block, 'text'):
                                text_parts.append(block.text)
                        thought_content = ' '.join(text_parts)
                    else:
                        thought_content = str(result.content)

                    if thought_content.strip():
                        # Clean up and format the thought
                        thought_content = thought_content.strip()
                        if len(thought_content) > 150:
                            preview = thought_content[:150] + "..."
                            print(f"💭 {name}: {preview}")
                        else:
                            print(f"💭 {name}: {thought_content}")

                # Display tool calls if any
                if hasattr(result, 'tool_calls') and result.tool_calls:
                    for tool_call in result.tool_calls:
                        tool_name = tool_call.get('name', 'unknown_tool')
                        print(f"🔧 {name}: Using tool '{tool_name}'")

            except Exception as e:
                logger.error(f"Failed to invoke {name} agent after 3 attempts: {traceback.format_exc()}")
                result = AIMessage(
                    content=f"I apologize, but I encountered an error and couldn't complete the task. Please try again or rephrase your request.".rstrip(),
                    name=name,
                )
                consecutive_visits = state.get("consecutive_visits", {})
                consecutive_visits[name] = consecutive_visits.get(name, 0) + 1

            # Handle different result types following reference pattern
            if not isinstance(result, ToolMessage):
                if isinstance(result, dict):
                    result_dict = result
                else:
                    result_dict = result.dict() if hasattr(result, 'dict') else {
                        'content': str(result.content) if hasattr(result, 'content') else str(result),
                        'tool_calls': getattr(result, 'tool_calls', [])
                    }

                # Fix: Strip trailing whitespace from content to prevent Anthropic API errors
                if 'content' in result_dict and isinstance(result_dict['content'], str):
                    result_dict['content'] = result_dict['content'].rstrip()

                result = AIMessage(
                    **{
                        k: v
                        for k, v in result_dict.items()
                        if k not in ["type", "name"]
                    },
                    name=name,
                )

            return {
                "messages": [result], 
                "sender": name,
                "consecutive_visits": consecutive_visits
            }

        return agent_node

    def _build_workflow(self) -> StateGraph:
        """Build the clean workflow graph with Planner approval."""
        # Create agents - all agents get access to all tools
        software_engineer_agent = self._create_agent(SOFTWARE_ENGINEER_PROMPT, self.all_agent_tools)
        code_analyzer_agent = self._create_agent(CODE_ANALYZER_PROMPT, self.all_agent_tools)
        editing_agent = self._create_agent(EDITING_AGENT_PROMPT, self.all_agent_tools)

        # Create agent nodes
        planner_node = self._create_planner_node()
        software_engineer_node = self._create_agent_node(software_engineer_agent, self.software_engineer_name)
        code_analyzer_node = self._create_agent_node(code_analyzer_agent, self.code_analyzer_name)
        editing_node = self._create_agent_node(editing_agent, self.editor_name)

        # Create workflow
        workflow = StateGraph(AgentState)

        # Add nodes
        if self.use_planner:
            workflow.add_node(self.planner_name, planner_node)
        workflow.add_node(self.software_engineer_name, software_engineer_node)
        workflow.add_node(self.code_analyzer_name, code_analyzer_node)
        workflow.add_node(self.editor_name, editing_node)
        workflow.add_node("tools", self.unified_tool_node)

        # Set entry point based on feature flag
        if self.use_planner:
            workflow.add_edge(START, self.planner_name)
        else:
            workflow.add_edge(START, self.software_engineer_name)

        # Add unified tool routing - all agents can use all tools
        workflow.add_conditional_edges(
            "tools",
            lambda x: x["sender"],
            {
                self.software_engineer_name: self.software_engineer_name,
                self.code_analyzer_name: self.code_analyzer_name,
                self.editor_name: self.editor_name,
            },
        )

        # Add Planner routing only if enabled
        if self.use_planner:
            workflow.add_conditional_edges(
                self.planner_name,
                self._planner_router,
                {
                    "approved": self.software_engineer_name,
                    "rejected": END,
                    "continue": self.planner_name,
                },
            )

        # Add main routing logic
        workflow.add_conditional_edges(
            self.software_engineer_name,
            self._software_engineer_router,
            {
                "continue": self.software_engineer_name,
                "analyze_code": self.code_analyzer_name,
                "edit_file": self.editor_name,
                "tools": "tools",
                "__end__": END,
            },
        )

        workflow.add_conditional_edges(
            self.code_analyzer_name,
            self._code_analyzer_router,
            {
                "continue": self.code_analyzer_name,
                "done": self.software_engineer_name,
                "edit_file": self.editor_name,
                "tools": "tools",
            },
        )

        workflow.add_conditional_edges(
            self.editor_name,
            self._editor_router,
            {
                "continue": self.editor_name,
                "done": self.software_engineer_name,
                "tools": "tools",
            },
        )

        return workflow

    def _create_planner_node(self):
        """Create the Planner agent node."""
        def planner_node(state):
            try:
                # Track consecutive visits to prevent infinite loops
                consecutive_visits = state.get("consecutive_visits", {})
                visit_count = consecutive_visits.get(self.planner_name, 0) + 1
                consecutive_visits[self.planner_name] = visit_count

                # Use the planner agent to evaluate and approve tasks
                messages = state.get("messages", [])
                if messages:
                    # Let the planner agent assess the task and create a plan
                    result = self.planner_agent.process(state) if self.planner_agent else {"messages": [], "sender": self.planner_name}

                    # Check if planner has created an approvable plan
                    last_message = result.get("messages", [])[-1] if result.get("messages") else None
                    if last_message and any(signal in last_message.content.lower() for signal in ["plan approved", "proceeding with implementation", "implementation ready"]):
                        # Auto-approve to move to software engineer
                        approval_msg = AIMessage(
                            content="Plan approved. Proceeding with implementation.".rstrip(),
                            name=self.planner_name
                        )
                        return {
                            "messages": result.get("messages", []) + [approval_msg],
                            "sender": self.planner_name,
                            "consecutive_visits": consecutive_visits
                        }

                    # Add consecutive visits to result and return
                    result["consecutive_visits"] = consecutive_visits
                    return result
            except Exception as e:
                logger.error(f"Error in Planner agent: {e}")
                error_msg = AIMessage(
                    content=f"Error in Planner: {str(e)}".rstrip(),
                    name=self.planner_name
                )
                consecutive_visits = state.get("consecutive_visits", {})
                return {
                    "messages": [error_msg], 
                    "sender": self.planner_name,
                    "consecutive_visits": consecutive_visits
                }

        return planner_node

    def _planner_router(self, state) -> Literal["approved", "rejected", "continue"]:
        """Route from Planner based on approval status."""
        messages = state["messages"]

        # Find the last AI message
        for message in reversed(messages):
            if isinstance(message, AIMessage):
                last_message = message
                break
        else:
            last_message = messages[-1]

        content = str(last_message.content).lower() if hasattr(last_message, 'content') else ""

        # Check for approval signals
        if any(signal in content for signal in ["approved", "plan approved", "proceeding with implementation"]):
            return "approved"
        elif any(signal in content for signal in ["rejected", "plan rejected", "need revision"]):
            return "rejected"
        elif any(signal in content for signal in ["plan complete", "implementation ready"]):
            return "approved"
        else:
            return "continue"

    def _software_engineer_router(self, state) -> Literal["continue", "analyze_code", "edit_file", "tools", "__end__"]:
        """Route from Software Engineer based on message content following reference pattern."""
        messages = state["messages"]

        # Find the last AI message
        for message in reversed(messages):
            if isinstance(message, AIMessage):
                last_ai_message = message
                break
        else:
            last_ai_message = messages[-1]

        # Check for tool calls first
        if hasattr(last_ai_message, 'tool_calls') and last_ai_message.tool_calls:
            return "tools"

        # Check for specific routing keywords (following reference pattern)
        content = str(last_ai_message.content).upper() if hasattr(last_ai_message, 'content') else ""

        # Enhanced completion detection
        completion_keywords = [
            "TASK COMPLETED", "PATCH COMPLETED", "IMPLEMENTATION COMPLETED",
            "FILE CREATED SUCCESSFULLY", "CODE GENERATED SUCCESSFULLY", 
            "CAESAR CIPHER IMPLEMENTATION COMPLETE", "COMPLETED",
            "DONE", "FINISHED", "SUCCESS"
        ]

        if any(keyword in content for keyword in completion_keywords):
            return "__end__"

        if "ANALYZE CODE" in content:
            return "analyze_code"
        if "EDIT FILE" in content:
            return "edit_file"

        # Check consecutive visits to prevent infinite loops - REDUCED THRESHOLD 
        consecutive_visits = state.get("consecutive_visits", {})
        swe_visits = consecutive_visits.get(self.software_engineer_name, 0)

        # More aggressive loop prevention - reduced from 5 to 3
        if swe_visits >= 3:
            logger.warning(f"🚨 Software Engineer visited {swe_visits} times consecutively. Terminating to prevent loop.")
            return "__end__"

        # Check if file was recently created - if so, consider task complete
        recent_tool_messages = [msg for msg in reversed(messages[-10:]) if isinstance(msg, ToolMessage)]
        if recent_tool_messages:
            for tool_msg in recent_tool_messages:
                if any(success_indicator in str(tool_msg.content).lower() for success_indicator in 
                       ["created successfully", "file created", "written to", "saved"]):
                    return "__end__"

        return "continue"

    def _code_analyzer_router(self, state) -> Literal["continue", "done", "edit_file", "tools"]:
        """Route from Code Analyzer based on message content following reference pattern."""
        messages = state["messages"]

        # Find the last AI message
        for message in reversed(messages):
            if isinstance(message, AIMessage):
                last_ai_message = message
                break
        else:
            last_ai_message = messages[-1]

        # Check for tool calls first
        if hasattr(last_ai_message, 'tool_calls') and last_ai_message.tool_calls:
            return "tools"

        # Check for specific routing keywords (following reference pattern)
        content = str(last_ai_message.content).upper() if hasattr(last_ai_message, 'content') else ""

        completion_keywords = ["ANALYSIS COMPLETE", "ANALYSIS DONE", "COMPLETE", "FINISHED"]
        if any(keyword in content for keyword in completion_keywords):
            return "done"
        if "EDIT FILE" in content:
            return "edit_file"

        # Check consecutive visits to prevent infinite loops - REDUCED THRESHOLD
        consecutive_visits = state.get("consecutive_visits", {})
        analyzer_visits = consecutive_visits.get(self.code_analyzer_name, 0)

        if analyzer_visits >= 2:  # More aggressive loop prevention - reduced from 5 to 2
            logger.warning(f"🚨 Code Analyzer visited {analyzer_visits} times consecutively. Returning to Software Engineer.")
            return "done"

        return "continue"

    def _editor_router(self, state) -> Literal["continue", "done", "tools"]:
        """Route from Editor based on message content following reference pattern."""
        messages = state["messages"]

        # Find the last AI message
        for message in reversed(messages):
            if isinstance(message, AIMessage):
                last_ai_message = message
                break
        else:
            last_ai_message = messages[-1]

        # Check for tool calls first
        if hasattr(last_ai_message, 'tool_calls') and last_ai_message.tool_calls:
            return "tools"

        # Check for specific routing keywords (following reference pattern)
        content = str(last_ai_message.content).upper() if hasattr(last_ai_message, 'content') else ""

        completion_keywords = ["EDITING COMPLETED", "EDIT COMPLETE", "DONE", "FINISHED", "COMPLETE"]
        if any(keyword in content for keyword in completion_keywords):
            return "done"

        # Check consecutive visits to prevent infinite loops - REDUCED THRESHOLD
        consecutive_visits = state.get("consecutive_visits", {})
        editor_visits = consecutive_visits.get(self.editor_name, 0)

        if editor_visits >= 2:  # More aggressive loop prevention - reduced from 5 to 2
            logger.warning(f"🚨 Editor visited {editor_visits} times consecutively. Returning to Software Engineer.")
            return "done"

        return "continue"

    def run_workflow(self, task_description: str) -> Dict[str, Any]:
        """Run the clean workflow."""
        logger.info(f"🚀 Starting clean SWE workflow with task: {task_description}")

        # Create initial state
        initial_state = create_initial_state(task_description)
        initial_state["messages"] = [HumanMessage(content=task_description)]

        # Compile and run workflow
        try:
            app = self.workflow.compile()
            config = {"recursion_limit": 50}  # Reduced limit to prevent infinite loops
            final_state = app.invoke(initial_state, config=config)

            return {
                "task_description": task_description,
                "final_state": final_state,
                "message_count": len(final_state["messages"]),
                "final_sender": final_state["sender"],
                "success": True
            }

        except Exception as e:
            logger.error(f"Clean workflow execution failed: {e}")
            return {
                "task_description": task_description,
                "error": str(e),
                "success": False
            }