"""
Context-Aware Code Generator for Metis Agent.

This module provides intelligent code generation based on project context,
existing patterns, and codebase structure.
"""
import os
import re
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple
import json

from .base import BaseTool
from .code_analysis import CodeAnalysisTool
from .project_context import ProjectContextTool


class ContextAwareCodeGenerator(BaseTool):
    """Tool for generating code based on project context and existing patterns."""
    
    def __init__(self):
        self.code_analyzer = CodeAnalysisTool()
        self.project_context = ProjectContextTool()
        
        # Code templates for different languages
        self.templates = {
            'python': {
                'function': '''def {name}({params}):
    """
    {docstring}
    """
    {body}''',
                'class': '''class {name}({base_classes}):
    """
    {docstring}
    """
    
    def __init__(self{init_params}):
        {init_body}
    
    {methods}''',
                'test': '''import unittest
{imports}

class Test{class_name}(unittest.TestCase):
    """Test cases for {class_name}."""
    
    def setUp(self):
        """Set up test fixtures before each test method."""
        {setup_code}
    
    {test_methods}

if __name__ == '__main__':
    unittest.main()'''
            },
            'javascript': {
                'function': '''function {name}({params}) {{
    /**
     * {docstring}
     */
    {body}
}}''',
                'class': '''class {name}{extends} {{
    /**
     * {docstring}
     */
    constructor({constructor_params}) {{
        {constructor_body}
    }}
    
    {methods}
}}''',
                'test': '''import {{ describe, it, expect{test_imports} }} from '{test_framework}';
{imports}

describe('{class_name}', () => {{
    {test_methods}
}});'''
            }
        }
    
    def can_handle(self, task: str) -> bool:
        """Check if this tool can handle the given task."""
        generation_keywords = [
            'generate', 'create', 'write', 'build', 'make', 'implement',
            'code', 'function', 'class', 'method', 'test', 'documentation',
            'refactor', 'improve', 'optimize', 'template', 'boilerplate',
            'scaffold', 'add', 'new'
        ]
        task_lower = task.lower()
        return any(keyword in task_lower for keyword in generation_keywords)
    
    def execute(self, task: str) -> Any:
        """Execute code generation based on the task."""
        task_lower = task.lower()
        
        try:
            # Get project context first
            project_context = self._get_project_context()
            
            if 'function' in task_lower:
                return self._generate_function(task, project_context)
            elif 'class' in task_lower:
                return self._generate_class(task, project_context)
            elif 'test' in task_lower:
                return self._generate_tests(task, project_context)
            elif 'documentation' in task_lower or 'docs' in task_lower:
                return self._generate_documentation(task, project_context)
            elif 'refactor' in task_lower:
                return self._suggest_refactoring(task, project_context)
            else:
                # Generic code generation
                return self._generate_code(task, project_context)
                
        except Exception as e:
            return {"error": f"Code generation failed: {str(e)}"}
    
    def _get_project_context(self, directory: str = ".") -> Dict[str, Any]:
        """Get comprehensive project context for code generation."""
        # Get project analysis
        project_analysis = self.project_context.analyze_directory(directory)
        code_analysis = self.code_analyzer.analyze_project(directory)
        
        context = {
            "project_type": project_analysis.get("project_type", "unknown"),
            "primary_language": code_analysis.get("summary", {}).get("primary_language", "python"),
            "languages": code_analysis.get("languages", {}),
            "common_imports": code_analysis.get("common_imports", {}),
            "patterns": code_analysis.get("project_patterns", []),
            "existing_classes": [],
            "existing_functions": [],
            "coding_style": self._analyze_coding_style(directory),
            "dependencies": project_analysis.get("dependencies", {}),
            "framework": project_analysis.get("frameworks", {})
        }
        
        # Extract existing code elements for context
        for language_files in code_analysis.get("files_by_language", {}).values():
            for file_path in language_files:
                file_analysis = self.code_analyzer.analyze_file(file_path)
                if file_analysis.get("success"):
                    context["existing_classes"].extend(file_analysis.get("classes", []))
                    context["existing_functions"].extend(file_analysis.get("functions", []))
        
        return context
    
    def _analyze_coding_style(self, directory: str = ".") -> Dict[str, Any]:
        """Analyze the coding style patterns in the project."""
        style = {
            "indentation": "4 spaces",  # default
            "naming_convention": "snake_case",  # default for Python
            "docstring_style": "google",  # default
            "import_style": "absolute",
            "line_length": 88  # default
        }
        
        try:
            dir_path = Path(directory).resolve()
            sample_files = []
            
            # Get a few sample Python files to analyze style
            for file_path in dir_path.rglob("*.py"):
                if file_path.is_file() and len(sample_files) < 5:
                    sample_files.append(file_path)
            
            if sample_files:
                # Analyze indentation
                indentations = []
                for file_path in sample_files:
                    try:
                        with open(file_path, 'r', encoding='utf-8') as f:
                            content = f.read()
                            lines = content.splitlines()
                            for line in lines:
                                if line.startswith(' ') and line.strip():
                                    leading_spaces = len(line) - len(line.lstrip(' '))
                                    if leading_spaces > 0:
                                        indentations.append(leading_spaces)
                    except Exception:
                        # Skip files that can't be read
                        continue
                
                if indentations:
                    # Find most common indentation
                    common_indent = max(set(indentations), key=indentations.count)
                    if common_indent == 2:
                        style["indentation"] = "2 spaces"
                    elif common_indent == 4:
                        style["indentation"] = "4 spaces"
                    elif common_indent == 8:
                        style["indentation"] = "8 spaces"
            
        except Exception:
            # Use defaults if analysis fails
            pass
        
        return style
    
    def _generate_function(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate a function based on the task description and project context."""
        primary_language = context["primary_language"]
        
        # Extract function details from task
        function_info = self._extract_function_info(task, context)
        
        if primary_language == "python":
            return self._generate_python_function(function_info, context)
        elif primary_language in ["javascript", "typescript"]:
            return self._generate_javascript_function(function_info, context)
        else:
            return {"error": f"Code generation not supported for {primary_language}"}
    
    def _extract_function_info(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Extract function information from task description."""
        # Simple NLP-like extraction
        info = {
            "name": "new_function",
            "description": task,
            "parameters": [],
            "return_type": None,
            "async": False
        }
        
        # Try to extract function name
        name_patterns = [
            r'function\s+called\s+(\w+)',
            r'function\s+named\s+(\w+)',
            r'create\s+(?:a\s+)?(?:simple\s+)?(?:python\s+)?function\s+(?:to\s+)?(?:calculate\s+)?(?:the\s+)?(\w+)',
            r'implement\s+(\w+)\s*\(',
            r'def\s+(\w+)\s*\(',
            r'write\s+(?:a\s+)?(\w+)\s+function',
            r'create\s+(?:a\s+)?(\w+)\s+function',
            r'build\s+(?:a\s+)?(\w+)\s+function',
            r'generate\s+(?:a\s+)?(\w+)\s+function',
        ]
        
        for pattern in name_patterns:
            match = re.search(pattern, task, re.IGNORECASE)
            if match:
                name = match.group(1)
                # Clean up common words that shouldn't be function names
                if name.lower() not in ['function', 'to', 'the', 'a', 'simple', 'python']:
                    info["name"] = name
                    break
        
        # Try to extract parameters
        param_patterns = [
            r'with\s+parameters?\s+(.+?)(?:\.|$)',
            r'takes?\s+(.+?)\s+as\s+(?:parameters?|arguments?)',
            r'\(([^)]+)\)'
        ]
        
        for pattern in param_patterns:
            match = re.search(pattern, task, re.IGNORECASE)
            if match:
                param_text = match.group(1)
                # Simple parameter extraction
                params = [p.strip() for p in param_text.split(',') if p.strip()]
                info["parameters"] = params[:5]  # Limit to 5 parameters
                break
        
        # Check for async
        if 'async' in task.lower() or 'asynchronous' in task.lower():
            info["async"] = True
        
        return info
    
    def _generate_python_function(self, function_info: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate a Python function."""
        # Generate appropriate imports based on context
        imports = []
        common_imports = context.get("common_imports", {})
        
        # Add common imports if relevant
        if any(word in function_info["description"].lower() for word in ["file", "path", "directory"]):
            if "pathlib" in common_imports or "os" in common_imports:
                imports.append("from pathlib import Path")
        
        if any(word in function_info["description"].lower() for word in ["json", "data"]):
            if "json" in common_imports:
                imports.append("import json")
        
        if any(word in function_info["description"].lower() for word in ["request", "http", "api"]):
            if "requests" in common_imports:
                imports.append("import requests")
        
        # Generate parameters with type hints
        params = function_info["parameters"]
        if not params:
            # Infer parameters based on task description
            if "factorial" in function_info["description"].lower():
                params = ["n"]
            elif "fibonacci" in function_info["description"].lower():
                params = ["n"]
            elif "sum" in function_info["description"].lower() or "add" in function_info["description"].lower():
                params = ["numbers"]
            else:
                params = ["data"]
        
        param_str = ", ".join(f"{param}: Any" for param in params)
        return_type = " -> Any"
        
        # Generate function body based on description
        body = self._generate_python_function_body(function_info, context)
        
        # Apply project coding style
        style = context.get("coding_style", {})
        indent = "    " if style.get("indentation") == "4 spaces" else "  "
        
        # Build the function code
        function_name = function_info["name"]
        if function_info["async"]:
            func_def = f"async def {function_name}({param_str}){return_type}:"
        else:
            func_def = f"def {function_name}({param_str}){return_type}:"
        
        docstring = f'"""\n{indent}{function_info["description"]}\n{indent}"""'
        
        code = f"{func_def}\n{indent}{docstring}\n{body}"
        
        result = {
            "success": True,
            "language": "python",
            "code": code,
            "imports": imports,
            "function_name": function_name,
            "description": function_info["description"],
            "context_used": {
                "common_imports": list(common_imports.keys())[:5],
                "coding_style": style,
                "project_patterns": context.get("patterns", [])[:3]
            }
        }
        
        return result
    
    def _generate_python_function_body(self, function_info: Dict[str, Any], context: Dict[str, Any]) -> str:
        """Generate Python function body based on description."""
        description = function_info["description"].lower()
        indent = "    "
        
        # Specific implementations for common functions
        if "factorial" in description:
            return f"{indent}if n < 0:\n{indent}    raise ValueError(\"Factorial is not defined for negative numbers\")\n{indent}if n == 0 or n == 1:\n{indent}    return 1\n{indent}result = 1\n{indent}for i in range(2, n + 1):\n{indent}    result *= i\n{indent}return result"
        
        elif "fibonacci" in description:
            return f"{indent}if n < 0:\n{indent}    raise ValueError(\"Fibonacci sequence is not defined for negative numbers\")\n{indent}if n == 0:\n{indent}    return 0\n{indent}elif n == 1:\n{indent}    return 1\n{indent}a, b = 0, 1\n{indent}for _ in range(2, n + 1):\n{indent}    a, b = b, a + b\n{indent}return b"
        
        elif "prime" in description and "check" in description:
            return f"{indent}if n < 2:\n{indent}    return False\n{indent}if n == 2:\n{indent}    return True\n{indent}if n % 2 == 0:\n{indent}    return False\n{indent}for i in range(3, int(n ** 0.5) + 1, 2):\n{indent}    if n % i == 0:\n{indent}        return False\n{indent}return True"
        
        elif "palindrome" in description:
            return f"{indent}return str(data).lower() == str(data).lower()[::-1]"
        
        elif "reverse" in description and "string" in description:
            return f"{indent}return str(data)[::-1]"
        
        elif "sum" in description or "add" in description:
            return f"{indent}if isinstance(numbers, (list, tuple)):\n{indent}    return sum(numbers)\n{indent}elif hasattr(numbers, '__iter__'):\n{indent}    return sum(numbers)\n{indent}else:\n{indent}    return numbers"
        
        elif "average" in description or "mean" in description:
            return f"{indent}if isinstance(numbers, (list, tuple)) and len(numbers) > 0:\n{indent}    return sum(numbers) / len(numbers)\n{indent}elif hasattr(numbers, '__iter__'):\n{indent}    nums = list(numbers)\n{indent}    return sum(nums) / len(nums) if nums else 0\n{indent}else:\n{indent}    return numbers"
        
        elif "max" in description or "maximum" in description:
            return f"{indent}if isinstance(numbers, (list, tuple)) and numbers:\n{indent}    return max(numbers)\n{indent}elif hasattr(numbers, '__iter__'):\n{indent}    return max(numbers)\n{indent}else:\n{indent}    return numbers"
        
        elif "min" in description or "minimum" in description:
            return f"{indent}if isinstance(numbers, (list, tuple)) and numbers:\n{indent}    return min(numbers)\n{indent}elif hasattr(numbers, '__iter__'):\n{indent}    return min(numbers)\n{indent}else:\n{indent}    return numbers"
        
        elif any(word in description for word in ["read", "file", "load"]):
            return f"{indent}try:\n{indent}    with open(data, 'r', encoding='utf-8') as file:\n{indent}        return file.read()\n{indent}except FileNotFoundError:\n{indent}    return f\"File '{{data}}' not found\"\n{indent}except Exception as e:\n{indent}    return f\"Error reading file: {{e}}\""
        
        elif any(word in description for word in ["save", "write", "store"]):
            return f"{indent}try:\n{indent}    with open(data, 'w', encoding='utf-8') as file:\n{indent}        file.write(str(data))\n{indent}    return True\n{indent}except Exception as e:\n{indent}    print(f'Error saving file: {{e}}')\n{indent}    return False"
        
        elif any(word in description for word in ["sort", "order"]):
            return f"{indent}if isinstance(data, (list, tuple)):\n{indent}    return sorted(data)\n{indent}elif isinstance(data, str):\n{indent}    return ''.join(sorted(data))\n{indent}else:\n{indent}    return data"
        
        elif any(word in description for word in ["calculate", "compute", "math"]):
            return f"{indent}# Implement your calculation logic here\n{indent}result = 0\n{indent}# Add your mathematical operations\n{indent}return result"
        
        else:
            return f"{indent}# Implement your function logic here\n{indent}return data"
    
    def _generate_class(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate a class based on the task description."""
        primary_language = context["primary_language"]
        
        class_info = self._extract_class_info(task, context)
        
        if primary_language == "python":
            return self._generate_python_class(class_info, context)
        else:
            return {"error": f"Class generation not supported for {primary_language}"}
    
    def _extract_class_info(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Extract class information from task description."""
        info = {
            "name": "NewClass",
            "description": task,
            "methods": [],
            "attributes": [],
            "base_classes": []
        }
        
        # Extract class name
        name_patterns = [
            r'class\s+called\s+(\w+)',
            r'class\s+named\s+(\w+)',
            r'create\s+(?:a\s+)?(\w+)\s+class',
            r'implement\s+(?:a\s+)?(\w+)\s+class'
        ]
        
        for pattern in name_patterns:
            match = re.search(pattern, task, re.IGNORECASE)
            if match:
                info["name"] = match.group(1).title()  # Capitalize for class names
                break
        
        # Extract methods
        if "methods" in task.lower():
            method_patterns = [
                r'methods?\s+(?:called\s+)?(.+?)(?:\.|$)',
                r'with\s+(.+?)\s+methods?'
            ]
            for pattern in method_patterns:
                match = re.search(pattern, task, re.IGNORECASE)
                if match:
                    methods_text = match.group(1)
                    methods = [m.strip() for m in methods_text.split(',') if m.strip()]
                    info["methods"] = methods[:5]  # Limit methods
                    break
        
        return info
    
    def _generate_python_class(self, class_info: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate a Python class."""
        class_name = class_info["name"]
        
        # Generate imports
        imports = []
        if any(word in class_info["description"].lower() for word in ["typing", "type"]):
            imports.append("from typing import Any, Optional, Dict, List")
        
        # Generate methods
        methods = []
        if class_info["methods"]:
            for method_name in class_info["methods"]:
                method_code = f'''    def {method_name.lower().replace(' ', '_')}(self) -> Any:
        """
        {method_name} method.
        """
        # TODO: Implement {method_name} logic
        pass'''
                methods.append(method_code)
        else:
            # Generate some default methods based on context
            methods.append('''    def process(self) -> Any:
        """
        Process method for main functionality.
        """
        # TODO: Implement processing logic
        pass''')
        
        # Build class code
        base_classes = ", ".join(class_info.get("base_classes", []))
        if base_classes:
            class_def = f"class {class_name}({base_classes}):"
        else:
            class_def = f"class {class_name}:"
        
        docstring = f'    """\n    {class_info["description"]}\n    """'
        
        init_method = '''    def __init__(self) -> None:
        """
        Initialize the class.
        """
        # TODO: Initialize attributes
        pass'''
        
        methods_code = "\n\n".join(methods)
        
        code = f"{class_def}\n{docstring}\n\n{init_method}\n\n{methods_code}"
        
        result = {
            "success": True,
            "language": "python",
            "code": code,
            "imports": imports,
            "class_name": class_name,
            "description": class_info["description"],
            "methods_generated": len(methods)
        }
        
        return result
    
    def _generate_tests(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate test code based on existing code or description."""
        # Try to find what to test
        test_target = self._extract_test_target(task)
        
        if test_target:
            # Find the target file/class/function
            target_analysis = self._analyze_test_target(test_target, context)
            if target_analysis:
                return self._generate_test_code(target_analysis, context)
        
        # Generic test generation
        return self._generate_generic_tests(task, context)
    
    def _generate_documentation(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate documentation for code or project."""
        doc_target = self._extract_documentation_target(task)
        
        if doc_target:
            # Generate documentation for specific file/module
            return self._generate_code_documentation(doc_target, context)
        else:
            # Generate project documentation
            return self._generate_project_documentation(context)
    
    def _suggest_refactoring(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Suggest refactoring improvements for code."""
        # This would analyze existing code and suggest improvements
        return {
            "success": True,
            "suggestions": [
                "Extract common functionality into utility functions",
                "Use type hints for better code clarity",
                "Consider using data classes for configuration",
                "Add error handling and logging",
                "Break down large functions into smaller ones"
            ],
            "priority": "medium",
            "estimated_effort": "2-4 hours"
        }
    
    def _generate_code(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate generic code based on task description."""
        primary_language = context["primary_language"]
        
        # Simple code generation based on keywords
        if primary_language == "python":
            return self._generate_python_snippet(task, context)
        else:
            return {"error": f"Generic code generation not supported for {primary_language}"}
    
    def _generate_python_snippet(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate a Python code snippet."""
        imports = []
        code_lines = []
        
        # Add imports based on task content
        task_lower = task.lower()
        
        if any(word in task_lower for word in ["file", "path", "directory"]):
            imports.append("from pathlib import Path")
        if any(word in task_lower for word in ["json", "data"]):
            imports.append("import json")
        if any(word in task_lower for word in ["datetime", "time", "date"]):
            imports.append("from datetime import datetime")
        if any(word in task_lower for word in ["requests", "http", "api"]):
            imports.append("import requests")
        
        # Generate basic code structure
        code_lines.append("# TODO: Implement the following logic:")
        code_lines.append(f"# {task}")
        code_lines.append("")
        
        if "function" in task_lower:
            code_lines.extend([
                "def main():",
                "    # Implementation here",
                "    pass",
                "",
                "if __name__ == '__main__':",
                "    main()"
            ])
        else:
            code_lines.append("# Implementation code here")
        
        imports_code = "\n".join(imports) if imports else ""
        main_code = "\n".join(code_lines)
        
        full_code = f"{imports_code}\n\n{main_code}" if imports else main_code
        
        return {
            "success": True,
            "language": "python",
            "code": full_code,
            "description": f"Generated code snippet for: {task}",
            "imports": imports
        }
    
    # Helper methods for extraction and analysis
    def _extract_test_target(self, task: str) -> Optional[str]:
        """Extract what should be tested from the task."""
        patterns = [
            r'test\s+(?:the\s+)?(\w+)\s+(?:class|function|module)',
            r'generate\s+tests?\s+for\s+(\w+)',
            r'create\s+tests?\s+for\s+(\w+)'
        ]
        
        for pattern in patterns:
            match = re.search(pattern, task, re.IGNORECASE)
            if match:
                return match.group(1)
        return None
    
    def _analyze_test_target(self, target: str, context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Analyze the target to be tested."""
        # This would analyze existing code to understand what to test
        return None  # Simplified for now
    
    def _generate_test_code(self, target_analysis: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate specific test code."""
        # This would generate tests based on the analysis
        return {"success": True, "code": "# Test code would go here"}
    
    def _generate_generic_tests(self, task: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate generic test structure."""
        return {
            "success": True,
            "language": "python",
            "code": '''import unittest

class TestGeneric(unittest.TestCase):
    """Generic test class."""
    
    def setUp(self):
        """Set up test fixtures."""
        pass
    
    def test_example(self):
        """Example test method."""
        # TODO: Implement test logic
        self.assertTrue(True)

if __name__ == '__main__':
    unittest.main()''',
            "description": "Generic test template"
        }
    
    def _extract_documentation_target(self, task: str) -> Optional[str]:
        """Extract documentation target from task."""
        patterns = [
            r'document\s+(?:the\s+)?(\w+(?:\.py)?)',
            r'documentation\s+for\s+(\w+)',
            r'docs\s+for\s+(\w+)'
        ]
        
        for pattern in patterns:
            match = re.search(pattern, task, re.IGNORECASE)
            if match:
                return match.group(1)
        return None
    
    def _generate_code_documentation(self, target: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate documentation for specific code."""
        return {
            "success": True,
            "documentation": f"# Documentation for {target}\n\nThis module provides functionality for {target}.",
            "format": "markdown"
        }
    
    def _generate_project_documentation(self, context: Dict[str, Any]) -> Dict[str, Any]:
        """Generate project-level documentation."""
        project_name = Path.cwd().name
        primary_lang = context.get("primary_language", "python")
        
        doc = f"""# {project_name}

## Overview
This is a {primary_lang} project with {len(context.get('languages', {}))} language(s).

## Structure
- Primary Language: {primary_lang}
- Total Files: {len(context.get('existing_functions', []))} functions found
- Classes: {len(context.get('existing_classes', []))} classes found

## Getting Started
[Add installation and usage instructions here]

## Contributing
[Add contribution guidelines here]
"""
        
        return {
            "success": True,
            "documentation": doc,
            "format": "markdown",
            "filename": "README.md"
        }
