"""
RedTeam - AI Agent Security Assessment.

Automated security testing for AI agents. Test your agents against 67+
threat patterns before deployment.
"""

from __future__ import annotations

import httpx
import time
import json
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Optional, Callable, Any, List, Dict, Union


DEFAULT_API_URL = "https://agent-trust-infrastructure-production.up.railway.app"


# ============================================================================
# Enums
# ============================================================================

class ScanMode(str, Enum):
    """Scan thoroughness mode."""
    QUICK = "quick"           # Fast scan, ~20 attacks
    STANDARD = "standard"     # Balanced, ~50 attacks
    COMPREHENSIVE = "comprehensive"  # Full scan, 100+ attacks


class ThreatCategory(str, Enum):
    """Categories of security threats."""
    PROMPT_INJECTION = "prompt_injection"
    JAILBREAK = "jailbreak"
    DATA_EXFILTRATION = "data_exfiltration"
    MEMORY_POISONING = "memory_poisoning"
    MCP_ATTACKS = "mcp_attacks"
    A2A_ATTACKS = "a2a_attacks"
    RAG_POISONING = "rag_poisoning"
    INDIRECT_INJECTION = "indirect_injection"


class Severity(str, Enum):
    """Threat severity levels."""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


class RiskLevel(str, Enum):
    """Overall risk assessment."""
    LOW = "LOW"
    MEDIUM = "MEDIUM"
    HIGH = "HIGH"
    CRITICAL = "CRITICAL"


class ScanStatus(str, Enum):
    """Status of a running scan."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


# ============================================================================
# Data Models
# ============================================================================

@dataclass
class ThreatInfo:
    """Information about a threat template."""
    id: str
    name: str
    category: ThreatCategory
    severity: Severity
    description: str


@dataclass
class Vulnerability:
    """A discovered vulnerability."""
    threat_id: str
    threat_name: str
    category: ThreatCategory
    severity: Severity
    attack_payload: str
    response_snippet: Optional[str] = None
    matched_indicators: List[str] = field(default_factory=list)
    confidence: float = 0.0


@dataclass
class ScanProgress:
    """Progress of a running scan."""
    scan_id: str
    status: ScanStatus
    progress_percent: float
    completed_attacks: int
    total_attacks: int
    successful_attacks: int
    current_category: Optional[ThreatCategory] = None
    elapsed_seconds: float = 0.0


@dataclass
class ScanResult:
    """
    Complete results of a security scan.
    
    Attributes:
        scan_id: Unique identifier for this scan
        target_url: The agent endpoint that was tested
        timestamp: When the scan was run
        duration_seconds: How long the scan took
        security_score: Overall security score (0-100, higher is more secure)
        risk_level: Overall risk assessment (LOW, MEDIUM, HIGH, CRITICAL)
        total_attacks: Number of attacks attempted
        successful_attacks: Number of attacks that succeeded (vulnerabilities)
        blocked_attacks: Number of attacks the agent defended against
        vulnerabilities: List of discovered vulnerabilities
        severity_breakdown: Count of vulnerabilities by severity
        category_breakdown: Results broken down by threat category
        recommendations: Suggested remediation steps
    """
    scan_id: str
    target_url: str
    timestamp: datetime
    duration_seconds: float
    security_score: float
    risk_level: RiskLevel
    total_attacks: int
    successful_attacks: int
    blocked_attacks: int
    vulnerabilities: List[Vulnerability] = field(default_factory=list)
    severity_breakdown: Dict[str, int] = field(default_factory=dict)
    category_breakdown: Dict[str, Dict[str, int]] = field(default_factory=dict)
    recommendations: List[str] = field(default_factory=list)
    error: Optional[str] = None
    
    @property
    def is_secure(self) -> bool:
        """Returns True if no vulnerabilities were found."""
        return self.successful_attacks == 0
    
    @property
    def pass_rate(self) -> float:
        """Returns the percentage of attacks blocked (0-100)."""
        if self.total_attacks == 0:
            return 100.0
        return (self.blocked_attacks / self.total_attacks) * 100
    
    @property
    def critical_count(self) -> int:
        """Number of critical vulnerabilities."""
        return self.severity_breakdown.get("critical", 0)
    
    @property
    def high_count(self) -> int:
        """Number of high severity vulnerabilities."""
        return self.severity_breakdown.get("high", 0)
    
    @property
    def has_critical_issues(self) -> bool:
        """Returns True if any critical vulnerabilities exist."""
        return self.critical_count > 0
    
    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        return {
            "scan_id": self.scan_id,
            "target_url": self.target_url,
            "timestamp": self.timestamp.isoformat() if self.timestamp else None,
            "duration_seconds": self.duration_seconds,
            "security_score": self.security_score,
            "risk_level": self.risk_level.value if isinstance(self.risk_level, RiskLevel) else self.risk_level,
            "total_attacks": self.total_attacks,
            "successful_attacks": self.successful_attacks,
            "blocked_attacks": self.blocked_attacks,
            "pass_rate": self.pass_rate,
            "severity_breakdown": self.severity_breakdown,
            "category_breakdown": self.category_breakdown,
            "recommendations": self.recommendations,
            "vulnerabilities": [
                {
                    "threat_id": v.threat_id,
                    "threat_name": v.threat_name,
                    "category": v.category.value if isinstance(v.category, ThreatCategory) else v.category,
                    "severity": v.severity.value if isinstance(v.severity, Severity) else v.severity,
                    "confidence": v.confidence,
                }
                for v in self.vulnerabilities
            ],
            "error": self.error,
        }
    
    def to_json(self, indent: int = 2) -> str:
        """Convert to JSON string."""
        return json.dumps(self.to_dict(), indent=indent)


# ============================================================================
# Exceptions
# ============================================================================

class RedTeamError(Exception):
    """Base exception for RedTeam."""
    pass


class APIError(RedTeamError):
    """API returned an error response."""
    def __init__(self, message: str, status_code: int = None, response: dict = None):
        super().__init__(message)
        self.status_code = status_code
        self.response = response


class ScanError(RedTeamError):
    """Error during scan execution."""
    pass


class TimeoutError(RedTeamError):
    """Scan timed out."""
    pass


# ============================================================================
# RedTeam Client
# ============================================================================

class RedTeam:
    """
    AI Agent Security Assessment Client.
    
    Test your AI agents against 67+ threat patterns including prompt injection,
    jailbreaks, data exfiltration, memory poisoning, and more.
    
    Args:
        api_key: API key for authentication (optional for free tier)
        api_url: Base URL for the API (default: production URL)
        timeout: Request timeout in seconds (default: 60)
    
    Example:
        from agent_trust import RedTeam
        
        redteam = RedTeam(api_key="ta_xxx...")
        
        # Run a security scan
        result = redteam.scan(
            target_url="https://my-agent.example.com/chat",
            mode="standard",
        )
        
        print(f"Security Score: {result.security_score}/100")
        print(f"Risk Level: {result.risk_level}")
        print(f"Vulnerabilities Found: {result.successful_attacks}")
        
        if result.has_critical_issues:
            print("⚠️ Critical vulnerabilities detected!")
            for vuln in result.vulnerabilities:
                if vuln.severity == Severity.CRITICAL:
                    print(f"  - {vuln.threat_name}")
        
        # Export results
        redteam.export(result, "security-report.json")
    """
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        api_url: str = DEFAULT_API_URL,
        timeout: float = 60.0,
    ):
        self.api_url = api_url.rstrip("/")
        self.timeout = timeout
        self.api_key = api_key
        self._client = httpx.Client(timeout=timeout)
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.close()
    
    def close(self):
        """Close the HTTP client."""
        self._client.close()
    
    def _headers(self) -> dict:
        """Build request headers."""
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        return headers
    
    def _request(self, method: str, path: str, **kwargs) -> dict:
        """Make an HTTP request to the API."""
        url = f"{self.api_url}{path}"
        headers = kwargs.pop("headers", {})
        headers.update(self._headers())
        
        response = self._client.request(method, url, headers=headers, **kwargs)
        
        if response.status_code >= 400:
            try:
                error_data = response.json()
                message = error_data.get("detail", str(error_data))
            except Exception:
                message = response.text or f"HTTP {response.status_code}"
            raise APIError(message, response.status_code)
        
        return response.json()
    
    # ========================================================================
    # Threat Library
    # ========================================================================
    
    def list_threats(
        self,
        category: Optional[ThreatCategory] = None,
        severity: Optional[Severity] = None,
    ) -> List[ThreatInfo]:
        """
        List available threat templates.
        
        Args:
            category: Filter by threat category
            severity: Filter by severity level
        
        Returns:
            List of threat templates
        
        Example:
            threats = redteam.list_threats(category=ThreatCategory.PROMPT_INJECTION)
            print(f"Found {len(threats)} prompt injection patterns")
        """
        params = {}
        if category:
            params["category"] = category.value if isinstance(category, ThreatCategory) else category
        if severity:
            params["severity"] = severity.value if isinstance(severity, Severity) else severity
        
        data = self._request("GET", "/redteam/threats", params=params)
        
        return [
            ThreatInfo(
                id=t["id"],
                name=t["name"],
                category=ThreatCategory(t["category"]),
                severity=Severity(t["severity"]),
                description=t["description"],
            )
            for t in data
        ]
    
    def threat_stats(self) -> dict:
        """
        Get threat library statistics.
        
        Returns:
            Dictionary with total threats, counts by category and severity
        
        Example:
            stats = redteam.threat_stats()
            print(f"Total threats: {stats['total_threats']}")
            print(f"Categories: {stats['by_category']}")
        """
        return self._request("GET", "/redteam/threats/stats")
    
    # ========================================================================
    # Scanning
    # ========================================================================
    
    def scan(
        self,
        target_url: str,
        *,
        mode: Union[ScanMode, str] = ScanMode.STANDARD,
        categories: Optional[List[ThreatCategory]] = None,
        severity_min: Union[Severity, str] = Severity.LOW,
        max_attacks: Optional[int] = None,
        auth_token: Optional[str] = None,
        payload_field: str = "message",
        response_field: Optional[str] = None,
        headers: Optional[Dict[str, str]] = None,
        timeout_seconds: int = 30,
        concurrency: int = 5,
        poll_interval: float = 2.0,
        max_wait: float = 300.0,
        on_progress: Optional[Callable[[ScanProgress], None]] = None,
    ) -> ScanResult:
        """
        Run a security scan against an AI agent.
        
        Starts a background scan and polls for completion. Use on_progress
        for real-time progress updates.
        
        Args:
            target_url: The agent's API endpoint (e.g., https://api.example.com/chat)
            mode: Scan thoroughness - "quick", "standard", or "comprehensive"
            categories: Limit testing to specific threat categories
            severity_min: Minimum severity level to test
            max_attacks: Limit the number of attack attempts
            auth_token: Bearer token for agent authentication
            payload_field: JSON field name for the message payload (default: "message")
            response_field: JSON field for response extraction (auto-detect if None)
            headers: Additional HTTP headers for agent requests
            timeout_seconds: Request timeout for each attack
            concurrency: Number of concurrent attack requests
            poll_interval: Seconds between status polls
            max_wait: Maximum seconds to wait for scan completion
            on_progress: Callback for progress updates
        
        Returns:
            ScanResult with security score, vulnerabilities, and recommendations
        
        Raises:
            APIError: If the API returns an error
            TimeoutError: If scan doesn't complete within max_wait
            ScanError: If scan fails
        
        Example:
            # Basic scan
            result = redteam.scan("https://my-agent.com/chat")
            
            # Scan with options
            result = redteam.scan(
                "https://my-agent.com/chat",
                mode="comprehensive",
                categories=[ThreatCategory.PROMPT_INJECTION, ThreatCategory.JAILBREAK],
                auth_token="Bearer sk-xxx",
                on_progress=lambda p: print(f"Progress: {p.progress_percent:.0f}%"),
            )
            
            if result.risk_level in (RiskLevel.HIGH, RiskLevel.CRITICAL):
                print("⚠️ Agent has serious security issues!")
        """
        # Normalize enums
        if isinstance(mode, str):
            mode = ScanMode(mode)
        if isinstance(severity_min, str):
            severity_min = Severity(severity_min)
        
        # Build request
        request_data = {
            "target_url": target_url,
            "mode": mode.value,
            "severity_min": severity_min.value,
            "payload_field": payload_field,
            "timeout_seconds": timeout_seconds,
            "concurrency": concurrency,
        }
        
        if categories:
            request_data["categories"] = [
                c.value if isinstance(c, ThreatCategory) else c
                for c in categories
            ]
        if max_attacks:
            request_data["max_attacks"] = max_attacks
        if auth_token:
            request_data["auth_token"] = auth_token
        if response_field:
            request_data["response_field"] = response_field
        if headers:
            request_data["headers"] = headers
        
        # Start scan
        start_response = self._request("POST", "/redteam/scan", json=request_data)
        scan_id = start_response["scan_id"]
        
        # Poll for completion
        start_time = time.time()
        while True:
            elapsed = time.time() - start_time
            if elapsed > max_wait:
                raise TimeoutError(f"Scan did not complete within {max_wait}s")
            
            status = self.get_scan_status(scan_id)
            
            if on_progress:
                on_progress(status)
            
            if status.status == ScanStatus.COMPLETED:
                break
            elif status.status == ScanStatus.FAILED:
                raise ScanError(f"Scan failed: {scan_id}")
            elif status.status == ScanStatus.CANCELLED:
                raise ScanError(f"Scan was cancelled: {scan_id}")
            
            time.sleep(poll_interval)
        
        # Get results
        return self.get_scan_result(scan_id)
    
    def scan_async_start(
        self,
        target_url: str,
        *,
        mode: Union[ScanMode, str] = ScanMode.STANDARD,
        categories: Optional[List[ThreatCategory]] = None,
        severity_min: Union[Severity, str] = Severity.LOW,
        max_attacks: Optional[int] = None,
        auth_token: Optional[str] = None,
        payload_field: str = "message",
        response_field: Optional[str] = None,
        headers: Optional[Dict[str, str]] = None,
        timeout_seconds: int = 30,
        concurrency: int = 5,
    ) -> str:
        """
        Start a scan without waiting for completion.
        
        Returns the scan_id immediately. Use get_scan_status() and 
        get_scan_result() to check progress and retrieve results.
        
        Returns:
            scan_id: Unique identifier for the scan
        
        Example:
            # Start scan
            scan_id = redteam.scan_async_start("https://my-agent.com/chat")
            
            # Do other work...
            
            # Check status
            status = redteam.get_scan_status(scan_id)
            if status.status == ScanStatus.COMPLETED:
                result = redteam.get_scan_result(scan_id)
        """
        # Normalize enums
        if isinstance(mode, str):
            mode = ScanMode(mode)
        if isinstance(severity_min, str):
            severity_min = Severity(severity_min)
        
        request_data = {
            "target_url": target_url,
            "mode": mode.value,
            "severity_min": severity_min.value,
            "payload_field": payload_field,
            "timeout_seconds": timeout_seconds,
            "concurrency": concurrency,
        }
        
        if categories:
            request_data["categories"] = [
                c.value if isinstance(c, ThreatCategory) else c
                for c in categories
            ]
        if max_attacks:
            request_data["max_attacks"] = max_attacks
        if auth_token:
            request_data["auth_token"] = auth_token
        if response_field:
            request_data["response_field"] = response_field
        if headers:
            request_data["headers"] = headers
        
        response = self._request("POST", "/redteam/scan", json=request_data)
        return response["scan_id"]
    
    def scan_mock(
        self,
        vulnerability_rate: float = 0.3,
        mode: Union[ScanMode, str] = ScanMode.QUICK,
        max_attacks: int = 20,
    ) -> ScanResult:
        """
        Run a scan against a mock target for testing/demo.
        
        Args:
            vulnerability_rate: Simulated vulnerability rate (0.0-1.0)
            mode: Scan mode
            max_attacks: Number of attacks to simulate
        
        Returns:
            ScanResult with simulated results
        
        Example:
            # Test the SDK integration
            result = redteam.scan_mock(vulnerability_rate=0.5)
            print(f"Mock score: {result.security_score}")
        """
        if isinstance(mode, str):
            mode = ScanMode(mode)
        
        data = self._request(
            "POST",
            "/redteam/scan/mock",
            params={
                "vulnerability_rate": vulnerability_rate,
                "mode": mode.value,
                "max_attacks": max_attacks,
            }
        )
        
        return self._parse_scan_result(data)
    
    # ========================================================================
    # Scan Status & Results
    # ========================================================================
    
    def get_scan_status(self, scan_id: str) -> ScanProgress:
        """
        Get the status of a running scan.
        
        Args:
            scan_id: The scan identifier returned by scan_async_start()
        
        Returns:
            ScanProgress with current status and progress
        """
        data = self._request("GET", f"/redteam/scan/{scan_id}/status")
        
        return ScanProgress(
            scan_id=data["scan_id"],
            status=ScanStatus(data["status"]),
            progress_percent=data["progress_percent"],
            completed_attacks=data["completed_attacks"],
            total_attacks=data["total_attacks"],
            successful_attacks=data["successful_attacks"],
            current_category=ThreatCategory(data["current_category"]) if data.get("current_category") else None,
            elapsed_seconds=data.get("elapsed_seconds", 0),
        )
    
    def get_scan_result(self, scan_id: str, full: bool = False) -> ScanResult:
        """
        Get the results of a completed scan.
        
        Args:
            scan_id: The scan identifier
            full: If True, include full attack details (can be large)
        
        Returns:
            ScanResult with security assessment
        
        Raises:
            APIError: If scan not found or still in progress
        """
        endpoint = f"/redteam/scan/{scan_id}/full" if full else f"/redteam/scan/{scan_id}/result"
        data = self._request("GET", endpoint)
        return self._parse_scan_result(data, scan_id=scan_id)
    
    def cancel_scan(self, scan_id: str) -> bool:
        """
        Cancel a running scan.
        
        Args:
            scan_id: The scan identifier
        
        Returns:
            True if cancellation was successful
        """
        data = self._request("POST", f"/redteam/scan/{scan_id}/cancel")
        return data.get("status") == "cancelled"
    
    # ========================================================================
    # Export
    # ========================================================================
    
    def export(
        self,
        result: ScanResult,
        output_path: str,
        format: str = "json",
    ) -> str:
        """
        Export scan results to a file.
        
        Args:
            result: The ScanResult to export
            output_path: Path to write the report
            format: Output format - "json" (more formats coming)
        
        Returns:
            The output path
        
        Example:
            result = redteam.scan("https://my-agent.com/chat")
            redteam.export(result, "security-report.json")
        """
        if format == "json":
            content = result.to_json(indent=2)
        else:
            raise ValueError(f"Unsupported format: {format}. Supported: json")
        
        with open(output_path, "w") as f:
            f.write(content)
        
        return output_path
    
    # ========================================================================
    # Internal
    # ========================================================================
    
    def _parse_scan_result(self, data: dict, scan_id: str = None) -> ScanResult:
        """Parse API response into ScanResult."""
        # Handle timestamp
        timestamp = data.get("timestamp")
        if isinstance(timestamp, str):
            try:
                timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
            except:
                timestamp = datetime.utcnow()
        elif timestamp is None:
            timestamp = datetime.utcnow()
        
        # Parse vulnerabilities if present
        vulnerabilities = []
        if "top_vulnerabilities" in data:
            for v in data["top_vulnerabilities"]:
                vulnerabilities.append(Vulnerability(
                    threat_id=v.get("threat_id", ""),
                    threat_name=v.get("threat_name", ""),
                    category=ThreatCategory(v["category"]) if v.get("category") else ThreatCategory.PROMPT_INJECTION,
                    severity=Severity(v["severity"]) if v.get("severity") else Severity.MEDIUM,
                    attack_payload=v.get("attack_payload", ""),
                    response_snippet=v.get("response_snippet"),
                    matched_indicators=v.get("matched_indicators", []),
                    confidence=v.get("confidence", 0.0),
                ))
        
        # Build severity breakdown
        severity_breakdown = data.get("severity_breakdown", {})
        if not severity_breakdown:
            severity_breakdown = {
                "critical": data.get("critical_count", 0),
                "high": data.get("high_count", 0),
                "medium": data.get("medium_count", 0),
                "low": data.get("low_count", 0),
            }
        
        return ScanResult(
            scan_id=data.get("scan_id") or scan_id or "unknown",
            target_url=data.get("target_url", "unknown"),
            timestamp=timestamp,
            duration_seconds=data.get("duration_seconds", 0),
            security_score=data.get("security_score", 0),
            risk_level=RiskLevel(data["risk_level"]) if data.get("risk_level") else RiskLevel.MEDIUM,
            total_attacks=data.get("total_attacks", 0),
            successful_attacks=data.get("successful_attacks", 0),
            blocked_attacks=data.get("blocked_attacks", 0),
            vulnerabilities=vulnerabilities,
            severity_breakdown=severity_breakdown,
            category_breakdown=data.get("results_by_category", {}),
            recommendations=data.get("recommendations", []),
            error=data.get("error"),
        )


# ============================================================================
# Async Client
# ============================================================================

class AsyncRedTeam:
    """
    Async version of the RedTeam client.
    
    Example:
        async with AsyncRedTeam(api_key="ta_xxx...") as redteam:
            result = await redteam.scan("https://my-agent.com/chat")
            print(f"Security Score: {result.security_score}")
    """
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        api_url: str = DEFAULT_API_URL,
        timeout: float = 60.0,
    ):
        self.api_url = api_url.rstrip("/")
        self.timeout = timeout
        self.api_key = api_key
        self._client = httpx.AsyncClient(timeout=timeout)
    
    async def __aenter__(self):
        return self
    
    async def __aexit__(self, *args):
        await self.close()
    
    async def close(self):
        """Close the HTTP client."""
        await self._client.aclose()
    
    def _headers(self) -> dict:
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        return headers
    
    async def _request(self, method: str, path: str, **kwargs) -> dict:
        url = f"{self.api_url}{path}"
        headers = kwargs.pop("headers", {})
        headers.update(self._headers())
        
        response = await self._client.request(method, url, headers=headers, **kwargs)
        
        if response.status_code >= 400:
            try:
                error_data = response.json()
                message = error_data.get("detail", str(error_data))
            except Exception:
                message = response.text or f"HTTP {response.status_code}"
            raise APIError(message, response.status_code)
        
        return response.json()
    
    async def list_threats(
        self,
        category: Optional[ThreatCategory] = None,
        severity: Optional[Severity] = None,
    ) -> List[ThreatInfo]:
        """List available threat templates."""
        params = {}
        if category:
            params["category"] = category.value if isinstance(category, ThreatCategory) else category
        if severity:
            params["severity"] = severity.value if isinstance(severity, Severity) else severity
        
        data = await self._request("GET", "/redteam/threats", params=params)
        
        return [
            ThreatInfo(
                id=t["id"],
                name=t["name"],
                category=ThreatCategory(t["category"]),
                severity=Severity(t["severity"]),
                description=t["description"],
            )
            for t in data
        ]
    
    async def threat_stats(self) -> dict:
        """Get threat library statistics."""
        return await self._request("GET", "/redteam/threats/stats")
    
    async def scan(
        self,
        target_url: str,
        *,
        mode: Union[ScanMode, str] = ScanMode.STANDARD,
        categories: Optional[List[ThreatCategory]] = None,
        severity_min: Union[Severity, str] = Severity.LOW,
        max_attacks: Optional[int] = None,
        auth_token: Optional[str] = None,
        payload_field: str = "message",
        response_field: Optional[str] = None,
        headers: Optional[Dict[str, str]] = None,
        timeout_seconds: int = 30,
        concurrency: int = 5,
        poll_interval: float = 2.0,
        max_wait: float = 300.0,
        on_progress: Optional[Callable[[ScanProgress], None]] = None,
    ) -> ScanResult:
        """Run a security scan against an AI agent."""
        import asyncio
        
        if isinstance(mode, str):
            mode = ScanMode(mode)
        if isinstance(severity_min, str):
            severity_min = Severity(severity_min)
        
        request_data = {
            "target_url": target_url,
            "mode": mode.value,
            "severity_min": severity_min.value,
            "payload_field": payload_field,
            "timeout_seconds": timeout_seconds,
            "concurrency": concurrency,
        }
        
        if categories:
            request_data["categories"] = [
                c.value if isinstance(c, ThreatCategory) else c
                for c in categories
            ]
        if max_attacks:
            request_data["max_attacks"] = max_attacks
        if auth_token:
            request_data["auth_token"] = auth_token
        if response_field:
            request_data["response_field"] = response_field
        if headers:
            request_data["headers"] = headers
        
        start_response = await self._request("POST", "/redteam/scan", json=request_data)
        scan_id = start_response["scan_id"]
        
        start_time = time.time()
        while True:
            elapsed = time.time() - start_time
            if elapsed > max_wait:
                raise TimeoutError(f"Scan did not complete within {max_wait}s")
            
            status = await self.get_scan_status(scan_id)
            
            if on_progress:
                on_progress(status)
            
            if status.status == ScanStatus.COMPLETED:
                break
            elif status.status == ScanStatus.FAILED:
                raise ScanError(f"Scan failed: {scan_id}")
            elif status.status == ScanStatus.CANCELLED:
                raise ScanError(f"Scan was cancelled: {scan_id}")
            
            await asyncio.sleep(poll_interval)
        
        return await self.get_scan_result(scan_id)
    
    async def scan_mock(
        self,
        vulnerability_rate: float = 0.3,
        mode: Union[ScanMode, str] = ScanMode.QUICK,
        max_attacks: int = 20,
    ) -> ScanResult:
        """Run a scan against a mock target."""
        if isinstance(mode, str):
            mode = ScanMode(mode)
        
        data = await self._request(
            "POST",
            "/redteam/scan/mock",
            params={
                "vulnerability_rate": vulnerability_rate,
                "mode": mode.value,
                "max_attacks": max_attacks,
            }
        )
        
        return self._parse_scan_result(data)
    
    async def get_scan_status(self, scan_id: str) -> ScanProgress:
        """Get the status of a running scan."""
        data = await self._request("GET", f"/redteam/scan/{scan_id}/status")
        
        return ScanProgress(
            scan_id=data["scan_id"],
            status=ScanStatus(data["status"]),
            progress_percent=data["progress_percent"],
            completed_attacks=data["completed_attacks"],
            total_attacks=data["total_attacks"],
            successful_attacks=data["successful_attacks"],
            current_category=ThreatCategory(data["current_category"]) if data.get("current_category") else None,
            elapsed_seconds=data.get("elapsed_seconds", 0),
        )
    
    async def get_scan_result(self, scan_id: str, full: bool = False) -> ScanResult:
        """Get the results of a completed scan."""
        endpoint = f"/redteam/scan/{scan_id}/full" if full else f"/redteam/scan/{scan_id}/result"
        data = await self._request("GET", endpoint)
        return self._parse_scan_result(data, scan_id=scan_id)
    
    async def cancel_scan(self, scan_id: str) -> bool:
        """Cancel a running scan."""
        data = await self._request("POST", f"/redteam/scan/{scan_id}/cancel")
        return data.get("status") == "cancelled"
    
    def export(
        self,
        result: ScanResult,
        output_path: str,
        format: str = "json",
    ) -> str:
        """Export scan results to a file."""
        if format == "json":
            content = result.to_json(indent=2)
        else:
            raise ValueError(f"Unsupported format: {format}")
        
        with open(output_path, "w") as f:
            f.write(content)
        
        return output_path
    
    def _parse_scan_result(self, data: dict, scan_id: str = None) -> ScanResult:
        """Parse API response into ScanResult."""
        timestamp = data.get("timestamp")
        if isinstance(timestamp, str):
            try:
                timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
            except:
                timestamp = datetime.utcnow()
        elif timestamp is None:
            timestamp = datetime.utcnow()
        
        vulnerabilities = []
        if "top_vulnerabilities" in data:
            for v in data["top_vulnerabilities"]:
                vulnerabilities.append(Vulnerability(
                    threat_id=v.get("threat_id", ""),
                    threat_name=v.get("threat_name", ""),
                    category=ThreatCategory(v["category"]) if v.get("category") else ThreatCategory.PROMPT_INJECTION,
                    severity=Severity(v["severity"]) if v.get("severity") else Severity.MEDIUM,
                    attack_payload=v.get("attack_payload", ""),
                    response_snippet=v.get("response_snippet"),
                    matched_indicators=v.get("matched_indicators", []),
                    confidence=v.get("confidence", 0.0),
                ))
        
        severity_breakdown = data.get("severity_breakdown", {})
        if not severity_breakdown:
            severity_breakdown = {
                "critical": data.get("critical_count", 0),
                "high": data.get("high_count", 0),
                "medium": data.get("medium_count", 0),
                "low": data.get("low_count", 0),
            }
        
        return ScanResult(
            scan_id=data.get("scan_id") or scan_id or "unknown",
            target_url=data.get("target_url", "unknown"),
            timestamp=timestamp,
            duration_seconds=data.get("duration_seconds", 0),
            security_score=data.get("security_score", 0),
            risk_level=RiskLevel(data["risk_level"]) if data.get("risk_level") else RiskLevel.MEDIUM,
            total_attacks=data.get("total_attacks", 0),
            successful_attacks=data.get("successful_attacks", 0),
            blocked_attacks=data.get("blocked_attacks", 0),
            vulnerabilities=vulnerabilities,
            severity_breakdown=severity_breakdown,
            category_breakdown=data.get("results_by_category", {}),
            recommendations=data.get("recommendations", []),
            error=data.get("error"),
        )
