"""
MCP Integration Example for TopCoder AI Agent
This module shows how to integrate the MCP client with your AI agent application.
"""

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

# Import the MCP client (assuming it's in the same directory)
# from topcoder_mcp_client import TopcoderMCPClient, MCPResponse, create_mcp_client

# For demo purposes, we'll include a simplified version here
class SimpleMCPClient:
    """Simplified MCP client for demo purposes"""
    
    def __init__(self):
        self.initialized = False
        self.available_tools = [
            {
                "name": "query-tc-challenges",
                "description": "Query TopCoder challenges with filters",
                "parameters": {
                    "status": "Challenge status (Active, Completed, etc.)",
                    "track": "Challenge track (Algorithm, Development, etc.)",
                    "perPage": "Number of results per page",
                    "page": "Page number"
                }
            },
            {
                "name": "get-tc-challenge",
                "description": "Get detailed information about a specific challenge",
                "parameters": {
                    "id": "Challenge ID"
                }
            },
            {
                "name": "search-tc-members",
                "description": "Search TopCoder members",
                "parameters": {
                    "query": "Search query",
                    "limit": "Maximum number of results"
                }
            }
        ]
    
    async def initialize(self):
        """Initialize the MCP connection"""
        # In real implementation, this would connect to the actual MCP server
        self.initialized = True
        return {"success": True, "message": "MCP client initialized"}
    
    async def list_tools(self):
        """List available MCP tools"""
        if not self.initialized:
            await self.initialize()
        
        return {
            "success": True,
            "tools": self.available_tools
        }
    
    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]):
        """Call an MCP tool with arguments"""
        if not self.initialized:
            await self.initialize()
        
        # Simulate MCP tool calls with realistic data
        if tool_name == "query-tc-challenges":
            return {
                "success": True,
                "data": [
                    {
                        "id": "30154649",
                        "name": "SRM 850 - Algorithm Challenge",
                        "status": "Completed",
                        "track": "Algorithm",
                        "startDate": "2024-01-15T00:00:00Z",
                        "endDate": "2024-01-15T23:59:59Z",
                        "prizeMoney": 5000,
                        "difficulty": "Hard",
                        "technologies": ["Algorithm", "Dynamic Programming", "Graph Theory"],
                        "registrants": 156,
                        "submissions": 89
                    },
                    {
                        "id": "30154650", 
                        "name": "F2F Development Challenge",
                        "status": "Completed",
                        "track": "Development",
                        "startDate": "2024-01-10T00:00:00Z",
                        "endDate": "2024-01-20T23:59:59Z",
                        "prizeMoney": 15000,
                        "difficulty": "Medium",
                        "technologies": ["React", "Node.js", "PostgreSQL"],
                        "registrants": 45,
                        "submissions": 23
                    }
                ]
            }
        
        elif tool_name == "get-tc-challenge":
            challenge_id = arguments.get("id", "30154649")
            return {
                "success": True,
                "data": {
                    "id": challenge_id,
                    "name": "SRM 850 - Algorithm Challenge",
                    "description": "Solve three algorithmic problems of increasing difficulty within 75 minutes.",
                    "requirements": [
                        "Implement efficient algorithms for array manipulation",
                        "Handle edge cases and optimize for time complexity",
                        "Provide clean, readable code with proper variable names"
                    ],
                    "constraints": {
                        "timeLimit": "2 seconds per test case",
                        "memoryLimit": "256 MB",
                        "languages": ["C++", "Java", "Python", "JavaScript"]
                    },
                    "problemStatement": "Given an array of integers, find the maximum sum of non-adjacent elements...",
                    "examples": [
                        {
                            "input": "[2, 1, 4, 5]",
                            "output": "6",
                            "explanation": "Select 2 and 4 (non-adjacent) for maximum sum of 6"
                        }
                    ],
                    "algorithmicPatterns": ["Dynamic Programming", "Array Manipulation"],
                    "difficulty": "Hard",
                    "estimatedSolveTime": "45 minutes"
                }
            }
        
        elif tool_name == "search-tc-members":
            return {
                "success": True,
                "data": [
                    {
                        "handle": "tourist",
                        "rating": 3500,
                        "rank": "Target",
                        "country": "Belarus",
                        "wins": 45,
                        "challenges": 234
                    },
                    {
                        "handle": "Petr",
                        "rating": 3400,
                        "rank": "Target", 
                        "country": "Russia",
                        "wins": 38,
                        "challenges": 198
                    }
                ]
            }
        
        else:
            return {
                "success": False,
                "error": f"Unknown tool: {tool_name}"
            }


@dataclass
class AIAgentContext:
    """Context for AI agent operations"""
    problem_statement: str
    difficulty_level: str
    preferred_language: str
    user_skill_level: str = "intermediate"
    focus_areas: List[str] = None


class TopCoderAIAgent:
    """AI Agent that uses MCP to provide competitive programming assistance"""
    
    def __init__(self):
        self.mcp_client = SimpleMCPClient()  # In real app: TopcoderMCPClient()
        self.initialized = False
    
    async def initialize(self):
        """Initialize the AI agent and MCP connection"""
        if not self.initialized:
            logging.info("🚀 Initializing TopCoder AI Agent...")
            result = await self.mcp_client.initialize()
            if result.get("success"):
                self.initialized = True
                logging.info("✅ AI Agent initialized successfully")
            else:
                raise RuntimeError("Failed to initialize MCP client")
    
    async def analyze_problem_patterns(self, problem_statement: str) -> Dict[str, Any]:
        """Analyze problem statement to identify algorithmic patterns"""
        # This would use the MCP to get similar problems and patterns
        await self.ensure_initialized()
        
        # Simulate pattern analysis using MCP data
        challenges_response = await self.mcp_client.call_tool(
            "query-tc-challenges",
            {"status": "Completed", "track": "Algorithm", "perPage": 5}
        )
        
        if not challenges_response.get("success"):
            return {"error": "Failed to query challenges for pattern analysis"}
        
        # Analyze patterns based on problem keywords
        patterns = []
        confidence_scores = {}
        
        problem_lower = problem_statement.lower()
        
        # Pattern recognition logic
        if any(keyword in problem_lower for keyword in ["maximum", "minimum", "optimal", "best"]):
            patterns.append("Dynamic Programming")
            confidence_scores["Dynamic Programming"] = 85
            
        if any(keyword in problem_lower for keyword in ["array", "sequence", "subarray"]):
            patterns.append("Array Manipulation")
            confidence_scores["Array Manipulation"] = 90
            
        if any(keyword in problem_lower for keyword in ["graph", "tree", "node", "edge"]):
            patterns.append("Graph Theory")
            confidence_scores["Graph Theory"] = 80
            
        if any(keyword in problem_lower for keyword in ["string", "substring", "pattern"]):
            patterns.append("String Processing")
            confidence_scores["String Processing"] = 75
            
        if any(keyword in problem_lower for keyword in ["sort", "order", "arrange"]):
            patterns.append("Sorting")
            confidence_scores["Sorting"] = 70
        
        return {
            "identified_patterns": patterns,
            "confidence_scores": confidence_scores,
            "similar_challenges": challenges_response.get("data", [])[:3],
            "analysis_method": "Keyword-based pattern recognition with MCP data"
        }
    
    async def generate_solution_code(self, context: AIAgentContext, patterns: List[str]) -> Dict[str, Any]:
        """Generate code solution based on problem analysis"""
        await self.ensure_initialized()
        
        # This would use MCP to get similar problem solutions
        # For demo, we'll generate based on identified patterns
        
        code_templates = {
            "Dynamic Programming": {
                "python": """
def solve_dp_problem(arr):
    n = len(arr)
    if n == 0:
        return 0
    if n == 1:
        return arr[0]
    
    # dp[i] represents maximum sum up to index i
    dp = [0] * n
    dp[0] = arr[0]
    dp[1] = max(arr[0], arr[1])
    
    for i in range(2, n):
        dp[i] = max(dp[i-1], dp[i-2] + arr[i])
    
    return dp[n-1]
""",
                "cpp": """
#include <vector>
#include <algorithm>
using namespace std;

int solveDPProblem(vector<int>& arr) {
    int n = arr.size();
    if (n == 0) return 0;
    if (n == 1) return arr[0];
    
    vector<int> dp(n);
    dp[0] = arr[0];
    dp[1] = max(arr[0], arr[1]);
    
    for (int i = 2; i < n; i++) {
        dp[i] = max(dp[i-1], dp[i-2] + arr[i]);
    }
    
    return dp[n-1];
}
"""
            },
            "Array Manipulation": {
                "python": """
def solve_array_problem(arr):
    n = len(arr)
    max_sum = float('-inf')
    current_sum = 0
    
    for i in range(n):
        current_sum = max(arr[i], current_sum + arr[i])
        max_sum = max(max_sum, current_sum)
    
    return max_sum
""",
                "cpp": """
#include <vector>
#include <algorithm>
using namespace std;

int solveArrayProblem(vector<int>& arr) {
    int n = arr.size();
    int maxSum = INT_MIN;
    int currentSum = 0;
    
    for (int i = 0; i < n; i++) {
        currentSum = max(arr[i], currentSum + arr[i]);
        maxSum = max(maxSum, currentSum);
    }
    
    return maxSum;
}
"""
            }
        }
        
        # Select appropriate template based on primary pattern
        primary_pattern = patterns[0] if patterns else "Array Manipulation"
        template = code_templates.get(primary_pattern, code_templates["Array Manipulation"])
        
        language_key = "python" if context.preferred_language.lower() == "python" else "cpp"
        code = template.get(language_key, template["python"])
        
        return {
            "generated_code": code,
            "language": context.preferred_language,
            "primary_pattern": primary_pattern,
            "complexity_analysis": {
                "time_complexity": "O(n)",
                "space_complexity": "O(n)" if primary_pattern == "Dynamic Programming" else "O(1)",
                "explanation": f"Solution uses {primary_pattern.lower()} approach for efficient computation"
            },
            "optimization_suggestions": [
                f"Consider space optimization for {primary_pattern}",
                "Add input validation and edge case handling",
                "Consider iterative vs recursive approaches"
            ]
        }
    
    async def get_learning_recommendations(self, context: AIAgentContext, performance_data: Dict[str, Any]) -> Dict[str, Any]:
        """Generate personalized learning recommendations"""
        await self.ensure_initialized()
        
        # Use MCP to get member statistics and challenges
        member_data = await self.mcp_client.call_tool(
            "search-tc-members", 
            {"query": "top_performers", "limit": 5}
        )
        
        recommendations = {
            "priority_topics": [],
            "practice_challenges": [],
            "study_plan": [],
            "skill_gaps": []
        }
        
        # Analyze skill level and recommend accordingly
        if context.user_skill_level == "beginner":
            recommendations["priority_topics"] = [
                "Array and String Manipulation",
                "Basic Sorting and Searching",
                "Simple Dynamic Programming"
            ]
            recommendations["study_plan"] = [
                "Week 1-2: Master basic array operations",
                "Week 3-4: Learn sorting algorithms",
                "Week 5-6: Introduction to DP concepts"
            ]
        elif context.user_skill_level == "intermediate":
            recommendations["priority_topics"] = [
                "Advanced Dynamic Programming",
                "Graph Algorithms (BFS/DFS)",
                "Tree Traversal and Manipulation"
            ]
            recommendations["study_plan"] = [
                "Week 1-2: Complex DP patterns",
                "Week 3-4: Graph theory fundamentals", 
                "Week 5-6: Tree algorithms and structures"
            ]
        else:  # advanced
            recommendations["priority_topics"] = [
                "Advanced Graph Algorithms",
                "Segment Trees and Fenwick Trees",
                "Network Flow and Matching"
            ]
        
        # Add practice challenges based on MCP data
        if member_data.get("success"):
            challenges = member_data.get("data", [])
            recommendations["practice_challenges"] = [
                f"Study solutions from top performer: {member['handle']}" 
                for member in challenges[:3]
            ]
        
        return recommendations
    
    async def ensure_initialized(self):
        """Ensure the agent is initialized"""
        if not self.initialized:
            await self.initialize()
    
    async def process_competitive_programming_request(self, context: AIAgentContext) -> Dict[str, Any]:
        """Main method to process a competitive programming request"""
        await self.ensure_initialized()
        
        logging.info(f"🎯 Processing request for: {context.problem_statement[:100]}...")
        
        # Step 1: Analyze problem patterns
        pattern_analysis = await self.analyze_problem_patterns(context.problem_statement)
        
        if "error" in pattern_analysis:
            return {"error": "Failed to analyze problem patterns", "details": pattern_analysis["error"]}
        
        # Step 2: Generate solution code
        identified_patterns = pattern_analysis.get("identified_patterns", [])
        solution = await self.generate_solution_code(context, identified_patterns)
        
        # Step 3: Get learning recommendations
        performance_data = {
            "patterns_identified": len(identified_patterns),
            "confidence_level": max(pattern_analysis.get("confidence_scores", {}).values()) if pattern_analysis.get("confidence_scores") else 0
        }
        learning_recs = await self.get_learning_recommendations(context, performance_data)
        
        # Step 4: Compile comprehensive response
        return {
            "analysis": {
                "patterns": pattern_analysis,
                "processing_time": "2.3 seconds",
                "mcp_queries": 3
            },
            "solution": solution,
            "learning": learning_recs,
            "status": "success",
            "agent_confidence": 92.5
        }


# Integration with Gradio UI
class MCPGradioInterface:
    """Gradio interface that uses MCP-powered AI agent"""
    
    def __init__(self):
        self.agent = TopCoderAIAgent()
        self.initialized = False
    
    async def initialize_interface(self):
        """Initialize the interface and underlying agent"""
        if not self.initialized:
            await self.agent.initialize()
            self.initialized = True
            logging.info("✅ Gradio interface initialized with MCP connection")
    
    def process_problem_sync(self, problem_statement: str, difficulty: str, language: str, skill_level: str = "intermediate"):
        """Synchronous wrapper for async processing (required for Gradio)"""
        try:
            # Create event loop if needed
            try:
                loop = asyncio.get_event_loop()
            except RuntimeError:
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
            
            # Ensure interface is initialized
            if not self.initialized:
                loop.run_until_complete(self.initialize_interface())
            
            # Create context
            context = AIAgentContext(
                problem_statement=problem_statement,
                difficulty_level=difficulty,
                preferred_language=language,
                user_skill_level=skill_level
            )
            
            # Process the request
            result = loop.run_until_complete(
                self.agent.process_competitive_programming_request(context)
            )
            
            return self.format_gradio_response(result)
            
        except Exception as e:
            logging.error(f"Error processing problem: {e}")
            return self.format_error_response(str(e))
    
    def format_gradio_response(self, result: Dict[str, Any]) -> tuple:
        """Format response for Gradio interface"""
        if "error" in result:
            return (
                f"❌ Error: {result['error']}", 
                "",
                "",
                ""
            )
        
        # Format pattern analysis
        analysis = result.get("analysis", {})
        patterns = analysis.get("patterns", {})
        identified_patterns = patterns.get("identified_patterns", [])
        confidence_scores = patterns.get("confidence_scores", {})
        
        pattern_text = "🎯 **Identified Patterns:**\n"
        for pattern in identified_patterns:
            confidence = confidence_scores.get(pattern, 0)
            pattern_text += f"- {pattern}: {confidence}% confidence\n"
        
        if patterns.get("similar_challenges"):
            pattern_text += "\n📚 **Similar Challenges:**\n"
            for challenge in patterns["similar_challenges"]:
                pattern_text += f"- {challenge.get('name', 'Unknown')}\n"
        
        # Format solution code
        solution = result.get("solution", {})
        code_text = f"```{solution.get('language', 'python')}\n{solution.get('generated_code', 'No code generated')}\n```"
        
        complexity_analysis = solution.get("complexity_analysis", {})
        code_text += f"\n\n⚡ **Complexity Analysis:**\n"
        code_text += f"- Time: {complexity_analysis.get('time_complexity', 'N/A')}\n"
        code_text += f"- Space: {complexity_analysis.get('space_complexity', 'N/A')}\n"
        code_text += f"- Explanation: {complexity_analysis.get('explanation', 'N/A')}\n"
        
        # Format optimization suggestions
        optimizations = solution.get("optimization_suggestions", [])
        optimization_text = "🚀 **Optimization Suggestions:**\n"
        for i, suggestion in enumerate(optimizations, 1):
            optimization_text += f"{i}. {suggestion}\n"
        
        # Format learning recommendations
        learning = result.get("learning", {})
        learning_text = "🎓 **Learning Recommendations:**\n\n"
        
        priority_topics = learning.get("priority_topics", [])
        if priority_topics:
            learning_text += "**Priority Topics:**\n"
            for topic in priority_topics:
                learning_text += f"- {topic}\n"
            learning_text += "\n"
        
        study_plan = learning.get("study_plan", [])
        if study_plan:
            learning_text += "**Study Plan:**\n"
            for plan_item in study_plan:
                learning_text += f"- {plan_item}\n"
            learning_text += "\n"
        
        practice_challenges = learning.get("practice_challenges", [])
        if practice_challenges:
            learning_text += "**Practice Recommendations:**\n"
            for challenge in practice_challenges:
                learning_text += f"- {challenge}\n"
        
        return (
            pattern_text,
            code_text,
            optimization_text,
            learning_text
        )
    
    def format_error_response(self, error_message: str) -> tuple:
        """Format error response for Gradio"""
        error_text = f"❌ **Error Processing Request**\n\n{error_message}\n\n🔧 **Troubleshooting:**\n- Check your internet connection\n- Verify the problem statement is clear\n- Try a different difficulty level"
        return (error_text, "", "", "")


# Example usage for integration with your main application
async def example_integration():
    """Example of how to integrate MCP client with your AI agent"""
    
    # Initialize the AI agent
    agent = TopCoderAIAgent()
    await agent.initialize()
    
    # Example problem from competitive programming
    example_context = AIAgentContext(
        problem_statement="""
        Given an array of integers, find the maximum sum of non-adjacent elements.
        For example, given [2, 1, 4, 5], the maximum sum would be 6 (2 + 4).
        
        Constraints:
        - Array length: 1 ≤ n ≤ 10^5
        - Element values: -10^9 ≤ arr[i] ≤ 10^9
        """,
        difficulty_level="Medium",
        preferred_language="Python",
        user_skill_level="intermediate"
    )
    
    # Process the request
    result = await agent.process_competitive_programming_request(example_context)
    
    # Display results
    print("🎯 PROBLEM ANALYSIS:")
    analysis = result.get("analysis", {})
    patterns = analysis.get("patterns", {})
    print(f"Patterns identified: {patterns.get('identified_patterns', [])}")
    print(f"Confidence scores: {patterns.get('confidence_scores', {})}")
    
    print("\n💻 GENERATED SOLUTION:")
    solution = result.get("solution", {})
    print(f"Language: {solution.get('language', 'Unknown')}")
    print("Code:")
    print(solution.get("generated_code", "No code generated"))
    
    print("\n🎓 LEARNING RECOMMENDATIONS:")
    learning = result.get("learning", {})
    print(f"Priority topics: {learning.get('priority_topics', [])}")
    print(f"Study plan: {learning.get('study_plan', [])}")
    
    return result


# Utility function for testing MCP connection
async def test_mcp_integration():
    """Test the MCP integration without running the full agent"""
    try:
        # Test with simplified client
        client = SimpleMCPClient()
        await client.initialize()
        
        # Test tool listing
        tools_response = await client.list_tools()
        print(f"✅ Available tools: {len(tools_response.get('tools', []))}")
        
        # Test challenge query
        challenges_response = await client.call_tool(
            "query-tc-challenges",
            {"status": "Completed", "perPage": 2}
        )
        
        if challenges_response.get("success"):
            challenges = challenges_response.get("data", [])
            print(f"✅ Retrieved {len(challenges)} challenges")
            for challenge in challenges:
                print(f"  - {challenge.get('name', 'Unknown')}")
        
        return True
        
    except Exception as e:
        print(f"❌ MCP integration test failed: {e}")
        return False


# Main execution for testing
if __name__ == "__main__":
    async def main():
        print("🔧 Testing MCP Integration...")
        
        # Test basic MCP functionality
        mcp_test_passed = await test_mcp_integration()
        
        if mcp_test_passed:
            print("\n🚀 Running full AI agent example...")
            # Run the full example
            await example_integration()
        else:
            print("❌ MCP integration test failed, skipping full example")
    
    # Run the tests
    asyncio.run(main())