#!/usr/bin/env python3
"""
Debug script to test MCP tool parameter passing in LangGraph
"""

import asyncio
from workflows.clean_swe_workflow import CleanSWEWorkflow
from pathlib import Path
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_mcp_tool_direct():
    """Test MCP tool directly"""
    print("=== DIRECT MCP TOOL TEST ===")
    
    workflow = CleanSWEWorkflow(Path.cwd(), Path.cwd() / 'output', False)
    
    # Find security_check tool
    security_tool = None
    for tool in workflow.mcp_tools:
        if tool.name == 'security_check':
            security_tool = tool
            break
    
    if not security_tool:
        print("❌ Security tool not found!")
        return
    
    print(f"✅ Found security tool: {security_tool.name}")
    print(f"Tool type: {type(security_tool)}")
    
    # Test with proper parameters
    test_args = {
        'code_files': [{'filename': 'test.py', 'content': 'print("hello")'}]
    }
    
    try:
        result = security_tool.invoke(test_args)
        print(f"✅ Direct invoke result: {str(result)[:100]}...")
    except Exception as e:
        print(f"❌ Direct invoke error: {e}")

def test_mcp_in_langraph():
    """Test MCP tool through LangGraph ToolNode"""
    print("\n=== LANGGRAPH TOOLNODE TEST ===")
    
    workflow = CleanSWEWorkflow(Path.cwd(), Path.cwd() / 'output', False)
    
    # Create a simple state graph to test tool calling
    from typing import TypedDict, Annotated
    from langchain_core.messages import BaseMessage
    from langgraph.prebuilt import ToolNode
    
    class TestState(TypedDict):
        messages: Annotated[list[BaseMessage], "The messages in the conversation"]
    
    # Find security tool
    security_tool = None
    for tool in workflow.mcp_tools:
        if tool.name == 'security_check':
            security_tool = tool
            break
    
    if not security_tool:
        print("❌ Security tool not found!")
        return
    
    # Create ToolNode with just the security tool
    tool_node = ToolNode([security_tool])
    
    # Create a mock tool call message (simulating what agent would send)
    tool_call_message = AIMessage(
        content="I'll scan the code for security issues",
        tool_calls=[{
            "name": "security_check",
            "args": {
                "code_files": [{"filename": "test.py", "content": "print('hello')"}]
            },
            "id": "test_call_123"
        }]
    )
    
    # Test state
    test_state = {
        "messages": [tool_call_message]
    }
    
    print(f"Tool call args: {tool_call_message.tool_calls[0]['args']}")
    
    try:
        result = tool_node.invoke(test_state)
        print(f"✅ ToolNode result: {result}")
        
        # Check the messages in result
        if 'messages' in result:
            for msg in result['messages']:
                if isinstance(msg, ToolMessage):
                    print(f"Tool message content: {msg.content[:100]}...")
    except Exception as e:
        print(f"❌ ToolNode error: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    test_mcp_tool_direct()
    test_mcp_in_langraph()