src/tools/project_handler.py

"""Project handling utilities for Ada files and archives."""

import tempfile
import uuid
import os
import zipfile
import glob
import shutil
from src.agents.code_analyzer import CodeAnalyzerAgent
from src.agents.ada_converter import AdaConverterAgent
from src.agents.unit_test_generator import UnitTestGeneratorAgent

class ProjectHandler:
    """Handles Ada project operations including extraction, analysis, conversion, and download."""
    
    def __init__(self, model):
        """Initialize ProjectHandler with model and create agents."""
        self.model = model
        self.analyzer = CodeAnalyzerAgent(model)
        self.converter = AdaConverterAgent(model)
        self.test_generator = UnitTestGeneratorAgent(model)
    
    @staticmethod
    def cleanup_old_projects(target_dir: str | None = None) -> None:
        """
        Clean up old ada_project_* directories in the temp directory.
        
        Args:
            target_dir: Directory to clean up (defaults to system temp directory)
        """
        if target_dir is None:
            target_dir = tempfile.gettempdir()
        
        # Find all ada_project_* directories
        pattern = os.path.join(target_dir, "ada_project_*")
        old_projects = glob.glob(pattern)
        
        # Remove each directory
        for project_dir in old_projects:
            if os.path.isdir(project_dir):
                try:
                    shutil.rmtree(project_dir)
                except OSError:
                    # Ignore errors if directory can't be removed
                    pass

    def unzip_project(self, zip_file: str, target_dir: str | None = None) -> str:
        """
        Extract a zip file containing Ada project files.
        
        Args:
            zip_file: Path to the zip file
            target_dir: Directory to extract to (defaults to system temp directory)
            
        Returns:
            Path to the extracted directory (unique UUID-based name)
        """
        if target_dir is None:
            target_dir = tempfile.gettempdir()
        
        # Clean up old ada_project_* directories
        self.cleanup_old_projects(target_dir)
        
        # Generate unique directory name using UUID
        unique_dir_name = f"ada_project_{uuid.uuid4().hex}"
        extract_path = os.path.join(target_dir, unique_dir_name)
        
        # Create the extraction directory
        os.makedirs(extract_path, exist_ok=True)
        
        # Extract zip file to extract_path
        with zipfile.ZipFile(zip_file, 'r') as zip_ref:
            zip_ref.extractall(extract_path)
        
        return extract_path

    @staticmethod
    def build_directory_tree(directory_path: str) -> str:
        """
        Build a visual directory tree representation.
        
        Args:
            directory_path: Path to the directory to build tree for
            
        Returns:
            String representation of directory tree suitable for Gradio display
        """
        def _build_tree(path: str, prefix: str = "") -> str:
            """Recursively build tree structure."""
            items = []
            
            if not os.path.exists(path):
                return ""
                
            try:
                # Get all items in directory
                entries = sorted(os.listdir(path))
                
                for i, entry in enumerate(entries):
                    entry_path = os.path.join(path, entry)
                    is_last_item = i == len(entries) - 1
                    
                    # Choose appropriate tree symbols
                    current_prefix = "└── " if is_last_item else "├── "
                    next_prefix = prefix + ("    " if is_last_item else "│   ")
                    
                    if os.path.isdir(entry_path):
                        items.append(f"{prefix}{current_prefix}{entry}/")
                        # Recursively add subdirectory contents
                        subtree = _build_tree(entry_path, next_prefix)
                        if subtree:
                            items.append(subtree)
                    else:
                        items.append(f"{prefix}{current_prefix}{entry}")
                        
            except PermissionError:
                items.append(f"{prefix}[Permission Denied]")
                
            return "\n".join(items)
        
        # Start with root directory name
        root_name = os.path.basename(directory_path) or directory_path
        tree = f"{root_name}/\n"
        tree += _build_tree(directory_path)
        
        return tree

    @staticmethod
    def find_ada_files(extracted_path: str) -> list[str]:
        """Find all Ada files (.ads and .adb) in the extracted project."""
        ada_files = []
        for root, _, files in os.walk(extracted_path):
            for file in files:
                if file.endswith(('.ads', '.adb')):
                    ada_files.append(os.path.join(root, file))
        return ada_files

    def analyze_business_logic(self, ada_code: str) -> str:
        """Analyze Ada code to extract business logic."""
        return self.analyzer.extract_business_logic(ada_code)

    def convert_to_python(self, ada_code: str) -> str:
        """Convert Ada code to Python."""
        return self.converter.convert_to_python(ada_code)

    def generate_unit_tests(self, python_code: str) -> str:
        """Generate unit tests for Python code."""
        return self.test_generator.generate_unit_tests(python_code)

    def process_ada_file(self, ada_code: str) -> tuple[str, str, str]:
        """Process Ada code: analyze, convert to Python, and generate unit tests.
        
        Args:
            ada_code: The Ada source code to process
            
        Returns:
            Tuple of (analysis_result, python_code, unit_tests)
        """
        # Analyze the Ada code
        analysis_result = self.analyze_business_logic(ada_code)
        
        # Convert Ada code to Python
        python_code = self.convert_to_python(ada_code)
        
        # Generate unit tests for the converted Python code
        unit_tests = self.generate_unit_tests(python_code)
        
        return analysis_result, python_code, unit_tests

    def convert_and_download_project(self, extracted_path: str) -> str | None:
        """Convert all Ada files to Python and create downloadable zip with tests"""
        if not extracted_path:
            return None
        
        try:
            # Create temporary directory for converted files
            temp_dir = tempfile.mkdtemp()
            converted_root = os.path.join(temp_dir, "converted_project")
            tests_dir = os.path.join(converted_root, "tests")
            os.makedirs(converted_root, exist_ok=True)
            os.makedirs(tests_dir, exist_ok=True)
            
            # Find all Ada files (.ads and .adb)
            ada_files = self.find_ada_files(extracted_path)
            
            converted_files = []
            
            # Convert each Ada file to Python
            for ada_file in ada_files:
                try:
                    # Read Ada file content
                    with open(ada_file, 'r', encoding='utf-8') as f:
                        ada_code = f.read()
                    
                    # Convert to Python
                    python_code = self.convert_to_python(ada_code)
                    
                    # Calculate relative path to preserve directory structure
                    rel_path = os.path.relpath(ada_file, extracted_path)
                    python_filename = os.path.splitext(rel_path)[0] + '.py'
                    python_file_path = os.path.join(converted_root, python_filename)
                    
                    # Create directories if needed
                    os.makedirs(os.path.dirname(python_file_path), exist_ok=True)
                    
                    # Write Python file
                    with open(python_file_path, 'w', encoding='utf-8') as f:
                        f.write(python_code)
                    
                    converted_files.append((python_file_path, python_filename, python_code))
                    
                except Exception as e:
                    print(f"Error converting {ada_file}: {str(e)}")
                    continue
            
            # Generate unit tests for each converted Python file
            for python_file_path, python_filename, python_code in converted_files:
                try:
                    # Generate unit tests from Python code string
                    unit_tests = self.generate_unit_tests(python_code)
                    
                    # Create test file name
                    test_filename = f"test_{os.path.basename(python_filename)}"
                    test_file_path = os.path.join(tests_dir, test_filename)
                    
                    # Write test file
                    with open(test_file_path, 'w', encoding='utf-8') as f:
                        f.write(unit_tests)
                        
                except Exception as e:
                    print(f"Error generating tests for {python_filename}: {str(e)}")
                    continue
            
            # Create zip file
            zip_path = os.path.join(temp_dir, "converted_ada_project.zip")
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for root, _, files in os.walk(converted_root):
                    for file in files:
                        file_path = os.path.join(root, file)
                        arcname = os.path.relpath(file_path, converted_root)
                        zipf.write(file_path, arcname)
            
            return zip_path
            
        except Exception as e:
            print(f"Error in convert_and_download_project: {str(e)}")
            return None