"""Pentest runner for red-teaming AI agents."""
from __future__ import annotations
import time
from typing import Any, List, Optional, TYPE_CHECKING

from .models import PENTEST_CATEGORIES

if TYPE_CHECKING:
    from .client import AshrLabsClient

_PENTEST_AVAILABLE = False

def _ensure_pentest_deps():
    global _PENTEST_AVAILABLE
    if _PENTEST_AVAILABLE:
        return
    try:
        import claude_agent_sdk
        _PENTEST_AVAILABLE = True
    except ImportError:
        raise ImportError("Pentest requires claude-agent-sdk. Install: pip install claude-agent-sdk")


def build_pentest_request_config(
    agent_name: str,
    agent_description: str = "",
    system_prompt: str = "",
    tools: list[dict[str, Any]] | None = None,
    categories: list[str] | None = None,
    num_scenarios_per_category: int = 3,
    difficulty: str = "mixed",
    include_multi_turn: bool = True,
) -> dict[str, Any]:
    """Build the request config dict for a pentest dataset generation request.

    This creates the properly-structured config that the Ashr Labs API
    expects when generating a security evaluation dataset.

    Args:
        agent_name: Name of the agent being tested.
        agent_description: Description of what the agent does.
        system_prompt: The agent's system prompt (helps generate
            targeted attack scenarios).
        tools: List of tool definitions the agent has access to.
            Each tool should have ``name``, ``description``, and
            optionally ``parameters``.
        categories: Security categories to include. Defaults to all 7:
            prompt_injection, privilege_escalation, data_exfiltration,
            jailbreak, tool_abuse, denial_of_service, social_engineering.
        num_scenarios_per_category: Number of attack scenarios per
            category (default: 3).
        difficulty: Difficulty level - "easy", "medium", "hard", or
            "mixed" (default: "mixed").
        include_multi_turn: Whether to include multi-turn attack
            scenarios (default: True).

    Returns:
        A config dict ready to pass to ``create_request`` or
        ``generate_pentest_dataset``.

    Raises:
        ValueError: If any category name is invalid.
    """
    if difficulty not in ("easy", "medium", "hard", "mixed"):
        raise ValueError(
            f"difficulty must be 'easy', 'medium', 'hard', or 'mixed', got '{difficulty}'"
        )

    resolved_categories = categories or list(PENTEST_CATEGORIES)
    invalid = set(resolved_categories) - set(PENTEST_CATEGORIES)
    if invalid:
        raise ValueError(
            f"Invalid pentest categories: {invalid}. "
            f"Valid categories: {PENTEST_CATEGORIES}"
        )

    return {
        "metadata": {
            "dataset_name": f"Pentest: {agent_name}",
            "description": f"Security evaluation dataset for {agent_name}",
            "dataset_type": "pentest",
        },
        "agent": {
            "name": agent_name,
            "description": agent_description,
            "system_prompt": system_prompt,
            "tools": tools or [],
            "accepted_inputs": {"text": True, "audio": False, "file": False, "image": False, "video": False},
            "output_format": {"type": "text"},
        },
        "context": {
            "domain": "cybersecurity",
            "use_case": "security_evaluation",
            "scenario_context": f"Red-team evaluation of {agent_name}: {agent_description}",
        },
        "test_config": {
            "num_variations": num_scenarios_per_category,
            "coverage": {
                "happy_path": False,
                "edge_cases": False,
                "adversarial": True,
                "security": True,
            },
        },
        "pentest_config": {
            "categories": resolved_categories,
            "num_scenarios_per_category": num_scenarios_per_category,
            "difficulty": difficulty,
            "include_multi_turn": include_multi_turn,
        },
        "generation_options": {
            "generate_audio": False,
            "generate_files": False,
            "generate_simulations": False,
        },
    }

class PentestResult:
    """Result wrapper for pentest reports."""
    def __init__(self, report):
        self._report = report

    @property
    def grade(self): return self._report.overall_grade

    @property
    def risk_score(self): return self._report.overall_risk_score

    @property
    def findings(self):
        findings = []
        for cat in self._report.categories:
            findings.extend(cat.findings)
        return findings

    @property
    def attack_surface(self): return self._report.attack_surface

    @property
    def report(self): return self._report

    def summary(self):
        return {
            "grade": self.grade, "risk_score": self.risk_score,
            "total_findings": self._report.total_findings,
            "critical": self._report.critical_count, "high": self._report.high_count,
            "medium": self._report.medium_count, "low": self._report.low_count,
            "budget_limited": self._report.budget_limited,
            "duration_seconds": self._report.duration_seconds,
            "cost_usd": self._report.cost_usd,
        }

    def deploy(self, client, **kwargs):
        """Submit results to ashr platform."""
        return client._make_request(
            "run_pentest",
            agent_name=self._report.agent_name,
            mode=kwargs.get("mode", "standard"),
            report=self._report.model_dump(mode="json"),
            categories_tested=[c.category for c in self._report.categories],
        )

class PentestRunner:
    """Runs security pentests against AI agents.

    Example:
        runner = PentestRunner(mode="standard")
        result = runner.run(my_agent, tools=[...])
        print(result.grade)
        result.deploy(client)
    """
    def __init__(self, categories="all", mode="standard", max_budget_usd=None):
        if mode not in ("quick", "standard", "thorough"):
            raise ValueError("mode must be 'quick', 'standard', or 'thorough'")
        self.mode = mode
        self.max_budget_usd = max_budget_usd
        self.categories = None if categories == "all" else categories

    def run(self, agent, tools=None):
        _ensure_pentest_deps()
        from backend.pentest.orchestrator import PentestOrchestrator
        import anyio

        orch = PentestOrchestrator(
            agent=agent, tools=tools or [], mode=self.mode,
            categories=self.categories, max_budget_usd=self.max_budget_usd,
        )
        report = anyio.run(orch.run)
        report.agent_name = type(agent).__name__
        return PentestResult(report)
