agent_core.py

import os
import uuid
from datetime import datetime
from typing import Dict, Optional
from github_client import GitHubClient
from gemini_ai import analyze_issue, generate_fix_plan, generate_code_patch
from validators import CodeValidator, check_safety_limits
from storage import StorageManager
import traceback


class CodeSentinelAgent:
    """Main agent that orchestrates issue-to-PR workflow."""
    
    def __init__(self, github_token: Optional[str] = None):
        self.github_client = GitHubClient(github_token)
        self.validator = CodeValidator()
        self.storage = StorageManager()
    
    def process_issue_to_pr(self, repo_url: str, issue_text: str, 
                           base_branch: str = "main",
                           max_files: int = 10,
                           max_lines: int = 500) -> Dict:
        """
        Main workflow: Issue → Analysis → Plan → Patch → Validate → PR
        """
        run_id = str(uuid.uuid4())[:8]
        run_data = {
            'run_id': run_id,
            'timestamp': datetime.now().isoformat(),
            'repo_url': repo_url,
            'issue_text': issue_text,
            'status': 'started',
            'steps': []
        }
        
        temp_dir = None
        fork_info = None
        
        try:
            # Step 0: Fork repository
            self._add_step(run_data, 'fork', 'Forking repository...')
            repo_full_name = self.github_client.parse_repo_url(repo_url)
            fork_info = self.github_client.fork_repository(repo_full_name)
            
            if fork_info['already_existed']:
                self._add_step(run_data, 'fork', f"Using existing fork: {fork_info['fork_repo']}", success=True)
            else:
                self._add_step(run_data, 'fork', f"Repository forked to: {fork_info['fork_repo']}", success=True)
            
            run_data['fork_info'] = fork_info
            
            # Step 1: Clone forked repository
            self._add_step(run_data, 'clone', 'Cloning forked repository...')
            temp_dir, repo = self.github_client.clone_repo(fork_info['fork_url'], base_branch)
            self._add_step(run_data, 'clone', 'Repository cloned', success=True)
            
            # Step 2: Get repository context
            self._add_step(run_data, 'context', 'Analyzing repository structure...')
            repo_structure = self.github_client.get_repo_structure(temp_dir)
            self._add_step(run_data, 'context', 'Structure analyzed', success=True)
            
            # Step 3: AI Analysis - Diagnose issue
            self._add_step(run_data, 'diagnosis', 'AI analyzing issue...')
            diagnosis = analyze_issue(issue_text, repo_structure)
            run_data['diagnosis'] = diagnosis.model_dump()
            self._add_step(run_data, 'diagnosis', f'Issue diagnosed: {diagnosis.summary}', success=True)
            
            # Step 4: Generate fix plan
            self._add_step(run_data, 'plan', 'AI generating fix plan...')
            plan = generate_fix_plan(diagnosis, repo_structure)
            run_data['plan'] = plan.model_dump()
            self._add_step(run_data, 'plan', f'Plan created: {plan.approach[:100]}...', success=True)
            
            # Step 5: Safety check
            estimated_changes = plan.estimated_changes
            files_count = len(plan.files_to_modify)
            safety_check = check_safety_limits(files_count, estimated_changes, max_files, max_lines)
            
            if not safety_check['within_limits']:
                raise Exception(f"Safety limits exceeded: {safety_check['issues']}")
            
            # Step 6: Generate patches for each file
            self._add_step(run_data, 'patch', f'Generating patches for {files_count} file(s)...')
            patches = []
            validation_results = []
            
            for file_path in plan.files_to_modify[:max_files]:  # Hard limit
                try:
                    # Read current file content
                    file_content = self.github_client.read_file(temp_dir, file_path)
                    
                    # Generate patch
                    patch = generate_code_patch(plan, file_content, file_path, diagnosis)
                    patches.append(patch.model_dump())
                    
                    # Validate patch
                    if file_path.endswith('.py'):
                        validation = self.validator.validate_patch(
                            file_path,
                            patch.original_code,
                            patch.modified_code
                        )
                        validation_results.append(validation)
                    
                except FileNotFoundError:
                    # File doesn't exist yet, will be created
                    patch = generate_code_patch(plan, "", file_path, diagnosis)
                    patches.append(patch.model_dump())
                except Exception as e:
                    self._add_step(run_data, 'patch', f'Error with {file_path}: {str(e)}', success=False)
            
            run_data['patches'] = patches
            run_data['validation_results'] = validation_results
            self._add_step(run_data, 'patch', f'Generated {len(patches)} patch(es)', success=True)
            
            # Step 7: Apply patches
            self._add_step(run_data, 'apply', 'Applying patches to repository...')
            branch_name = f"codesentinel-fix-{run_id}"
            self.github_client.create_branch(repo, branch_name)
            
            for patch_data in patches:
                file_path = patch_data['file_path']
                modified_code = patch_data['modified_code']
                self.github_client.write_file(temp_dir, file_path, modified_code)
            
            self._add_step(run_data, 'apply', 'Patches applied', success=True)
            
            # Step 8: Commit and push
            self._add_step(run_data, 'commit', 'Committing changes...')
            commit_message = f"Fix: {diagnosis.summary}\n\nGenerated by CodeSentinel AI Agent\nRun ID: {run_id}"
            self.github_client.commit_changes(repo, commit_message)
            self.github_client.push_branch(repo, branch_name)
            self._add_step(run_data, 'commit', 'Changes pushed to remote', success=True)
            
            # Step 9: Create pull request
            self._add_step(run_data, 'pr', 'Creating pull request...')
            pr_title = f"[CodeSentinel] {diagnosis.summary}"
            pr_body = self._generate_pr_description(diagnosis, plan, patches, validation_results)
            
            # Get fork owner for head branch specification
            fork_owner = fork_info['fork_repo'].split('/')[0]
            head_branch = f"{fork_owner}:{branch_name}"
            
            pr_url = self.github_client.create_pull_request(
                repo_full_name,  # Target the original repo
                pr_title,
                pr_body,
                head_branch,  # From fork:branch
                base_branch
            )
            
            run_data['pr_url'] = pr_url
            run_data['status'] = 'completed'
            self._add_step(run_data, 'pr', f'PR created: {pr_url}', success=True)
            
        except Exception as e:
            run_data['status'] = 'failed'
            run_data['error'] = str(e)
            run_data['traceback'] = traceback.format_exc()
            self._add_step(run_data, 'error', f'Error: {str(e)}', success=False)
        
        finally:
            # Cleanup
            if temp_dir:
                self.github_client.cleanup_temp_dir(temp_dir)
            
            # Save run data
            self.storage.save_run(run_id, run_data)
        
        return run_data
    
    def _add_step(self, run_data: Dict, step_name: str, message: str, success: Optional[bool] = None):
        """Add a step to the run data."""
        step: Dict = {
            'name': step_name,
            'message': message,
            'timestamp': datetime.now().isoformat()
        }
        if success is not None:
            step['success'] = success
        
        run_data['steps'].append(step)
    
    def _generate_pr_description(self, diagnosis, plan, patches, validation_results) -> str:
        """Generate PR description."""
        description = f"""## AI-Generated Fix

**Issue Summary:** {diagnosis.summary}

**Root Cause:** {diagnosis.root_cause}

**Complexity:** {diagnosis.complexity}

## Approach
{plan.approach}

## Changes Made
- Files modified: {len(patches)}
- Estimated changes: {plan.estimated_changes} lines

## Files Changed
"""
        for patch in patches:
            description += f"- `{patch['file_path']}`: {patch['explanation']}\n"
        
        description += "\n## Validation Results\n"
        if validation_results:
            for validation in validation_results:
                modified = validation['modified']
                description += f"- **{modified['file']}**: "
                if modified['passed']:
                    description += "✅ Passed"
                else:
                    description += "⚠️ Needs review"
                description += f" (Lint: {modified['lint_score']:.1f}, Complexity: {modified['complexity_score']:.1f})\n"
        else:
            description += "No Python files to validate.\n"
        
        description += "\n## Risks\n"
        for risk in plan.risks:
            description += f"- {risk}\n"
        
        description += "\n---\n*Generated by CodeSentinel - Autonomous GitHub AI Agent*"
        
        return description
    
    def get_run_history(self, limit: int = 10):
        """Get recent run history."""
        return self.storage.get_recent_runs(limit)
    
    def get_run_details(self, run_id: str):
        """Get details of a specific run."""
        return self.storage.load_run(run_id)