metagpt_function_generator.py

#!/usr/bin/env python3
"""
MetaGPT-based Function Generator from Descriptions

This script generates Python functions based on AI-generated descriptions using MetaGPT framework,
maintaining compatibility with the original generate_functions_from_descriptions.py input/output format.
"""

import json
import sys
import time
import argparse
import re
import os
import random
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime
from tqdm import tqdm
from pathlib import Path
import asyncio
from dataclasses import dataclass
import yaml

# Add MetaGPT to Python path so it can find its own modules
current_dir = os.path.dirname(os.path.abspath(__file__))
metagpt_path = os.path.join(current_dir, 'baselines', 'MetaGPT')
if metagpt_path not in sys.path:
    sys.path.insert(0, metagpt_path)

# Set a default API key for MetaGPT if not already set
if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = 'dummy-key-for-testing'

from pydantic import BaseModel, Field

# Try to import MetaGPT components with selective error handling
METAGPT_AVAILABLE = False
try:
    # Temporarily mock problematic provider modules to allow core imports
    import sys
    

    
    # Test core imports first
    import metagpt
    
    # Import core classes individually with error handling
    try:
        from metagpt.actions.action import Action
        print("✅ Action import successful")
    except ImportError as e:
        print(f"⚠️  Could not import Action: {e}")
        raise
    
    try:
        from metagpt.roles.role import Role, RoleReactMode
        print("✅ Role imports successful")
    except ImportError as e:
        print(f"⚠️  Could not import Role: {e}")
        raise
    
    try:
        from metagpt.schema import Message
        print("✅ Message import successful")
    except ImportError as e:
        print(f"⚠️  Could not import Message: {e}")
        raise
    
    try:
        from metagpt.logs import logger
        print("✅ Logger import successful")
    except ImportError as e:
        print(f"⚠️  Could not import logger: {e}")
        raise
    
    try:
        from metagpt.context import Context
        print("✅ Context import successful")
    except ImportError as e:
        print(f"⚠️  Could not import Context: {e}")
        raise
        
    try:
        from metagpt.team import Team
        print("✅ Team import successful")
    except ImportError as e:
        print(f"⚠️  Could not import Team: {e}")
        raise
    
    METAGPT_AVAILABLE = True
    print("🎉 All MetaGPT core components loaded successfully!")
    
except Exception as e:
    print(f"⚠️  MetaGPT imports failed: {e}")
    print("⚠️  Falling back to basic implementation...")
    METAGPT_AVAILABLE = False
    
    # Create dummy classes for fallback
    class Action:
        def __init__(self, **kwargs):
            self.name = kwargs.get('name', 'DummyAction')
            self.llm = None
        
        async def _aask(self, prompt, **kwargs):
            return "# Dummy implementation - MetaGPT not available"
        
        async def run(self, *args, **kwargs):
            return await self._aask(args[0] if args else "")
    
    class Role:
        def __init__(self, **kwargs):
            self.name = kwargs.get('name', 'DummyRole')
            self.profile = kwargs.get('profile', 'DummyProfile')
            self.goal = kwargs.get('goal', 'Dummy goal')
            self.constraints = kwargs.get('constraints', 'No constraints')
            self.llm = None
            self.context = None
            self.stats = {}
        
        def set_actions(self, actions):
            self.actions = actions
        
        def _set_react_mode(self, **kwargs):
            pass
    
    class RoleReactModeValue:
        def __init__(self, value):
            self._value = value
        
        @property
        def value(self):
            return self._value
    
    class RoleReactMode:
        BY_ORDER = RoleReactModeValue("BY_ORDER")
    
    class Message:
        def __init__(self, content="", **kwargs):
            self.content = content
    
    class logger:
        @staticmethod
        def info(msg):
            print(f"INFO: {msg}")
        
        @staticmethod
        def error(msg):
            print(f"ERROR: {msg}")
        
        @staticmethod
        def warning(msg):
            print(f"WARNING: {msg}")
    
    class Context:
        def __init__(self):
            pass
    
    class Team:
        def __init__(self, **kwargs):
            pass


@dataclass
class FunctionRequest:
    """Data class for function generation request"""
    name: str
    signature: str
    description: str
    parameters: List[Dict[str, Any]] = None
    complexity_score: int = 0
    calls_made: List[str] = None
    variables_used: List[str] = None
    control_structures: List[str] = None
    file_path: str = ""
    line_start: int = 0
    
    def __post_init__(self):
        if self.parameters is None:
            self.parameters = []
        if self.calls_made is None:
            self.calls_made = []
        if self.variables_used is None:
            self.variables_used = []
        if self.control_structures is None:
            self.control_structures = []


class GenerateFunctionAction(Action):
    """Action for generating a single Python function from description"""
    
    name: str = "GenerateFunctionAction"
    
    PROMPT_TEMPLATE: str = """
You are an expert Python developer specialized in game development and code generation.
Generate a complete, functional Python function based on the provided description and metadata.

**FUNCTION TO IMPLEMENT:**
Name: {function_name}
Signature: {function_signature}

**FUNCTION DESCRIPTION:**
{description}

**FUNCTION METADATA:**
- Complexity Score: {complexity_score}
- Expected Control Structures: {control_structures}
- Expected Function Calls: {calls_made}
- Expected Variables: {variables_used}

**PARAMETERS:**
{parameters_info}

**IMPORTANT REQUIREMENTS:**
1. Use the exact signature provided: `{function_signature}`
2. Implement the behavior described in the function description
3. Follow Python coding best practices and PEP 8
4. Generate complete, functional Python code
5. Include proper error handling where appropriate
6. Include necessary imports at the top
7. Add comprehensive docstring with Args and Returns sections

**OUTPUT FORMAT:**
Return only the complete Python function implementation with any required imports:

```python
# Required imports (if any)
import ...

def {function_name}(parameters):
    \"\"\"
    Function description.
    
    Args:
        param: Parameter description
        
    Returns:
        Return description
    \"\"\"
    # Implementation
    pass
```

Generate the functional Python code:
"""

    async def run(self, function_request: FunctionRequest, top_k: int = 3) -> Tuple[List[str], int]:
        """Generate function code from request with top-k sampling
        
        Returns:
            Tuple[List[str], int]: (generated_codes, total_tokens_used)
        """
        
        if not METAGPT_AVAILABLE:
            # Fallback implementation - generate multiple variations for top-k
            base_code = f"""def {function_request.name}(parameters):
    \"\"\"
    {function_request.description}
    \"\"\"
    # This is a fallback implementation - MetaGPT not available
    # Please install and configure MetaGPT properly for full functionality
    pass"""
            
            # Create variations for top-k sampling in fallback mode
            variations = []
            for i in range(top_k):
                if i == 0:
                    variations.append(base_code)
                else:
                    # Add slight variations with different comments for top-k
                    variant = f"""def {function_request.name}(parameters):
    \"\"\"
    {function_request.description}
    \"\"\"
    # Fallback implementation variant {i+1} - MetaGPT not available
    # Variation: Different implementation approach {i+1}
    pass"""
                    variations.append(variant)
            
            # Estimate tokens for fallback mode (rough approximation: 1 token ≈ 4 characters)
            total_chars = len(function_request.description) + sum(len(code) for code in variations)
            estimated_tokens = total_chars // 4
            
            return variations, estimated_tokens
        
        # Build parameter information
        param_info = []
        for param in function_request.parameters:
            param_desc = f"- {param.get('name', 'unknown')}"
            if param.get('type_annotation'):
                param_desc += f": {param['type_annotation']}"
            if param.get('default'):
                param_desc += f" = {param['default']}"
            param_desc += f" ({param.get('kind', 'unknown')})"
            param_info.append(param_desc)
        
        param_text = "\n".join(param_info) if param_info else "No parameters"
        
        # Format the prompt
        prompt = self.PROMPT_TEMPLATE.format(
            function_name=function_request.name,
            function_signature=function_request.signature,
            description=function_request.description,
            complexity_score=function_request.complexity_score,
            control_structures=', '.join(function_request.control_structures) if function_request.control_structures else 'None',
            calls_made=', '.join(function_request.calls_made[:10]) if function_request.calls_made else 'None',
            variables_used=', '.join(function_request.variables_used[:10]) if function_request.variables_used else 'None',
            parameters_info=param_text
        )
        
        # Generate multiple function code candidates (top-k sampling)
        candidates, tokens_used = await self._generate_top_k_candidates(prompt, top_k)
        return candidates, tokens_used
    
    async def _generate_top_k_candidates(self, prompt: str, top_k: int) -> Tuple[List[str], int]:
        """Generate top-k function code candidates using different strategies based on model
        
        Returns:
            Tuple[List[str], int]: (candidates, total_tokens_used)
        """
        
        if not METAGPT_AVAILABLE or not hasattr(self, 'llm') or self.llm is None:
            # Fallback - use simple aask method
            response = await self._aask(prompt)
            cleaned_code = self.clean_generated_code(response)
            
            # Create variations for top-k
            candidates = [cleaned_code]
            total_tokens = 0
            
            # Estimate tokens for initial prompt and response
            total_tokens += self._estimate_tokens(prompt) + self._estimate_tokens(cleaned_code)
            
            for i in range(1, top_k):
                # Add prompt variations for diversity
                variant_prompt = prompt + f" Make it unique and creative (variant {i+1})."
                variant_response = await self._aask(variant_prompt)
                variant_code = self.clean_generated_code(variant_response)
                candidates.append(variant_code)
                
                # Estimate tokens for variant
                total_tokens += self._estimate_tokens(variant_prompt) + self._estimate_tokens(variant_code)
            
            return candidates, total_tokens
        
        # If we have access to the underlying LLM configuration, use model-specific approaches
        # For now, implement a basic multi-candidate generation
        candidates = []
        total_tokens = 0
        
        # Generate multiple candidates with different approaches
        prompt_variations = [
            prompt,
            prompt + " Make it unique and creative.",
            prompt + " Use different code structure.",
            prompt + " Focus on different development patterns.",
            prompt + " Be more original.",
            prompt + " Use alternative implementation approach.",
            prompt + " Optimize for readability.",
            prompt + " Focus on error handling.",
            prompt + " Use modern Python features.",
        ]
        
        for i in range(top_k):
            try:
                current_prompt = prompt_variations[i % len(prompt_variations)]
                response = await self._aask(current_prompt)
                cleaned_code = self.clean_generated_code(response)
                if cleaned_code and len(cleaned_code.strip()) > 20:
                    candidates.append(cleaned_code)
                    # Estimate tokens for this generation
                    total_tokens += self._estimate_tokens(current_prompt) + self._estimate_tokens(cleaned_code)
            except Exception as e:
                print(f"⚠️  Error generating candidate {i+1}: {e}")
                continue
        
        # Ensure we have at least one candidate
        if not candidates:
            # Fallback to simple generation
            response = await self._aask(prompt)
            cleaned_code = self.clean_generated_code(response)
            candidates = [cleaned_code]
            total_tokens += self._estimate_tokens(prompt) + self._estimate_tokens(cleaned_code)
        
        # Pad with duplicates if we don't have enough candidates (no additional tokens for duplicates)
        while len(candidates) < top_k:
            if candidates:
                candidates.append(candidates[0])
            else:
                break
        
        return candidates[:top_k], total_tokens
    
    def clean_generated_code(self, code: str) -> str:
        """Clean up generated code by removing markdown formatting"""
        # Remove markdown code blocks
        code = re.sub(r'^```python\s*\n', '', code, flags=re.MULTILINE)
        code = re.sub(r'^```\s*$', '', code, flags=re.MULTILINE)
        code = re.sub(r'^```.*\n', '', code, flags=re.MULTILINE)
        
        # Remove leading/trailing whitespace
        code = code.strip()
        
        return code
    
    def validate_python_syntax(self, code: str) -> bool:
        """Validate if the generated code has valid Python syntax"""
        try:
            compile(code, '<string>', 'exec')
            return True
        except SyntaxError:
            return False
        except Exception:
            # Other compilation errors (like undefined names) are OK for syntax validation
            return True
    
    def _estimate_tokens(self, text: str) -> int:
        """Estimate token count for given text
        
        Uses a simple approximation: 1 token ≈ 4 characters for English text
        This is a rough estimate commonly used in the industry.
        
        Args:
            text: Input text to estimate tokens for
            
        Returns:
            Estimated number of tokens
        """
        if not text:
            return 0
        
        # Simple token estimation: ~4 characters per token for English text
        # This accounts for spaces, punctuation, and typical word lengths
        return max(1, len(text) // 4)


class FunctionGeneratorRole(Role):
    """Role responsible for generating Python functions from descriptions"""
    
    name: str = "FunctionGenerator"
    profile: str = "Expert Python Function Generator"
    goal: str = "Generate high-quality, functional Python functions from AI-generated descriptions"
    constraints: str = "Must maintain exact function signatures and implement described functionality with proper error handling"
    
    def __init__(self, top_k: int = 3, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([GenerateFunctionAction])
        self._set_react_mode(react_mode=RoleReactMode.BY_ORDER.value)
        self.top_k = top_k
        
        # Statistics tracking
        self.stats = {
            'total_functions': 0,
            'successful_generations': 0,
            'failed_generations': 0,
            'syntax_valid_functions': 0,
            'functions_with_imports': 0,
            'processing_time': 0.0,
            'total_chars_generated': 0,
            'total_candidates_generated': 0,
            'total_tokens_used': 0
        }
    
    async def generate_single_function(self, function_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Generate a single function with top-k sampling and return result"""
        try:
            # Create function request
            request = FunctionRequest(
                name=function_data.get('name', 'unknown'),
                signature=function_data.get('signature', ''),
                description=function_data.get('ai_generated_description', ''),
                parameters=function_data.get('parameters', []),
                complexity_score=function_data.get('complexity_score', 0),
                calls_made=function_data.get('calls_made', []),
                variables_used=function_data.get('variables_used', []),
                control_structures=function_data.get('control_structures', []),
                file_path=function_data.get('file_path', ''),
                line_start=function_data.get('line_start', 0)
            )
            
            # Generate multiple function code candidates (top-k sampling)
            action = GenerateFunctionAction()
            action.llm = self.llm  # Use the role's LLM
            generated_codes, tokens_used = await action.run(request, top_k=self.top_k)
            
            if not generated_codes or all(not code or len(code) < 20 for code in generated_codes):
                raise ValueError("Generated codes too short or empty")
            
            # Process all candidates
            processed_codes = []
            syntax_valid_list = []
            has_imports_list = []
            code_length_list = []
            line_count_list = []
            
            candidates_with_scores = []
            
            for i, generated_code in enumerate(generated_codes):
                if not generated_code or len(generated_code.strip()) < 20:
                    continue
                    
                # Validate syntax
                syntax_valid = action.validate_python_syntax(generated_code)
                
                # Check for imports
                has_imports = 'import ' in generated_code
                
                # Calculate quality score (higher is better)
                quality_score = self._calculate_code_quality_score(
                    generated_code, syntax_valid, has_imports, request
                )
                
                candidates_with_scores.append({
                    'code': generated_code,
                    'syntax_valid': syntax_valid,
                    'has_imports': has_imports,
                    'quality_score': quality_score,
                    'code_length': len(generated_code),
                    'line_count': len(generated_code.split('\n'))
                })
            
            # Sort candidates by quality score (descending)
            candidates_with_scores.sort(key=lambda x: x['quality_score'], reverse=True)
            
            # Extract sorted data
            for candidate in candidates_with_scores:
                processed_codes.append(candidate['code'])
                syntax_valid_list.append(candidate['syntax_valid'])
                has_imports_list.append(candidate['has_imports'])
                code_length_list.append(candidate['code_length'])
                line_count_list.append(candidate['line_count'])
            
            # Update statistics
            self.stats['total_candidates_generated'] += len(candidates_with_scores)
            if any(syntax_valid_list):
                self.stats['syntax_valid_functions'] += 1
            if any(has_imports_list):
                self.stats['functions_with_imports'] += 1
            
            self.stats['successful_generations'] += 1
            self.stats['total_chars_generated'] += sum(code_length_list)
            self.stats['total_tokens_used'] += tokens_used
            
            # Create result with multiple candidates
            result = {
                'original_function': function_data,
                'generated_code': processed_codes,  # Multiple candidates
                'syntax_valid': syntax_valid_list,
                'has_imports': has_imports_list,
                'code_length': code_length_list,
                'line_count': line_count_list,
                'generation_timestamp': datetime.now().isoformat(),
                'model_used': 'metagpt',
                'used_repo_context': False,
                'used_scot': False,
                'raw_response': None,
                'top_k': self.top_k,
                'candidates_count': len(processed_codes),
                'tokens_used': tokens_used  # Add token usage for this function
            }
            
            return result
            
        except Exception as e:
            logger.error(f"Error generating function {function_data.get('name', 'unknown')}: {e}")
            self.stats['failed_generations'] += 1
            return None
    
    def _calculate_code_quality_score(self, code: str, syntax_valid: bool, has_imports: bool, request: FunctionRequest) -> float:
        """Calculate a quality score for generated code to enable proper ranking"""
        score = 0.0
        
        # Base score for syntax validity
        if syntax_valid:
            score += 10.0
        else:
            score -= 5.0
        
        # Code length scoring (prefer reasonable length, penalize too short or too long)
        code_length = len(code)
        if 50 < code_length < 2000:
            score += 5.0
        elif code_length < 50:
            score -= 3.0
        elif code_length > 3000:
            score -= 2.0
        
        # Line count scoring
        line_count = len(code.split('\n'))
        if 5 < line_count < 100:
            score += 2.0
        
        # Docstring presence
        if '"""' in code or "'''" in code:
            score += 3.0
        
        # Has proper function signature
        if f"def {request.name}" in code:
            score += 5.0
        
        # Import usage (appropriate use of imports is good)
        if has_imports:
            score += 2.0
        
        # Check for common good practices
        if 'return ' in code:
            score += 2.0
        if 'if ' in code or 'for ' in code or 'while ' in code:
            score += 1.0  # Has control structures
        if 'try:' in code and 'except' in code:
            score += 2.0  # Has error handling
        
        # Penalize obvious placeholders/incomplete implementations
        if 'pass' in code.lower():
            score -= 2.0
        if 'todo' in code.lower() or 'fixme' in code.lower():
            score -= 1.0
        if 'not implemented' in code.lower():
            score -= 3.0
        
        # Bonus for complexity matching
        expected_complexity = request.complexity_score
        if expected_complexity > 0:
            # Rough estimation based on control structures and function calls
            estimated_complexity = code.count('if ') + code.count('for ') + code.count('while ') + code.count('(')
            if abs(estimated_complexity - expected_complexity) <= 2:
                score += 1.0
        
        return score
    
    async def process_functions_batch(self, functions_batch: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Process a batch of functions"""
        results = []
        
        for func_data in functions_batch:
            result = await self.generate_single_function(func_data)
            if result:
                results.append(result)
        
        return results


class MetaGPTFunctionGenerator:
    """Main class for MetaGPT-based function generation"""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None, config_file: Optional[str] = None, top_k: int = 3):
        self.config = self.load_config(config, config_file)
        self.context = Context()
        self.generator_role = None
        self.top_k = top_k
        
        # Statistics
        self.start_time = None
        self.end_time = None
    
    def load_config(self, config: Optional[Dict[str, Any]], config_file: Optional[str]) -> Dict[str, Any]:
        """Load configuration from dict or file"""
        if config:
            return config
        
        if config_file and os.path.exists(config_file):
            try:
                with open(config_file, 'r', encoding='utf-8') as f:
                    return yaml.safe_load(f)
            except Exception as e:
                logger.warning(f"Failed to load config file {config_file}: {e}")
        
        # Return default config
        return {
            'generation': {
                'batch_size': 5,
                'max_functions': None,
                'enable_syntax_validation': True,
                'include_imports': True
            },
            'performance': {
                'delay_between_batches': 0.1,
                'max_retries': 3
            },
            'io': {
                'input_file': 'Jsons/extracted_functions_with_comments.json',
                'output_file': 'Jsons/metagpt_patches.json'
            }
        }
    
    def setup_generator_role(self):
        """Setup the function generator role"""
        self.generator_role = FunctionGeneratorRole(top_k=self.top_k)
        self.generator_role.context = self.context
        
        # Log MetaGPT availability status
        if METAGPT_AVAILABLE:
            logger.info("✅ MetaGPT framework loaded successfully")
        else:
            logger.warning("⚠️  MetaGPT not available - using fallback implementation")
        
        logger.info(f"🎯 Top-k sampling enabled: generating {self.top_k} candidates per function")
        
        # Configure LLM if provided in config
        if 'llm_config' in self.config:
            # This would need to be configured based on your LLM setup
            pass
    
    async def process_functions_file(self, 
                                   input_file: str = "Jsons/extracted_functions_with_comments.json",
                                   output_file: str = "Jsons/patches.json",
                                   max_functions: Optional[int] = None,
                                   batch_size: int = 5) -> Dict[str, Any]:
        """Process all functions in the enhanced functions file"""
        
        logger.info(f"🚀 Starting MetaGPT function generation...")
        logger.info(f"📂 Input file: {input_file}")
        logger.info(f"📄 Output file: {output_file}")
        
        self.start_time = time.time()
        
        # Setup role
        if not self.generator_role:
            self.setup_generator_role()
        
        # Load functions data
        try:
            with open(input_file, 'r', encoding='utf-8') as f:
                functions_data = json.load(f)
            logger.info(f"✅ Loaded function data from {input_file}")
        except Exception as e:
            logger.error(f"❌ Error loading functions file: {e}")
            return {}
        
        # Extract functions with descriptions
        enhanced_functions = functions_data.get('enhanced_functions', [])
        
        # Filter functions with AI descriptions
        functions_with_descriptions = [
            func for func in enhanced_functions 
            if func.get('ai_generated_description') and len(func.get('ai_generated_description', '')) > 50
        ]
        
        if max_functions:
            functions_with_descriptions = functions_with_descriptions[:max_functions]
        
        self.generator_role.stats['total_functions'] = len(functions_with_descriptions)
        logger.info(f"📊 Processing {len(functions_with_descriptions)} functions with descriptions...")
        
        # Process functions in batches
        generated_patches = []
        
        for i in tqdm(range(0, len(functions_with_descriptions), batch_size), desc="Generating functions"):
            batch = functions_with_descriptions[i:i + batch_size]
            batch_results = await self.generator_role.process_functions_batch(batch)
            generated_patches.extend(batch_results)
            
            # Add small delay between batches to avoid overwhelming the API
            if i + batch_size < len(functions_with_descriptions):
                await asyncio.sleep(0.1)
        
        # Calculate processing time
        processing_time = time.time() - self.start_time
        self.generator_role.stats['processing_time'] = processing_time
        
        # Compile results
        patches_data = {
            'generation_metadata': {
                'timestamp': datetime.now().isoformat(),
                'input_file': input_file,
                'model_used': 'metagpt' if METAGPT_AVAILABLE else 'fallback',
                'framework': 'MetaGPT' if METAGPT_AVAILABLE else 'MetaGPT-Fallback',
                'metagpt_available': METAGPT_AVAILABLE,
                'base_url': 'local',
                'repo_path': None,
                'hcp_integration_enabled': False,
                'tree_sitter_available': False,
                'scot_prompting_enabled': False,
                'total_functions_processed': len(functions_with_descriptions),
                'successful_generations': self.generator_role.stats['successful_generations'],
                'failed_generations': self.generator_role.stats['failed_generations'],
                'syntax_valid_functions': self.generator_role.stats['syntax_valid_functions'],
                'functions_with_imports': self.generator_role.stats['functions_with_imports'],
                'scot_responses': 0,
                'api_calls_made': self.generator_role.stats['successful_generations'] + self.generator_role.stats['failed_generations'],
                'total_tokens_used': self.generator_role.stats['total_tokens_used'],  # Total tokens used across all functions
                'processing_time_seconds': processing_time,
                'success_rate': (self.generator_role.stats['successful_generations'] / len(functions_with_descriptions) * 100) if functions_with_descriptions else 0,
                'syntax_valid_rate': (self.generator_role.stats['syntax_valid_functions'] / self.generator_role.stats['successful_generations'] * 100) if self.generator_role.stats['successful_generations'] > 0 else 0
            },
            'original_metadata': {
                'source_extraction': functions_data.get('original_metadata', {}),
                'description_generation': functions_data.get('generation_metadata', {})
            },
            'generated_patches': generated_patches,
            'statistics': {
                'total_patches': len(generated_patches),
                'average_code_length': sum(p['code_length'][0] for p in generated_patches) / len(generated_patches) if generated_patches else 0,
                'average_line_count': sum(p['line_count'][0] for p in generated_patches) / len(generated_patches) if generated_patches else 0,
                'syntax_error_rate': ((self.generator_role.stats['successful_generations'] - self.generator_role.stats['syntax_valid_functions']) / self.generator_role.stats['successful_generations'] * 100) if self.generator_role.stats['successful_generations'] > 0 else 0,
                'functions_with_repo_context': 0,
                'functions_with_scot': 0,
                'total_chars_generated': self.generator_role.stats['total_chars_generated']
            }
        }
        
        # Save patches data
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(patches_data, f, indent=2, ensure_ascii=False)
            logger.info(f"✅ Generated function patches saved to {output_file}")
        except Exception as e:
            logger.error(f"❌ Error saving patches data: {e}")
        
        # Print summary
        self.print_generation_summary()
        
        return patches_data
    
    def print_generation_summary(self):
        """Print comprehensive generation summary"""
        if not self.generator_role:
            return
            
        stats = self.generator_role.stats
        
        print("\n" + "="*80)
        print("📊 METAGPT FUNCTION GENERATION SUMMARY")
        print("="*80)
        
        print(f"\n📈 PROCESSING STATISTICS:")
        print(f"Total Functions: {stats['total_functions']}")
        print(f"Successful Generations: {stats['successful_generations']}")
        print(f"Failed Generations: {stats['failed_generations']}")
        print(f"Success Rate: {(stats['successful_generations'] / stats['total_functions'] * 100) if stats['total_functions'] > 0 else 0:.1f}%")
        
        print(f"\n🔧 CODE QUALITY:")
        print(f"Syntax Valid Functions: {stats['syntax_valid_functions']}")
        print(f"Functions with Imports: {stats['functions_with_imports']}")
        if stats['successful_generations'] > 0:
            print(f"Syntax Valid Rate: {(stats['syntax_valid_functions'] / stats['successful_generations'] * 100):.1f}%")
            print(f"Import Usage Rate: {(stats['functions_with_imports'] / stats['successful_generations'] * 100):.1f}%")
        
        print(f"\n🎯 TOP-K SAMPLING:")
        print(f"Top-K Value: {self.top_k}")
        print(f"Total Candidates Generated: {stats['total_candidates_generated']:,}")
        if stats['successful_generations'] > 0:
            avg_candidates_per_function = stats['total_candidates_generated'] / stats['successful_generations']
            print(f"Average Candidates per Function: {avg_candidates_per_function:.1f}")
        
        print(f"\n⏱️ PERFORMANCE:")
        print(f"Processing Time: {stats['processing_time']:.1f} seconds")
        print(f"Total Characters Generated: {stats['total_chars_generated']:,}")
        print(f"Total Tokens Used: {stats['total_tokens_used']:,}")
        
        if stats['successful_generations'] > 0:
            avg_chars_per_function = stats['total_chars_generated'] / stats['successful_generations']
            avg_tokens_per_function = stats['total_tokens_used'] / stats['successful_generations']
            print(f"Average Characters per Function: {avg_chars_per_function:.1f}")
            print(f"Average Tokens per Function: {avg_tokens_per_function:.1f}")
        
        print(f"\n⚙️  FRAMEWORK:")
        print(f"Framework: MetaGPT")
        print(f"Generation Mode: Single Role with Action + Top-K Sampling")
        
        print("="*80)


async def main_async():
    """Async main function"""
    parser = argparse.ArgumentParser(description="Generate Python functions from AI descriptions using MetaGPT with top-k sampling")
    parser.add_argument("--input-file", 
                       help="Input JSON file with function descriptions")
    parser.add_argument("--output-file", 
                       help="Output JSON file for generated function patches")
    parser.add_argument("--config-file", default="metagpt_config.yaml",
                       help="Configuration file (YAML/config format)")
    parser.add_argument("--max-functions", type=int,
                       help="Maximum number of functions to process")
    parser.add_argument("--batch-size", type=int,
                       help="Number of functions to process in each batch")
    parser.add_argument("--top-k", type=int, default=3,
                       help="Number of candidate functions to generate per request (top-k sampling)")
    parser.add_argument("--analyze", action="store_true",
                       help="Analyze existing generated patches")
    parser.add_argument("--sample", type=int, default=3,
                       help="Number of sample functions to display")
    
    args = parser.parse_args()
    
    if args.analyze:
        # TODO: Implement analysis functionality
        print("Analysis functionality not implemented yet")
        return
    
    # Create generator with configuration and top-k parameter
    generator = MetaGPTFunctionGenerator(config_file=args.config_file, top_k=args.top_k)
    
    # Use command line args to override config
    config = generator.config
    input_file = args.input_file or config.get('io', {}).get('input_file', 'Jsons/extracted_functions_with_comments.json')
    output_file = args.output_file or config.get('io', {}).get('output_file', 'Jsons/metagpt_patches.json')
    max_functions = args.max_functions or config.get('generation', {}).get('max_functions')
    batch_size = args.batch_size or config.get('generation', {}).get('batch_size', 5)
    
    result = await generator.process_functions_file(
        input_file=input_file,
        output_file=output_file,
        max_functions=max_functions,
        batch_size=batch_size
    )
    
    print(f"\n✅ MetaGPT function generation complete!")
    print(f"📄 Generated patches saved to: {output_file}")


def main():
    """Main function"""
    asyncio.run(main_async())


if __name__ == "__main__":
    main()