import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import random
import json
from typing import Dict, List, Optional
import numpy as np
import hashlib
import datetime
import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity

class UserManager:
    def __init__(self):
        # Use a simple in-memory dictionary instead of Firebase
        self.users = {}
        self.progress = {}
    
    def register_user(self, username: str, password: str) -> bool:
        if username in self.users:
            return False
        
        # Hash password
        hashed_password = hashlib.sha256(password.encode()).hexdigest()
        
        # Store user data
        self.users[username] = {
            'password_hash': hashed_password,
            'join_date': datetime.datetime.now(),
            'solved_problems': [],
            'skill_level': 'beginner'
        }
        self.progress[username] = {
            'arrays': 0,
            'graphs': 0,
            'dynamic_programming': 0,
            'trees': 0
        }
        return True
    
    def authenticate(self, username: str, password: str) -> bool:
        if username not in self.users:
            return False
        
        hashed_password = hashlib.sha256(password.encode()).hexdigest()
        return self.users[username]['password_hash'] == hashed_password
    
    def update_progress(self, username: str, topic: str, success: bool):
        if success:
            self.progress[username][topic] += 1
        
        # Update skill level based on progress
        total_solved = sum(self.progress[username].values())
        if total_solved > 50:
            self.users[username]['skill_level'] = 'advanced'
        elif total_solved > 20:
            self.users[username]['skill_level'] = 'intermediate'

class EnhancedProblemGenerator:
    def __init__(self):
        self.templates = {
            "arrays": [
                {
                    "title": "Two Sum",
                    "difficulty": "Easy",
                    "description": "Given an array of integers nums and an integer target, return indices of the two numbers that add up to target.",
                    "constraints": ["2 <= nums.length <= 104", "-109 <= nums[i] <= 109"],
                    "test_cases": [
                        {"input": "[2,7,11,15], target=9", "output": "[0,1]"},
                        {"input": "[3,2,4], target=6", "output": "[1,2]"}
                    ],
                    "hints": [
                        "Consider using a hash map to store previously seen numbers",
                        "Think about the complement of each number relative to the target"
                    ]
                },
                {
                    "title": "Maximum Subarray",
                    "difficulty": "Easy",
                    "description": "Find the contiguous subarray with the largest sum.",
                    "constraints": ["1 <= nums.length <= 105", "-104 <= nums[i] <= 104"],
                    "test_cases": [
                        {"input": "[-2,1,-3,4,-1,2,1,-5,4]", "output": "6"},
                        {"input": "[1]", "output": "1"}
                    ],
                    "hints": [
                        "Consider Kadane's algorithm",
                        "Think about when to reset your current sum"
                    ]
                }
            ],
            "dynamic_programming": [
                {
                    "title": "Climbing Stairs",
                    "difficulty": "Easy",
                    "description": "You are climbing a staircase. It takes n steps to reach the top. Each time you can either climb 1 or 2 steps. In how many distinct ways can you climb to the top?",
                    "constraints": ["1 <= n <= 45"],
                    "test_cases": [
                        {"input": "2", "output": "2"},
                        {"input": "3", "output": "3"}
                    ],
                    "hints": [
                        "Think about the Fibonacci sequence",
                        "Consider the last step you take"
                    ]
                }
            ],
            "trees": [
                {
                    "title": "Maximum Depth of Binary Tree",
                    "difficulty": "Easy",
                    "description": "Given the root of a binary tree, return its maximum depth.",
                    "constraints": ["The number of nodes in the tree is in the range [0, 104]"],
                    "test_cases": [
                        {"input": "[3,9,20,null,null,15,7]", "output": "3"},
                        {"input": "[1,null,2]", "output": "2"}
                    ],
                    "hints": [
                        "Consider using recursion",
                        "Think about the base case of an empty tree"
                    ]
                }
            ]
        }
        
        # Load AI model for hint generation
        self.tokenizer = AutoTokenizer.from_pretrained("microsoft/CodeGPT-small-py")
        self.model = AutoModelForCausalLM.from_pretrained("microsoft/CodeGPT-small-py")
    
    def generate_ai_hint(self, problem: Dict, user_code: str) -> str:
        prompt = f"""
Problem: {problem['description']}
User's code: {user_code}
Hint: Let me help you think about this problem. 
"""
        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
        outputs = self.model.generate(**inputs, max_length=200, num_return_sequences=1)
        hint = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return hint.split("Hint: ")[-1].strip()
    
    def get_problem_for_user(self, username: str, user_manager: UserManager, topic: str) -> Dict:
        user_level = user_manager.users[username]['skill_level']
        available_problems = self.templates.get(topic.lower(), [])
        
        if not available_problems:
            return {"error": "Topic not found"}
        
        # Filter problems based on user level
        if user_level == 'beginner':
            problems = [p for p in available_problems if p['difficulty'] == 'Easy']
        elif user_level == 'intermediate':
            problems = [p for p in available_problems if p['difficulty'] in ['Easy', 'Medium']]
        else:
            problems = available_problems
        
        return random.choice(problems)

class EnhancedCodeEvaluator:
    def __init__(self):
        self.test_environment = {}
    
    def analyze_complexity(self, code: str) -> str:
        # Simple complexity analysis based on nested loops
        if "while" not in code and "for" not in code:
            return "O(1)"
        elif code.count("for") + code.count("while") == 1:
            return "O(n)"
        elif code.count("for") + code.count("while") == 2:
            return "O(n²)"
        else:
            return "O(n³) or higher"
    
    def evaluate_code(self, code: str, test_cases: List[Dict]) -> Dict:
        try:
            # Create safe execution environment
            local_dict = {}
            exec(code, {"__builtins__": {}}, local_dict)
            
            results = []
            passed = 0
            
            for test in test_cases:
                try:
                    # Execute test case
                    test_result = eval(f"solution{test['input']}", local_dict)
                    expected = eval(test['output'])
                    
                    if test_result == expected:
                        passed += 1
                        results.append({"status": "PASS", "input": test['input']})
                    else:
                        results.append({
                            "status": "FAIL",
                            "input": test['input'],
                            "expected": expected,
                            "got": test_result
                        })
                except Exception as e:
                    results.append({
                        "status": "ERROR",
                        "input": test['input'],
                        "error": str(e)
                    })
            
            # Analyze code complexity
            complexity = self.analyze_complexity(code)
            
            return {
                "success": True,
                "results": results,
                "passed": passed,
                "total": len(test_cases),
                "complexity": complexity
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }

class EnhancedLeetCodeEducator:
    def __init__(self):
        self.problem_generator = EnhancedProblemGenerator()
        self.code_evaluator = EnhancedCodeEvaluator()
        self.user_manager = UserManager()
        
        # Check for GPU availability
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")
        
        # Initialize hint generator
        self.hint_generator = pipeline("text-generation", model="microsoft/CodeGPT-small-py")
    
    def register_user(self, username: str, password: str) -> str:
        if self.user_manager.register_user(username, password):
            return f"Successfully registered user: {username}"
        return "Username already exists"
    
    def login(self, username: str, password: str) -> bool:
        return self.user_manager.authenticate(username, password)
    
    def get_problem(self, username: str, topic: str) -> str:
        if not self.user_manager.authenticate(username, ""):  # Check if user exists
            return "Please log in first"
        
        problem = self.problem_generator.get_problem_for_user(username, self.user_manager, topic)
        if "error" in problem:
            return "Topic not found. Available topics: arrays, graphs, dynamic_programming, trees"
        
        return f"""
Problem: {problem['title']}
Difficulty: {problem['difficulty']}

Description:
{problem['description']}

Constraints:
{chr(10).join(['- ' + c for c in problem['constraints']])}

Example Test Cases:
{chr(10).join([f'Input: {tc["input"]}, Output: {tc["output"]}' for tc in problem['test_cases']])}

Your current skill level: {self.user_manager.users[username]['skill_level']}
Problems solved in this topic: {self.user_manager.progress[username][topic]}
"""
    
    def get_hint(self, username: str, topic: str, code: str) -> str:
        if not self.user_manager.authenticate(username, ""):
            return "Please log in first"
        
        problem = self.problem_generator.get_problem_for_user(username, self.user_manager, topic)
        if "error" in problem:
            return "Cannot generate hint: problem not found"
        
        if not code:
            # If no code provided, return a general hint
            return random.choice(problem['hints'])
        
        # Generate personalized hint based on user's code
        return self.problem_generator.generate_ai_hint(problem, code)
    
    def evaluate_submission(self, username: str, code: str, topic: str) -> str:
        if not self.user_manager.authenticate(username, ""):
            return "Please log in first"
        
        problem = self.problem_generator.get_problem_for_user(username, self.user_manager, topic)
        if "error" in problem:
            return "Error: Could not find problem for evaluation"
        
        results = self.code_evaluator.evaluate_code(code, problem['test_cases'])
        
        if not results['success']:
            return f"Error in code execution: {results['error']}"
        
        # Update user progress
        self.user_manager.update_progress(username, topic, results['passed'] == results['total'])
        
        output = f"""
Evaluation Results:
Passed: {results['passed']}/{results['total']} test cases
Time Complexity: {results['complexity']}

Details:
"""
        for idx, result in enumerate(results['results']):
            output += f"\nTest Case {idx + 1}: {result['status']}"
            if result['status'] == 'FAIL':
                output += f"\n  Expected: {result['expected']}"
                output += f"\n  Got: {result['got']}"
            elif result['status'] == 'ERROR':
                output += f"\n  Error: {result['error']}"
        
        if results['passed'] < results['total']:
            output += f"\n\nNeed a hint? Type 'hint' in the code box to get personalized help!"
        
        return output

def create_enhanced_gradio_interface():
    educator = EnhancedLeetCodeEducator()
    
    def process_interaction(username, password, action, topic, code):
        if action == "Register":
            return educator.register_user(username, password)
        elif action == "Login":
            if educator.login(username, password):
                return "Login successful!"
            return "Invalid credentials"
        elif action == "Get Problem":
            return educator.get_problem(username, topic)
        elif action == "Submit Solution":
            if code.strip().lower() == "hint":
                return educator.get_hint(username, topic, "")
            else:
                return educator.evaluate_submission(username, code, topic)
        else:
            return "Invalid action selected."
    
    # Create the interface
    iface = gr.Interface(
        fn=process_interaction,
        inputs=[
            gr.Textbox(label="Username"),
            gr.Textbox(label="Password", type="password"),
            gr.Radio(
                choices=["Register", "Login", "Get Problem", "Submit Solution"],
                label="Action",
                value="Register"
            ),
            gr.Dropdown(
                choices=["arrays", "dynamic_programming", "trees"],
                label="Topic",
                value="arrays"
            ),
            gr.Code(language="python", label="Your Solution (type 'hint' for help)")
        ],
        outputs=gr.Textbox(label="Output", lines=10),
        title="Enhanced LeetCode Educational Assistant",
        description="""
        Welcome to your personalized coding practice platform!
        
        1. Register or login to get started
        2. Choose a topic and get a problem
        3. Write your solution or type 'hint' for help
        4. Submit your code for evaluation
        
        Your progress is tracked and problems are tailored to your skill level!
        """,
        examples=[
            ["new_user", "password123", "Register", "arrays", ""],
            ["new_user", "password123", "Login", "arrays", ""],
            ["new_user", "password123", "Get Problem", "arrays", ""],
            ["new_user", "password123", "Submit Solution", "arrays", "def solution(nums, target):\n    seen = {}\n    for i, num in enumerate(nums):\n        complement = target - num\n        if complement in seen:\n            return [seen[complement], i]\n        seen[num] = i\n    return []"]
        ],
        theme="default"
    )
    
    return iface

# Create and launch the interface
interface = create_enhanced_gradio_interface()
interface.launch()