"""Adversaria Python SDK - Adversarial Testing Harness for LLMs"""

import subprocess
import json
import os
from typing import List, Optional, Dict, Any
from dataclasses import dataclass
from datetime import datetime

__version__ = "0.1.0"


@dataclass
class TestResult:
    """Test execution results"""
    id: str
    model: str
    provider: str
    timestamp: str
    risk_score: int
    total_attacks: int
    successful_attacks: int
    failed_attacks: int
    duration_ms: int
    raw_data: Dict[str, Any]

    def save_report(self, directory: str = "./reports") -> str:
        """Save report to directory"""
        os.makedirs(directory, exist_ok=True)
        filename = f"adversaria_report_{self.id}.json"
        filepath = os.path.join(directory, filename)
        
        with open(filepath, 'w') as f:
            json.dump(self.raw_data, f, indent=2)
        
        return filepath


@dataclass
class Suite:
    """Attack suite information"""
    id: str
    name: str
    description: str
    category: str
    payload_count: int
    enabled: bool


class Adversaria:
    """Adversaria client for running LLM security tests"""
    
    def __init__(self, config_path: Optional[str] = None):
        """
        Initialize Adversaria client
        
        Args:
            config_path: Optional path to configuration file
        """
        self.config_path = config_path
        self._check_cli_installed()
    
    def _check_cli_installed(self):
        """Check if adversaria CLI is installed"""
        try:
            subprocess.run(
                ["adversaria", "--version"],
                capture_output=True,
                check=True
            )
        except (subprocess.CalledProcessError, FileNotFoundError):
            raise RuntimeError(
                "Adversaria CLI not found. Install with: cargo install adversaria"
            )
    
    def run(
        self,
        provider: str,
        model: str,
        suites: Optional[List[str]] = None,
        api_key: Optional[str] = None
    ) -> TestResult:
        """
        Run security tests against an LLM
        
        Args:
            provider: Provider name ('openai', 'anthropic', 'ollama')
            model: Model name
            suites: Optional list of suite IDs to run
            api_key: Optional API key (uses env var if not provided)
        
        Returns:
            TestResult object with test results
        """
        # Build command
        cmd = ["adversaria", "run", "--provider", provider, "--model", model]
        
        if suites:
            cmd.extend(["--suites", ",".join(suites)])
        
        # Set API key if provided
        env = os.environ.copy()
        if api_key:
            if provider == "openai":
                env["OPENAI_API_KEY"] = api_key
            elif provider == "anthropic":
                env["ANTHROPIC_API_KEY"] = api_key
        
        # Run command
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True,
                env=env
            )
            
            # Parse output to get report path
            output = result.stdout
            
            # Get latest report
            reports = self.list_reports()
            if not reports:
                raise RuntimeError("No reports found after test execution")
            
            latest_report = reports[0]
            
            # Load report
            with open(latest_report, 'r') as f:
                data = json.load(f)
            
            return TestResult(
                id=data.get("id", ""),
                model=data.get("model", ""),
                provider=data.get("provider", ""),
                timestamp=data.get("timestamp", ""),
                risk_score=data.get("overall_risk_score", 0),
                total_attacks=data.get("total_attacks", 0),
                successful_attacks=data.get("successful_attacks", 0),
                failed_attacks=data.get("failed_attacks", 0),
                duration_ms=data.get("duration_ms", 0),
                raw_data=data
            )
            
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"Test execution failed: {e.stderr}")
    
    def list_suites(self) -> List[Suite]:
        """
        List available attack suites
        
        Returns:
            List of Suite objects
        """
        try:
            result = subprocess.run(
                ["adversaria", "list"],
                capture_output=True,
                text=True,
                check=True
            )
            
            # Parse output (simplified - in production would parse table output)
            # For now, return empty list as CLI output is formatted table
            return []
            
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"Failed to list suites: {e.stderr}")
    
    def list_reports(self, directory: str = "./reports") -> List[str]:
        """
        List available reports
        
        Args:
            directory: Reports directory
        
        Returns:
            List of report file paths
        """
        if not os.path.exists(directory):
            return []
        
        reports = [
            os.path.join(directory, f)
            for f in os.listdir(directory)
            if f.endswith('.json') and f.startswith('adversaria_report_')
        ]
        
        # Sort by modification time (newest first)
        reports.sort(key=lambda x: os.path.getmtime(x), reverse=True)
        
        return reports
    
    def load_report(self, filepath: str) -> TestResult:
        """
        Load a report from file
        
        Args:
            filepath: Path to report file
        
        Returns:
            TestResult object
        """
        with open(filepath, 'r') as f:
            data = json.load(f)
        
        return TestResult(
            id=data.get("id", ""),
            model=data.get("model", ""),
            provider=data.get("provider", ""),
            timestamp=data.get("timestamp", ""),
            risk_score=data.get("overall_risk_score", 0),
            total_attacks=data.get("total_attacks", 0),
            successful_attacks=data.get("successful_attacks", 0),
            failed_attacks=data.get("failed_attacks", 0),
            duration_ms=data.get("duration_ms", 0),
            raw_data=data
        )


__all__ = ["Adversaria", "TestResult", "Suite"]
