codesmith_engine.py

"""
Enhanced CodeSmith Engine - Autonomous Multi-File Code Generation
Handles vast, complex projects with self-validation and correction
"""
import os
import json
import asyncio
from typing import Dict, List, Optional, Any, Tuple
from hybrid_core import HybridCore, Field

class ProjectPlanner:
    """Plans complex projects by breaking them into tasks with dependencies"""
    
    def __init__(self, hybrid_core: HybridCore):
        self.hybrid_core = hybrid_core
    
    async def analyze_and_plan(self, prompt: str) -> Dict[str, Any]:
        """
        Analyze project requirements and create detailed execution plan
        
        Returns:
            {
                "project_type": str,
                "tech_stack": list,
                "files": list of dicts with {name, purpose, dependencies},
                "execution_order": list of file names in order,
                "complexity": str (simple/medium/complex/vast)
            }
        """
        analysis_prompt = f"""Analyze this project request and create a comprehensive plan:

Request: {prompt}

Provide a detailed JSON response with:
1. project_type: (web_app/api/fullstack/mobile/desktop/game/other)
2. tech_stack: list of technologies needed
3. files: array of objects with:
   - name: filename with path
   - purpose: what this file does
   - dependencies: list of other files it depends on
4. database_needed: boolean
5. api_endpoints: list if applicable
6. complexity: (simple/medium/complex/vast)

Think through the entire architecture. For complex projects, include:
- Configuration files (package.json, requirements.txt, .env.example)
- Database models/schemas
- API routes/controllers
- Frontend components
- Utility functions
- Tests (if applicable)

Return ONLY valid JSON, no explanations."""

        response = await self.hybrid_core.call_llm("openai", "gpt-4o", analysis_prompt, max_tokens=3000)
        
        try:
            plan = json.loads(response)
            # Calculate execution order based on dependencies
            plan["execution_order"] = self._calculate_execution_order(plan.get("files", []))
            return plan
        except:
            # Fallback plan
            return {
                "project_type": "simple",
                "tech_stack": ["python"],
                "files": [{"name": "main.py", "purpose": "Main file", "dependencies": []}],
                "execution_order": ["main.py"],
                "complexity": "simple"
            }
    
    def _calculate_execution_order(self, files: List[Dict]) -> List[str]:
        """Calculate optimal file generation order based on dependencies"""
        # Create dependency graph
        graph = {f["name"]: f.get("dependencies", []) for f in files}
        
        # Topological sort
        visited = set()
        order = []
        
        def visit(node):
            if node in visited:
                return
            visited.add(node)
            for dep in graph.get(node, []):
                if dep in graph:
                    visit(dep)
            order.append(node)
        
        for file in graph:
            visit(file)
        
        return order

class CodeValidator:
    """Validates generated code for syntax errors and common issues"""
    
    def __init__(self, hybrid_core: HybridCore):
        self.hybrid_core = hybrid_core
    
    async def validate_code(self, filename: str, code: str) -> Dict[str, Any]:
        """
        Validate code for syntax errors and issues
        
        Returns:
            {
                "valid": bool,
                "errors": list of error messages,
                "warnings": list of warnings,
                "suggestions": list of improvement suggestions
            }
        """
        # Basic syntax validation
        errors = []
        warnings = []
        
        # Check for common issues
        if "TODO" in code or "PLACEHOLDER" in code or "..." in code:
            errors.append("Code contains placeholders (TODO, PLACEHOLDER, or ...)")
        
        if code.strip() == "":
            errors.append("Code is empty")
        
        # Language-specific validation
        if filename.endswith('.py'):
            errors.extend(self._validate_python(code))
        elif filename.endswith(('.js', '.jsx', '.ts', '.tsx')):
            errors.extend(self._validate_javascript(code))
        
        # AI-powered validation
        if errors or len(code) > 100:
            ai_validation = await self._ai_validate(filename, code)
            errors.extend(ai_validation.get("errors", []))
            warnings.extend(ai_validation.get("warnings", []))
        
        return {
            "valid": len(errors) == 0,
            "errors": errors,
            "warnings": warnings,
            "suggestions": []
        }
    
    def _validate_python(self, code: str) -> List[str]:
        """Basic Python syntax validation"""
        errors = []
        try:
            compile(code, '<string>', 'exec')
        except SyntaxError as e:
            errors.append(f"Python syntax error: {e.msg} at line {e.lineno}")
        return errors
    
    def _validate_javascript(self, code: str) -> List[str]:
        """Basic JavaScript validation"""
        errors = []
        # Check for common issues
        if code.count('{') != code.count('}'):
            errors.append("Mismatched curly braces")
        if code.count('(') != code.count(')'):
            errors.append("Mismatched parentheses")
        if code.count('[') != code.count(']'):
            errors.append("Mismatched square brackets")
        return errors
    
    async def _ai_validate(self, filename: str, code: str) -> Dict:
        """Use AI to validate code"""
        prompt = f"""Validate this code for {filename}:

```
{code[:2000]}  # Limit to first 2000 chars
```

Check for:
1. Syntax errors
2. Missing imports
3. Undefined variables
4. Logic errors
5. Security issues

Return JSON:
{{
    "errors": ["list of critical errors"],
    "warnings": ["list of warnings"]
}}"""

        try:
            response = await self.hybrid_core.call_llm("openai", "gpt-4o-mini", prompt, max_tokens=500)
            return json.loads(response)
        except:
            return {"errors": [], "warnings": []}

class AutonomousExecutor:
    """Executes code generation with self-correction and iteration"""
    
    def __init__(self, hybrid_core: HybridCore, validator: CodeValidator):
        self.hybrid_core = hybrid_core
        self.validator = validator
        self.max_retries = 3
    
    async def generate_file_with_validation(
        self, 
        filename: str, 
        purpose: str, 
        context: str,
        dependencies: List[str] = None
    ) -> Tuple[str, Dict]:
        """
        Generate file with automatic validation and correction
        
        Returns:
            (code, metadata) tuple
        """
        attempt = 0
        last_errors = []
        
        while attempt < self.max_retries:
            # Generate code
            code = await self._generate_code(filename, purpose, context, dependencies, last_errors)
            
            # Validate
            validation = await self.validator.validate_code(filename, code)
            
            if validation["valid"]:
                return code, {
                    "attempts": attempt + 1,
                    "status": "success",
                    "warnings": validation["warnings"]
                }
            
            # Store errors for next attempt
            last_errors = validation["errors"]
            attempt += 1
            
            if attempt < self.max_retries:
                # Try to fix errors
                code = await self._fix_errors(filename, code, validation["errors"])
                # Validate fixed code
                validation = await self.validator.validate_code(filename, code)
                if validation["valid"]:
                    return code, {
                        "attempts": attempt + 1,
                        "status": "success_after_fix",
                        "warnings": validation["warnings"]
                    }
        
        # Return best attempt even if not perfect
        return code, {
            "attempts": self.max_retries,
            "status": "partial",
            "errors": last_errors
        }
    
    async def _generate_code(
        self, 
        filename: str, 
        purpose: str, 
        context: str,
        dependencies: List[str],
        previous_errors: List[str]
    ) -> str:
        """Generate code for a file"""
        deps_context = ""
        if dependencies:
            deps_context = f"\nThis file depends on: {', '.join(dependencies)}"
        
        error_context = ""
        if previous_errors:
            error_context = f"\n\nPrevious attempt had these errors:\n" + "\n".join(f"- {e}" for e in previous_errors)
            error_context += "\n\nFix these errors in the new version."
        
        prompt = f"""Generate complete, production-ready code for this file:

Filename: {filename}
Purpose: {purpose}
Project Context: {context}{deps_context}{error_context}

CRITICAL REQUIREMENTS:
1. Write COMPLETE, WORKING code - NO placeholders, NO TODOs, NO "..."
2. Include ALL necessary imports
3. Implement ALL functions fully
4. Add proper error handling
5. Include comments for complex logic
6. Follow best practices for the language
7. Ensure code is syntactically correct
8. Make it production-ready

Return ONLY the code, no explanations, no markdown formatting."""

        code = await self.hybrid_core.call_llm("openai", "gpt-4o", prompt, max_tokens=4000)
        
        # Clean up markdown code blocks if present
        if "```" in code:
            lines = code.split("\n")
            code_lines = []
            in_code_block = False
            for line in lines:
                if line.strip().startswith("```"):
                    in_code_block = not in_code_block
                    continue
                if in_code_block or not any(line.strip().startswith(x) for x in ["```", "#", "Here", "This"]):
                    code_lines.append(line)
            code = "\n".join(code_lines)
        
        return code.strip()
    
    async def _fix_errors(self, filename: str, code: str, errors: List[str]) -> str:
        """Attempt to fix errors in code"""
        prompt = f"""Fix the errors in this code for {filename}:

```
{code}
```

Errors to fix:
{chr(10).join(f"- {e}" for e in errors)}

Return the corrected code ONLY, no explanations."""

        fixed_code = await self.hybrid_core.call_llm("openai", "gpt-4o", prompt, max_tokens=4000)
        
        # Clean up markdown
        if "```" in fixed_code:
            lines = fixed_code.split("\n")
            code_lines = []
            in_code_block = False
            for line in lines:
                if line.strip().startswith("```"):
                    in_code_block = not in_code_block
                    continue
                if in_code_block:
                    code_lines.append(line)
            fixed_code = "\n".join(code_lines)
        
        return fixed_code.strip()

class CodeSmithEngine:
    """Enhanced autonomous code generation engine"""
    
    def __init__(self):
        self.hybrid_core = HybridCore()
        self.planner = ProjectPlanner(self.hybrid_core)
        self.validator = CodeValidator(self.hybrid_core)
        self.executor = AutonomousExecutor(self.hybrid_core, self.validator)
    
    async def generate_project(self, prompt: str, progress_callback=None) -> Dict[str, Any]:
        """
        Generate a complete project autonomously
        
        Args:
            prompt: Project description
            progress_callback: Optional callback function(step, total, message)
        
        Returns:
            {
                "plan": project plan,
                "files": dict of filename -> code,
                "metadata": generation metadata,
                "project_structure": file tree
            }
        """
        # Phase 1: Planning
        if progress_callback:
            progress_callback(1, 5, "Analyzing project requirements...")
        
        plan = await self.planner.analyze_and_plan(prompt)
        
        # Phase 2: Generate files in dependency order
        if progress_callback:
            progress_callback(2, 5, "Generating project files...")
        
        files = {}
        metadata = {}
        context = f"Project: {prompt}\nType: {plan['project_type']}\nTech Stack: {', '.join(plan['tech_stack'])}"
        
        total_files = len(plan["execution_order"])
        for idx, filename in enumerate(plan["execution_order"]):
            if progress_callback:
                progress_callback(2, 5, f"Generating {filename} ({idx+1}/{total_files})...")
            
            # Find file info
            file_info = next((f for f in plan["files"] if f["name"] == filename), None)
            if not file_info:
                continue
            
            # Generate with validation
            code, file_metadata = await self.executor.generate_file_with_validation(
                filename,
                file_info["purpose"],
                context,
                file_info.get("dependencies", [])
            )
            
            files[filename] = code
            metadata[filename] = file_metadata
        
        # Phase 3: Generate integration files if needed
        if progress_callback:
            progress_callback(3, 5, "Generating configuration files...")
        
        if plan["project_type"] in ["fullstack", "web_app", "api"]:
            config_files = await self._generate_config_files(plan, files)
            files.update(config_files)
        
        # Phase 4: Validate integration
        if progress_callback:
            progress_callback(4, 5, "Validating project integration...")
        
        integration_report = await self._validate_integration(files, plan)
        
        # Phase 5: Generate project structure
        if progress_callback:
            progress_callback(5, 5, "Finalizing project...")
        
        structure = self._generate_structure(list(files.keys()))
        
        return {
            "plan": plan,
            "files": files,
            "metadata": metadata,
            "integration_report": integration_report,
            "project_structure": structure
        }
    
    async def _generate_config_files(self, plan: Dict, existing_files: Dict) -> Dict[str, str]:
        """Generate configuration files based on project type"""
        config_files = {}
        
        # Generate package.json for Node.js projects
        if any(tech in plan["tech_stack"] for tech in ["react", "node", "express", "next"]):
            package_json = await self._generate_package_json(plan, existing_files)
            config_files["package.json"] = package_json
        
        # Generate requirements.txt for Python projects
        if any(tech in plan["tech_stack"] for tech in ["python", "fastapi", "flask", "django"]):
            requirements = await self._generate_requirements(plan, existing_files)
            config_files["requirements.txt"] = requirements
        
        # Generate .env.example
        env_example = await self._generate_env_example(plan)
        config_files[".env.example"] = env_example
        
        # Generate README.md
        readme = await self._generate_readme(plan, existing_files)
        config_files["README.md"] = readme
        
        return config_files
    
    async def _generate_package_json(self, plan: Dict, files: Dict) -> str:
        """Generate package.json"""
        prompt = f"""Generate a complete package.json for this project:

Project Type: {plan['project_type']}
Tech Stack: {', '.join(plan['tech_stack'])}
Files: {', '.join(files.keys())}

Include:
- Appropriate dependencies
- Dev dependencies
- Scripts (start, build, test)
- Proper project metadata

Return ONLY valid JSON."""

        response = await self.hybrid_core.call_llm("openai", "gpt-4o-mini", prompt, max_tokens=1000)
        return response
    
    async def _generate_requirements(self, plan: Dict, files: Dict) -> str:
        """Generate requirements.txt"""
        prompt = f"""Generate requirements.txt for this Python project:

Project Type: {plan['project_type']}
Tech Stack: {', '.join(plan['tech_stack'])}

List all required packages with versions.
Return ONLY the requirements list, one per line."""

        response = await self.hybrid_core.call_llm("openai", "gpt-4o-mini", prompt, max_tokens=500)
        return response
    
    async def _generate_env_example(self, plan: Dict) -> str:
        """Generate .env.example"""
        env_vars = []
        
        if "database" in str(plan).lower():
            env_vars.append("DATABASE_URL=your_database_url_here")
        
        if "api" in plan["project_type"]:
            env_vars.append("API_KEY=your_api_key_here")
            env_vars.append("SECRET_KEY=your_secret_key_here")
        
        return "\n".join(env_vars)
    
    async def _generate_readme(self, plan: Dict, files: Dict) -> str:
        """Generate README.md"""
        prompt = f"""Generate a comprehensive README.md for this project:

Project Type: {plan['project_type']}
Tech Stack: {', '.join(plan['tech_stack'])}
Files: {len(files)} files

Include:
1. Project title and description
2. Features
3. Installation instructions
4. Usage examples
5. File structure
6. Contributing guidelines

Return markdown format."""

        response = await self.hybrid_core.call_llm("openai", "gpt-4o-mini", prompt, max_tokens=1500)
        return response
    
    async def _validate_integration(self, files: Dict, plan: Dict) -> Dict:
        """Validate that all files work together"""
        issues = []
        
        # Check for missing dependencies
        for filename, code in files.items():
            # Extract imports (simple check)
            if filename.endswith('.py'):
                imports = [line for line in code.split('\n') if line.strip().startswith(('import ', 'from '))]
                # Check if imported modules exist as files
                for imp in imports:
                    if 'from .' in imp or 'from ..' in imp:
                        # Local import
                        module = imp.split('import')[0].replace('from', '').replace('.', '').strip()
                        if not any(module in f for f in files.keys()):
                            issues.append(f"{filename}: Missing local module {module}")
        
        return {
            "valid": len(issues) == 0,
            "issues": issues
        }
    
    def _generate_structure(self, files: List[str]) -> Dict[str, Any]:
        """Generate project structure tree"""
        structure = {"type": "directory", "name": "project", "children": []}
        
        for filepath in sorted(files):
            parts = filepath.split("/")
            current = structure
            
            for i, part in enumerate(parts):
                if i == len(parts) - 1:
                    # It's a file
                    current["children"].append({"type": "file", "name": part})
                else:
                    # It's a directory
                    existing = next((c for c in current.get("children", []) if c["name"] == part), None)
                    if not existing:
                        new_dir = {"type": "directory", "name": part, "children": []}
                        current["children"].append(new_dir)
                        current = new_dir
                    else:
                        current = existing
        
        return structure
    
    # Keep existing methods for backward compatibility
    async def explain_code(self, code: str) -> str:
        """Explain what a piece of code does"""
        prompt = f"""Explain this code in simple terms:

```
{code}
```

Provide:
1. What it does (2-3 sentences)
2. Key components
3. How it works"""

        explanation = await self.hybrid_core.call_llm("openai", "gpt-4o-mini", prompt, max_tokens=1000)
        return explanation
    
    async def optimize_code(self, code: str) -> Dict[str, str]:
        """Optimize code for performance and readability"""
        prompt = f"""Optimize this code:

```
{code}
```

Provide:
1. Optimized version
2. List of improvements made

Format:
OPTIMIZED CODE:
[code here]

IMPROVEMENTS:
- [improvement 1]
- [improvement 2]"""

        response = await self.hybrid_core.call_llm("openai", "gpt-4o", prompt, max_tokens=2000)
        
        # Parse response
        parts = response.split("IMPROVEMENTS:")
        optimized = parts[0].replace("OPTIMIZED CODE:", "").strip()
        if "```" in optimized:
            optimized = optimized.split("```")[1].strip()
        improvements = parts[1].strip() if len(parts) > 1 else "Various optimizations applied"
        
        return {
            "optimized_code": optimized,
            "improvements": improvements
        }
    
    async def debug_code(self, code: str, error: str = "") -> Dict[str, str]:
        """Debug code and suggest fixes"""
        prompt = f"""Debug this code:

```
{code}
```

Error (if any): {error}

Provide:
1. Identified issues
2. Fixed code
3. Explanation of fixes"""

        response = await self.hybrid_core.call_llm("openai", "gpt-4o", prompt, max_tokens=2000)
        
        return {
            "debug_response": response,
            "status": "analyzed"
        }