import subprocess
import json
from dataclasses import dataclass, field
from typing import List, Optional
from rich.console import Console
from rich.panel import Panel
from pathlib import Path
from rich.table import Table
from rich.text import Text
from rich.rule import Rule
from rich.markup import escape

console = Console()

@dataclass
class ScanResult:
    score: float  # 0.0 to 1.0 (1.0 = safe, 0.0 = critical issues)
    issues: List[str] = field(default_factory=list)
    critical_found: bool = False
    agent_risks_found: bool = False # New flag for Agent specific risks

class RepoScanner:
    """Scans a repository for configuration and documentation files."""
    
    def __init__(self, root_dir: Path):
        self.root_dir = root_dir

    def scan_for_config_files(self) -> List[str]:
        """
        Scans for docker-compose, Dockerfile, and READMEs up to depth 3.
        Returns a sorted list of relative paths.
        """
        patterns = [
            "docker-compose.yml", "docker-compose.yaml", "compose.yml", "compose.yaml",
            "Dockerfile", "Containerfile",
            "README.md", "README.txt", "Makefile"
        ]
        found_files = []
        
        try:
            # Depth 0 (Root)
            for p in patterns:
                found_files.extend(self.root_dir.glob(p))
            
            # Depth 1 
            for p in patterns:
                found_files.extend(self.root_dir.glob(f"*/{p}"))
                
            # Depth 2
            for p in patterns:
                found_files.extend(self.root_dir.glob(f"*/*/{p}"))

            # Depth 3
            for p in patterns:
                found_files.extend(self.root_dir.glob(f"*/*/*/{p}"))

            # Dedup and sort
            unique_files = sorted(list(set([str(f.relative_to(self.root_dir)) for f in found_files if f.is_file()])))
            return unique_files
        except Exception:
            return []

class VulnerabilityScanner:
    def __init__(self, target_dir: str = "."):
        self.target_dir = target_dir

    def render_report(self, result: ScanResult, app_name: str):
        """Render a beautiful security report."""
        console.print()
        console.print(Rule(f"🛡️  Security Audit Report: {app_name}", style="bold cyan"))
        
        if not result.issues:
            console.print(Panel(
                f"[bold green]No critical vulnerabilities detected in {app_name}.[/bold green]\n"
                "[dim]Scan included dependency checks (Safety) and static analysis (Bandit).[/dim]",
                border_style="green",
                expand=False
            ))
            return

        table = Table(title="Security Findings", border_style="yellow", show_lines=True)
        table.add_column("Type", style="bold")
        table.add_column("Severity", justify="center")
        table.add_column("Details")

        for issue in result.issues:
            severity = "[red]HIGH[/red]" if "HIGH" in issue or "AGENT RISK" in issue else "[yellow]MED[/yellow]"
            type_str = "Code Risk" if "BANDIT" in issue else "Dependency"
            table.add_row(type_str, severity, issue)

        console.print(table)
        
        if result.critical_found:
            console.print(Panel(
                "[bold red]⚠  CRITICAL RISKS DETECTED[/bold red]\n"
                "This agent contains potential security flaws. Proceed with caution.",
                border_style="red",
                expand=False
            ))
        console.print()

    def scan(self) -> 'ScanResult':
        issues = []
        critical = False
        agent_risks = False
        
        # 1. Run Safety (Dependency check)
        try:
            # We run it on the current project directory
            subprocess.run(["safety", "check", "--json"], capture_output=True, text=True, check=True)
        except subprocess.CalledProcessError as e:
            if e.returncode != 0:
                try:
                    data = json.loads(e.stdout)
                    # safety check returns non-zero if vulnerabilities found
                    for vuln in data.get("vulnerabilities", []):
                        issues.append(f"SAFETY: {vuln.get('package_name')} - {vuln.get('vulnerability_id')}")
                        critical = True
                except:
                    pass
        except FileNotFoundError:
            pass # Safety not installed

        # 2. Run Bandit (Static Analysis)
        try:
            # Exclude large/irrelevant dirs to prevent "stuck" behavior
            exclude = ".venv,.git,installed,node_modules,__pycache__,tests"
            res = subprocess.run(
                ["bandit", "-r", self.target_dir, "-f", "json", "-q", "-x", exclude],
                capture_output=True, text=True
            )
            data = json.loads(res.stdout)
            
            # Check results for agent-specific risks (e.g. B602 subprocess shell=True)
            for result in data.get("results", []):
                if result.get("test_id") in ["B602", "B604"]:
                    issues.append(f"BANDIT AGENT RISK: {result.get('issue_text')} in {result.get('filename')}")
                    agent_risks = True
                    critical = True

            # Check metrics for HIGH severity
            metrics = data.get("metrics", {}).get("_totals", {})
            high_count = metrics.get("SEVERITY.HIGH", 0)
            if high_count > 0:
                issues.append(f"BANDIT: Found {high_count} HIGH severity issues.")
                critical = True
                    
        except (FileNotFoundError, json.JSONDecodeError):
            pass

        score = 0.0 if critical else 1.0
        return ScanResult(score=score, issues=issues, critical_found=critical, agent_risks_found=agent_risks)
