from src.ai_client import AIClient
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown

console = Console()

class CodeAnalyzer:
    def __init__(self):
        self.ai = AIClient()

    def analyze_security(self, file_path: str, code: str) -> str:
        system_prompt = "You are a senior security engineer. Analyze the provided code for vulnerabilities (SQLi, XSS, Secrets, etc.). Return a brief report."
        user_prompt = f"File: {file_path}\n\nCode:\n{code}"
        return self.ai.ask(system_prompt, user_prompt)

    def generate_tests(self, file_path: str, code: str) -> str:
        system_prompt = "You are a QA engineer. Generate unit test cases for the provided code. Use the appropriate testing framework for the language detected."
        user_prompt = f"File: {file_path}\n\nCode:\n{code}"
        return self.ai.ask(system_prompt, user_prompt)

    def suggest_improvements(self, file_path: str, code: str) -> str:
        system_prompt = "You are a lead developer. Suggest refactoring, performance improvements, and best practices for the provided code."
        user_prompt = f"File: {file_path}\n\nCode:\n{code}"
        return self.ai.ask(system_prompt, user_prompt)

    def run_full_audit(self, files_data):
        console.print(Panel("[bold magenta]Starting AI Audit...[/bold magenta]"))
        
        # Limit analysis to first 5 files for this MVP to save tokens/time
        # In production, you would queue this or run selectively
        files_to_analyze = files_data[:5] 
        
        for path, code in files_to_analyze:
            console.print(f"\n[bold cyan]Analyzing: {path}[/bold cyan]")
            
            # 1. Security
            console.print("[yellow]🔒 Checking Security...[/yellow]")
            security_report = self.analyze_security(path, code)
            console.print(Markdown(security_report))
            
            # 2. Tests
            console.print("[yellow]🧪 Generating Tests...[/yellow]")
            test_report = self.generate_tests(path, code)
            console.print(Markdown(test_report))
            
            # 3. Improvements
            console.print("[yellow]🚀 Suggesting Improvements...[/yellow]")
            improvement_report = self.suggest_improvements(path, code)
            console.print(Markdown(improvement_report))
            
            console.print("-" * 50)
        
        console.print("\n[bold green]Audit Complete![/bold green]")