"""
Ashr Labs API Client.

This module provides the main client for interacting with the Ashr Labs API.
"""

import copy
import json
import time
from typing import Any
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

from .exceptions import (
    AshrLabsError,
    AuthenticationError,
    AuthorizationError,
    NotFoundError,
    ValidationError,
    RateLimitError,
    ServerError,
)
from .models import Dataset, Run, Request as RequestModel, APIKey, Session


_DEFAULT_BASE_URL = "https://api.ashr.io/testing-platform-api"


class AshrLabsClient:
    """
    Client for interacting with the Ashr Labs API.

    This client supports all API key accessible endpoints:
    - Datasets: get, list
    - Runs: create, get, list, delete
    - Requests: create, get, list
    - API Keys: list, revoke

    The client automatically resolves your ``tenant_id`` and ``user_id``
    from the API key on first use, so you never need to pass them manually.

    Example:
        ```python
        from ashr_labs import AshrLabsClient

        client = AshrLabsClient(api_key="tp_your_api_key_here")

        # No tenant_id needed — auto-resolved from your API key
        datasets = client.list_datasets()

        # Create a run
        run = client.create_run(
            dataset_id=42,
            result={"score": 0.95, "status": "passed"}
        )
        ```
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = _DEFAULT_BASE_URL,
        timeout: int = 30,
    ):
        """
        Initialize the Ashr Labs client.

        Args:
            api_key: Your Ashr Labs API key (starts with 'tp_').
            base_url: The base URL of the API (default: production URL).
            timeout: Request timeout in seconds (default: 30).

        Raises:
            ValueError: If the API key format is invalid.
        """
        if not api_key or not api_key.startswith("tp_"):
            raise ValueError("Invalid API key format. API keys must start with 'tp_'")

        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

        # Lazy-loaded session context (populated on first call that needs it)
        self._tenant_id: int | None = None
        self._user_id: int | None = None
        self._session: dict | None = None

    @classmethod
    def from_env(cls, timeout: int = 30) -> "AshrLabsClient":
        """
        Create a client from environment variables.

        Reads ``ASHR_LABS_API_KEY`` (required) and ``ASHR_LABS_BASE_URL``
        (optional, defaults to production) from the environment.

        Returns:
            An authenticated AshrLabsClient.

        Raises:
            RuntimeError: If ASHR_LABS_API_KEY is not set.
        """
        import os

        api_key = os.environ.get("ASHR_LABS_API_KEY")
        if not api_key:
            raise RuntimeError(
                "ASHR_LABS_API_KEY environment variable is not set. "
                "Create an API key at https://app.ashr.io → API Keys."
            )
        base_url = os.environ.get("ASHR_LABS_BASE_URL", _DEFAULT_BASE_URL)
        return cls(api_key=api_key, base_url=base_url, timeout=timeout)

    def _ensure_session(self) -> None:
        """Lazily call init() and cache tenant_id / user_id."""
        if self._session is not None:
            return
        self._session = self._make_request("init")
        self._tenant_id = self._session["tenant"]["id"]
        self._user_id = self._session["user"]["id"]

    def _resolve_tenant_id(self, tenant_id: int | None) -> int:
        """Return the explicit tenant_id, or auto-resolve from session."""
        if tenant_id is not None:
            return tenant_id
        self._ensure_session()
        return self._tenant_id  # type: ignore[return-value]

    def _resolve_user_id(self, user_id: int | None) -> int:
        """Return the explicit user_id, or auto-resolve from session."""
        if user_id is not None:
            return user_id
        self._ensure_session()
        return self._user_id  # type: ignore[return-value]

    def _make_request(self, function: str, **params: Any) -> dict:
        """
        Make an authenticated request to the API.

        Args:
            function: The API function to call.
            **params: Additional parameters for the function.

        Returns:
            The JSON response from the API.

        Raises:
            AshrLabsError: If the request fails.
        """
        payload = {"function": function, **params}

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }

        data = json.dumps(payload).encode("utf-8")
        request = Request(self.base_url, data=data, headers=headers, method="POST")

        try:
            with urlopen(request, timeout=self.timeout) as response:
                response_data = json.loads(response.read().decode("utf-8"))
        except HTTPError as e:
            status_code = e.code
            try:
                error_body = json.loads(e.read().decode("utf-8"))
                message = error_body.get("message", str(e))
            except (json.JSONDecodeError, AttributeError):
                message = str(e)

            self._raise_for_status(status_code, message, error_body if "error_body" in dir() else None)
        except URLError as e:
            raise AshrLabsError(f"Network error: {e.reason}")
        except TimeoutError:
            raise AshrLabsError("Request timed out")

        # Check for API-level errors
        if response_data.get("status") == "error":
            message = response_data.get("message", "Unknown error")
            self._raise_for_status(400, message, response_data)

        return response_data

    def _raise_for_status(
        self, status_code: int, message: str, response: dict | None = None
    ) -> None:
        """Raise the appropriate exception based on status code."""
        if status_code == 401:
            raise AuthenticationError(message, status_code, response)
        elif status_code == 403:
            raise AuthorizationError(message, status_code, response)
        elif status_code == 404:
            raise NotFoundError(message, status_code, response)
        elif status_code == 422:
            raise ValidationError(message, status_code, response)
        elif status_code == 429:
            raise RateLimitError(message, status_code, response)
        elif status_code >= 500:
            raise ServerError(message, status_code, response)
        else:
            raise AshrLabsError(message, status_code, response)

    # =========================================================================
    # Dataset Operations
    # =========================================================================

    def get_dataset(
        self,
        dataset_id: int,
        include_signed_urls: bool = False,
        url_expires_seconds: int = 3600,
    ) -> Dataset:
        """
        Retrieve a dataset by ID.

        Args:
            dataset_id: The ID of the dataset to retrieve.
            include_signed_urls: Whether to include signed S3 URLs for media files.
            url_expires_seconds: Expiration time for signed URLs in seconds.

        Returns:
            The dataset object.

        Raises:
            NotFoundError: If the dataset is not found.
            AuthorizationError: If you don't have access to this dataset.
        """
        response = self._make_request(
            "get_dataset",
            dataset_id=dataset_id,
            include_signed_urls=include_signed_urls,
            url_expires_seconds=url_expires_seconds,
        )
        return response["dataset"]

    def list_datasets(
        self,
        tenant_id: int | None = None,
        limit: int = 50,
        cursor: int | None = None,
        include_signed_urls: bool = False,
        url_expires_seconds: int = 3600,
    ) -> dict:
        """
        List datasets for a tenant.

        Args:
            tenant_id: The ID of the tenant (auto-resolved if omitted).
            limit: Maximum number of datasets to return (default: 50).
            cursor: ID of last item from previous page for cursor-based pagination.
                Pass the ``next_cursor`` value from the previous response.
            include_signed_urls: Whether to include signed S3 URLs for media files.
            url_expires_seconds: Expiration time for signed URLs in seconds.

        Returns:
            A dict containing 'datasets' list and 'next_cursor' for pagination.
        """
        params: dict[str, Any] = {
            "tenant_id": self._resolve_tenant_id(tenant_id),
            "limit": limit,
            "include_signed_urls": include_signed_urls,
            "url_expires_seconds": url_expires_seconds,
        }
        if cursor is not None:
            params["cursor"] = cursor
        return self._make_request("list_datasets", **params)

    # =========================================================================
    # Run Operations
    # =========================================================================

    def create_run(
        self,
        dataset_id: int,
        result: dict[str, Any],
        tenant_id: int | None = None,
        runner_id: int | None = None,
    ) -> Run:
        """
        Create a new test run.

        Args:
            dataset_id: The ID of the dataset this run is for.
            result: The run result data (metrics, status, etc.).
            tenant_id: The ID of the tenant (auto-resolved if omitted).
            runner_id: Optional ID of the user who ran the test.

        Returns:
            The created run object.

        Example:
            ```python
            run = client.create_run(
                dataset_id=42,
                result={
                    "score": 0.95,
                    "status": "passed",
                    "metrics": {"accuracy": 0.98, "latency_ms": 150}
                }
            )
            ```
        """
        params: dict[str, Any] = {
            "tenant_id": self._resolve_tenant_id(tenant_id),
            "dataset_id": dataset_id,
            "result": result,
        }
        if runner_id is not None:
            params["runner_id"] = runner_id

        response = self._make_request("create_run", **params)
        return response["run"]

    def delete_run(self, run_id: int) -> dict:
        """
        Delete a test run.

        Args:
            run_id: The ID of the run to delete.

        Returns:
            Confirmation of deletion.

        Raises:
            NotFoundError: If the run is not found.
        """
        return self._make_request("delete_run", run_id=run_id)

    def get_run(self, run_id: int) -> Run:
        """
        Retrieve a run by ID.

        Args:
            run_id: The ID of the run to retrieve.

        Returns:
            The run object.

        Raises:
            NotFoundError: If the run is not found.
        """
        response = self._make_request("get_run", run_id=run_id)
        return response["run"]

    def list_runs(
        self,
        dataset_id: int | None = None,
        tenant_id: int | None = None,
        limit: int = 50,
    ) -> dict:
        """
        List runs for a tenant or dataset.

        Args:
            dataset_id: Filter by dataset ID.
            tenant_id: Filter by tenant ID (auto-resolved if omitted).
            limit: Maximum number of runs to return.

        Returns:
            A dict containing 'runs' list.
        """
        params: dict[str, Any] = {
            "tenant_id": self._resolve_tenant_id(tenant_id),
            "limit": limit,
        }
        if dataset_id is not None:
            params["dataset_id"] = dataset_id

        return self._make_request("list_runs", **params)

    # =========================================================================
    # Request Operations
    # =========================================================================

    def create_request(
        self,
        request_name: str,
        request: dict[str, Any],
        request_input_schema: dict[str, Any] | None = None,
        tenant_id: int | None = None,
        requestor_id: int | None = None,
    ) -> RequestModel:
        """
        Create a new dataset generation request.

        The ``request`` dict contains the generation config that describes
        your agent and the test scenarios to generate.  The
        ``request_input_schema`` is a JSON Schema that the backend uses to
        validate that config; if you include agent tools, put them in
        ``request_input_schema["tools"]`` so they're auto-saved as skill
        templates.

        Args:
            request_name: A name/title for the request.
            request: The generation config (see full schema below).
            request_input_schema: JSON Schema for validating ``request``.
                If omitted a permissive default schema is sent and tools
                are auto-populated from ``request["agent"]["tools"]``.
            tenant_id: The ID of the tenant (auto-resolved if omitted).
            requestor_id: The ID of the user making the request (auto-resolved if omitted).

        Returns:
            The created request object.

        **Generation config structure** (the ``request`` dict)::

            {
                # ── metadata (optional) ──────────────────────────────
                "metadata": {
                    "dataset_name": str,       # Name for the generated dataset
                    "description": str,        # Description of the dataset
                },

                # ── agent (required) ─────────────────────────────────
                # At least one of name, description, or system_prompt is required.
                "agent": {
                    "name": str,               # Agent name
                    "description": str,        # What the agent does (min 20 chars recommended)
                    "system_prompt": str,      # System prompt given to the agent

                    # Tools the agent can call
                    "tools": [
                        {
                            "name": str,           # Tool name (snake_case)
                            "description": str,    # What the tool does
                            "parameters": {        # JSON Schema for tool params
                                "type": "object",
                                "properties": {
                                    "param_name": {"type": str, "description": str},
                                },
                                "required": [str, ...],
                            },
                            "returns": {           # (optional) Return value schema
                                "type": str,
                                "description": str,
                            },
                        },
                    ],

                    # What input modalities the agent accepts
                    # Values can be bool or {"enabled": bool}
                    "accepted_inputs": {
                        "text": True,              # (default True) Text input
                        "audio": False,            # Audio input (mp3, wav, m4a, ogg, webm)
                        "file": False,             # File input (pdf, txt, csv, json, xml, etc.)
                        "image": False,            # Image input (jpg, png, gif, webp)
                        "video": False,            # Video input (mp4, webm, mov, avi)
                        "conversation": False,     # Multi-participant conversations
                    },

                    # Output format
                    "output_format": {
                        "type": "text",            # "text" or "json"/"structured"
                        "schema": {...},           # (optional) JSON schema if type is structured
                    },

                    # Custom structured input schema (optional)
                    "input_schema": {
                        "name": str,               # Schema name (e.g., "OrderInput")
                        "description": str,        # What this input represents
                        "fields": [
                            {
                                "name": str,       # Field name
                                "type": str,       # string, number, integer, boolean, array, object
                                "description": str,
                                "required": True,  # (default True)
                                "enum": [...],     # (optional) Allowed values
                                "default": ...,    # (optional) Default value
                            },
                        ],
                        "example": {...},          # (optional) Example input object
                    },
                },

                # ── context (required) ───────────────────────────────
                # At least one of domain, use_case, or scenario_context is required.
                "context": {
                    "domain": str,             # Domain: "banking", "healthcare", "e-commerce",
                                               #   "legal", "education", "customer_service",
                                               #   "technology", "travel", "insurance", "other"
                    "use_case": str,           # Specific use case (min 10 chars recommended)
                    "scenario_context": str,   # Additional scenario context

                    "user_persona": {          # (optional) Who is interacting with the agent
                        "type": str,           # Persona type/role
                        "description": str,    # Persona description
                    },

                    "sample_data": {           # (optional) Example data for realistic tests
                        "examples": [{...}],   # List of example data objects
                    },
                },

                # ── test_config (optional, defaults provided) ────────
                "test_config": {
                    "num_variations": 5,           # 1-50, number of test scenarios
                    "strategy": "balanced",        # "focused", "diverse", or "balanced"

                    "coverage": {                  # Which aspects to test
                        "happy_path": True,        # Standard success scenarios
                        "edge_cases": True,        # Unusual/boundary inputs
                        "error_handling": True,    # Error and failure scenarios
                        "boundary_values": True,   # Boundary value testing
                    },

                    "complexity_distribution": {   # Must sum to ~1.0
                        "simple": 0.3,             # Simple scenarios (0-1)
                        "moderate": 0.5,           # Medium complexity (0-1)
                        "complex": 0.2,            # Complex scenarios (0-1)
                    },

                    "focus_areas": [str, ...],     # Specific areas to focus on
                    "exclude": [str, ...],         # Scenarios to exclude
                },

                # ── expected_behaviors (optional) ────────────────────
                "expected_behaviors": {
                    "must_include": [str, ...],    # Keywords that must appear in response
                    "must_not_include": [str, ...],# Keywords that must NOT appear
                    "expected_tools": [str, ...],  # Tools expected to be called
                },

                # ── generation_options (optional) ────────────────────
                "generation_options": {
                    "generate_audio": False,       # Generate audio test inputs
                    "generate_files": False,       # Generate file test inputs
                    "generate_images": False,      # Generate image test inputs
                    "generate_videos": False,      # Generate video test inputs
                    "generate_simulations": False,  # Generate website simulation videos
                },

                # ── website_simulation (optional) ────────────────────
                "website_simulation": {
                    "enabled": False,
                    "url": str,                    # URL of the website
                    "domain": str,                 # Domain category
                    "simulation_config": {
                        "html_content": str,       # Custom HTML (auto-generated if empty)
                        "cursor_size": int,        # Cursor diameter in px
                        "cursor_color": str,       # CSS color (e.g., "#EF4444")
                        "click_ring_color": str,   # Click ring CSS color
                        "viewport_width": 1280,    # Viewport width in px
                        "viewport_height": 720,    # Viewport height in px
                        "typing_delay_ms": 80,     # Delay between keystrokes
                        "video_fps": 25,           # Video frames per second
                    },
                },
            }

        Example:
            ```python
            req = client.create_request(
                request_name="Support Agent Eval",
                request={
                    "metadata": {"dataset_name": "Support Eval"},
                    "agent": {
                        "name": "Support Bot",
                        "description": "Answers customer questions about orders and refunds",
                        "system_prompt": "You are a helpful support agent.",
                        "tools": [
                            {
                                "name": "lookup_order",
                                "description": "Look up an order by ID",
                                "parameters": {
                                    "type": "object",
                                    "properties": {
                                        "order_id": {"type": "string", "description": "The order ID"}
                                    },
                                    "required": ["order_id"],
                                },
                            },
                        ],
                        "accepted_inputs": {"text": True, "audio": False, "file": False},
                        "output_format": {"type": "text"},
                    },
                    "context": {
                        "domain": "e-commerce",
                        "use_case": "Customers asking about order status, requesting refunds",
                        "scenario_context": "An online retail store called ShopWave",
                        "user_persona": {"type": "customer", "description": "Online shoppers"},
                    },
                    "test_config": {
                        "num_variations": 10,
                        "strategy": "diverse",
                        "coverage": {
                            "happy_path": True,
                            "edge_cases": True,
                            "error_handling": True,
                        },
                        "complexity_distribution": {"simple": 0.3, "moderate": 0.5, "complex": 0.2},
                    },
                    "expected_behaviors": {
                        "must_include": ["order"],
                        "expected_tools": ["lookup_order"],
                    },
                },
            )
            ```
        """
        self._validate_config_structure(request)

        # The backend requires request_input_schema; send a permissive
        # default so callers aren't forced to build one manually.
        if request_input_schema is None:
            request_input_schema = {
                "type": "object",
                "properties": {},
            }
            # Auto-populate tools from the request's agent config
            tools = (request.get("agent") or {}).get("tools")
            if tools:
                request_input_schema["tools"] = tools

        params: dict[str, Any] = {
            "tenant_id": self._resolve_tenant_id(tenant_id),
            "requestor_id": self._resolve_user_id(requestor_id),
            "request_name": request_name,
            "request": request,
            "request_input_schema": request_input_schema,
        }

        response = self._make_request("create_request", **params)
        return response["request"]

    def get_request(self, request_id: int) -> RequestModel:
        """
        Retrieve a request by ID.

        Args:
            request_id: The ID of the request to retrieve.

        Returns:
            The request object.

        Raises:
            NotFoundError: If the request is not found.
        """
        response = self._make_request("get_request", request_id=request_id)
        return response["request"]

    def list_requests(
        self,
        tenant_id: int | None = None,
        status: str | None = None,
        limit: int = 50,
        cursor: int | None = None,
    ) -> dict:
        """
        List requests for a tenant.

        Args:
            tenant_id: The ID of the tenant (auto-resolved if omitted).
            status: Filter by request status (e.g., 'pending', 'completed').
            limit: Maximum number of requests to return.
            cursor: ID of last item from previous page for cursor-based pagination.
                Pass the ``next_cursor`` value from the previous response.

        Returns:
            A dict containing 'requests' list and 'next_cursor' for pagination.
        """
        params: dict[str, Any] = {
            "tenant_id": self._resolve_tenant_id(tenant_id),
            "limit": limit,
        }
        if status is not None:
            params["status"] = status
        if cursor is not None:
            params["cursor"] = cursor

        return self._make_request("list_requests", **params)

    # =========================================================================
    # API Key Operations
    # =========================================================================

    def list_api_keys(self, include_inactive: bool = False) -> list[APIKey]:
        """
        List API keys for your tenant.

        Note: For security, only the key prefix is returned, not the full key.

        Args:
            include_inactive: Whether to include revoked/inactive keys.

        Returns:
            A list of API key objects.
        """
        response = self._make_request(
            "list_api_keys",
            include_inactive=include_inactive,
        )
        return response["api_keys"]

    def revoke_api_key(self, api_key_id: int) -> dict:
        """
        Revoke an API key.

        Args:
            api_key_id: The ID of the API key to revoke.

        Returns:
            Confirmation of revocation.

        Raises:
            NotFoundError: If the API key is not found.
        """
        return self._make_request("revoke_api_key", api_key_id=api_key_id)

    # =========================================================================
    # Session Operations
    # =========================================================================

    def init(self) -> Session:
        """
        Initialize a session and validate authentication.

        This method validates the API key and returns information about
        the authenticated user and tenant. Use this to verify credentials
        and retrieve user/tenant context at the start of a session.

        Returns:
            A dict containing:
            - status: "ok" if successful
            - user: User information (id, email, name, etc.)
            - tenant: Tenant information (id, tenant_name, etc.)

        Raises:
            AuthenticationError: If the API key is invalid or expired.

        Example:
            ```python
            client = AshrLabsClient(api_key="tp_...")

            # Validate credentials and get user info
            session = client.init()
            print(f"Logged in as: {session['user']['email']}")
            print(f"Tenant: {session['tenant']['tenant_name']}")
            ```
        """
        return self._make_request("init")

    # =========================================================================
    # Utility Methods
    # =========================================================================

    def health_check(self) -> dict:
        """
        Check if the API is reachable and responding.

        Returns:
            A dict with status information.
        """
        return self._make_request("keep_alive")

    # =========================================================================
    # Convenience / Workflow Methods
    # =========================================================================

    def wait_for_request(
        self,
        request_id: int,
        timeout: int = 600,
        poll_interval: int = 5,
    ) -> RequestModel:
        """Block until a request reaches a terminal state (completed/failed).

        Args:
            request_id: The ID of the request to poll.
            timeout: Maximum seconds to wait before raising TimeoutError.
            poll_interval: Seconds between polls.

        Returns:
            The final request object.

        Raises:
            TimeoutError: If the request doesn't finish within ``timeout`` seconds.
            AshrLabsError: If the request fails.
        """
        start = time.time()
        while time.time() - start < timeout:
            req = self.get_request(request_id=request_id)
            status = req.get("request_status", "")
            if status == "completed":
                return req
            if status == "failed":
                error_msg = req.get("error", "Request failed")
                raise AshrLabsError(f"Request {request_id} failed: {error_msg}")
            time.sleep(poll_interval)
        raise TimeoutError(
            f"Request {request_id} did not complete within {timeout}s"
        )

    def poll_run(
        self,
        run_id: int,
        timeout: int = 300,
        poll_interval: int = 20,
        on_poll: Any | None = None,
    ) -> Run:
        """Block until backend grading completes for a run.

        After ``deploy()``, the backend grades tool arguments and text
        responses asynchronously (typically 1-3 minutes). This method
        polls ``get_run()`` until ``aggregate_metrics.tests_passed`` is
        populated, meaning grading is finished.

        Args:
            run_id: The ID of the run to poll.
            timeout: Maximum seconds to wait (default: 300).
            poll_interval: Seconds between polls (default: 20).
            on_poll: Optional callback ``(elapsed_seconds, run_dict) -> None``
                called after each poll.

        Returns:
            The fully graded run object.

        Raises:
            TimeoutError: If grading doesn't finish within ``timeout`` seconds.
        """
        start = time.time()
        while True:
            run = self.get_run(run_id)
            metrics = run.get("result", {}).get("aggregate_metrics", {})
            if metrics.get("tests_passed") is not None:
                return run
            elapsed = time.time() - start
            if on_poll:
                on_poll(int(elapsed), run)
            if elapsed >= timeout:
                raise TimeoutError(
                    f"Run {run_id} grading did not complete within {timeout}s"
                )
            time.sleep(poll_interval)

    @staticmethod
    def _enrich_config(config: dict[str, Any]) -> dict[str, Any]:
        """Fill in missing context fields so the backend has enough to generate.

        If ``context.use_case`` and ``context.scenario_context`` are both
        missing, synthesize them from the agent's name and description so
        the generation prompt is meaningful.
        """
        config = copy.deepcopy(config)
        agent = config.get("agent", {})
        context = config.get("context", {})

        has_use_case = bool(context.get("use_case"))
        has_scenario = bool(context.get("scenario_context"))

        if not has_use_case and not has_scenario:
            name = agent.get("name", "")
            desc = agent.get("description", "")
            domain = context.get("domain", "")

            if desc:
                context["use_case"] = desc
            elif name:
                context["use_case"] = f"Testing the {name} agent"

            if name and desc:
                parts = [f"A user interacting with {name}"]
                if domain and domain != "general":
                    parts.append(f"in the {domain} domain")
                parts.append(f"— {desc}")
                context["scenario_context"] = " ".join(parts)

            config["context"] = context

        # Default test_config if missing
        if "test_config" not in config:
            config["test_config"] = {
                "num_variations": 5,
                "coverage": {
                    "happy_path": True,
                    "edge_cases": True,
                    "error_handling": True,
                },
            }

        return config

    def generate_dataset(
        self,
        request_name: str,
        config: dict[str, Any],
        request_input_schema: dict[str, Any] | None = None,
        timeout: int = 600,
        poll_interval: int = 5,
    ) -> tuple[int, dict]:
        """Create a dataset generation request, wait for it, and fetch the result.

        **Prefer reusing existing datasets** with ``EvalRunner.from_dataset()``
        instead of generating new ones each time. Only generate a new dataset
        when the agent's tools, inputs, or domain have changed. Reusing
        datasets lets you compare agent performance across code changes
        against the same scenarios.

        This is a high-level convenience method that combines
        ``create_request`` + ``wait_for_request`` + ``get_dataset``.

        Missing context fields (``use_case``, ``scenario_context``) are
        auto-filled from the agent's name and description so the backend
        always has enough information to generate meaningful scenarios.
        A default ``test_config`` is added if not provided.

        Args:
            request_name: A name/title for the request.
            config: The generation config dict — same structure as
                ``create_request``'s ``request`` parameter. See
                ``create_request`` for the full schema reference.

                Required sections:
                    - **agent**: At least one of ``name``, ``description``, or ``system_prompt``.
                      Optionally include ``tools``, ``accepted_inputs``, ``output_format``,
                      ``input_schema``, ``system_prompt``.
                    - **context**: At least one of ``domain``, ``use_case``, or ``scenario_context``.
                      Optionally include ``user_persona``, ``sample_data``.

                Optional sections:
                    - **metadata**: ``dataset_name``, ``description``
                    - **test_config**: ``num_variations`` (1-50), ``strategy`` (focused/diverse/balanced),
                      ``coverage`` (happy_path, edge_cases, error_handling, boundary_values),
                      ``complexity_distribution`` (simple, moderate, complex), ``focus_areas``, ``exclude``
                    - **expected_behaviors**: ``must_include``, ``must_not_include``, ``expected_tools``
                    - **generation_options**: ``generate_audio``, ``generate_files``, ``generate_images``,
                      ``generate_videos``, ``generate_simulations``
                    - **website_simulation**: ``enabled``, ``url``, ``domain``, ``simulation_config``

            request_input_schema: Optional JSON Schema for validation.
                If omitted, auto-populated from ``config["agent"]["tools"]``.
            timeout: Maximum seconds to wait for generation (default: 600).
            poll_interval: Seconds between status polls (default: 5).

        Returns:
            A tuple of ``(dataset_id, dataset_source)`` where
            ``dataset_source`` is the dict containing ``"runs"``.

        Raises:
            TimeoutError: If generation doesn't finish in time.
            AshrLabsError: If generation fails.

        Example:
            ```python
            dataset_id, source = client.generate_dataset(
                "My Agent Eval",
                {
                    "agent": {
                        "name": "Support Bot",
                        "description": "Handles customer orders and refunds",
                        "tools": [
                            {"name": "lookup_order", "description": "Look up an order",
                             "parameters": {"type": "object", "properties": {"order_id": {"type": "string"}}, "required": ["order_id"]}},
                        ],
                    },
                    "context": {
                        "domain": "e-commerce",
                        "use_case": "Customer support for online retail",
                    },
                    "test_config": {"num_variations": 10, "strategy": "diverse"},
                },
            )
            ```
        """
        config = self._enrich_config(config)
        req = self.create_request(
            request_name=request_name,
            request=config,
            request_input_schema=request_input_schema,
        )
        request_id = req["id"]

        self.wait_for_request(
            request_id, timeout=timeout, poll_interval=poll_interval
        )

        # Find the dataset created by this request
        resp = self.list_datasets(limit=10)
        datasets = resp.get("datasets", [])
        match = next(
            (d for d in datasets if d.get("request_id") == request_id),
            None,
        )
        if not match:
            raise AshrLabsError(
                f"No dataset found for request {request_id}"
            )

        dataset_id = match["id"]
        full_ds = self.get_dataset(dataset_id=dataset_id, include_signed_urls=False)
        source = full_ds.get("dataset_source", {})
        return dataset_id, source

    # =========================================================================
    # Pentest Dataset Operations
    # =========================================================================

    def request_pentest_dataset(
        self,
        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,
        tenant_id: int | None = None,
        requestor_id: int | None = None,
    ) -> RequestModel:
        """Create a pentest dataset generation request.

        Submits a request to generate a security evaluation dataset
        tailored to your agent. The generated dataset will contain
        adversarial attack scenarios across the specified security
        categories.

        To wait for the dataset to be generated and fetch it, use
        :meth:`generate_pentest_dataset` instead.

        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 attacks).
            tools: Tool definitions the agent has access to.
            categories: Security categories to test. Defaults to all 7:
                prompt_injection, privilege_escalation, data_exfiltration,
                jailbreak, tool_abuse, denial_of_service, social_engineering.
            num_scenarios_per_category: Scenarios per category (default: 3).
            difficulty: "easy", "medium", "hard", or "mixed" (default).
            include_multi_turn: Include multi-turn attacks (default: True).
            tenant_id: Tenant ID (auto-resolved if omitted).
            requestor_id: Requestor user ID (auto-resolved if omitted).

        Returns:
            The created request object. Use ``request["id"]`` to poll
            status with :meth:`get_request` or :meth:`wait_for_request`.

        Example:
            ```python
            req = client.request_pentest_dataset(
                agent_name="Support Bot",
                agent_description="Answers customer support questions",
                system_prompt="You are a helpful support agent.",
                tools=[{"name": "lookup_order", "description": "Look up an order"}],
                categories=["prompt_injection", "data_exfiltration"],
            )
            print(f"Request created: {req['id']}")

            # Poll until complete
            finished = client.wait_for_request(req["id"])
            ```
        """
        from .pentest import build_pentest_request_config

        config = build_pentest_request_config(
            agent_name=agent_name,
            agent_description=agent_description,
            system_prompt=system_prompt,
            tools=tools,
            categories=categories,
            num_scenarios_per_category=num_scenarios_per_category,
            difficulty=difficulty,
            include_multi_turn=include_multi_turn,
        )

        # Build schema with tools for auto-saving as skill templates
        request_input_schema: dict[str, Any] = {
            "type": "object",
            "properties": {
                "agent": {"type": "object"},
                "context": {"type": "object"},
                "pentest_config": {"type": "object"},
            },
            "required": ["agent", "context", "pentest_config"],
        }
        if tools:
            request_input_schema["tools"] = tools

        return self.create_request(
            request_name=f"Pentest: {agent_name}",
            request=config,
            request_input_schema=request_input_schema,
            tenant_id=tenant_id,
            requestor_id=requestor_id,
        )

    def generate_pentest_dataset(
        self,
        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,
        timeout: int = 600,
        poll_interval: int = 5,
    ) -> tuple[int, dict]:
        """Request a pentest dataset, wait for generation, and return it.

        This is the high-level convenience method that combines
        :meth:`request_pentest_dataset` + :meth:`wait_for_request` +
        :meth:`get_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.
            tools: Tool definitions the agent has access to.
            categories: Security categories to test (default: all 7).
            num_scenarios_per_category: Scenarios per category (default: 3).
            difficulty: "easy", "medium", "hard", or "mixed" (default).
            include_multi_turn: Include multi-turn attacks (default: True).
            timeout: Max seconds to wait for generation (default: 600).
            poll_interval: Seconds between polls (default: 5).

        Returns:
            A tuple of ``(dataset_id, dataset_source)`` where
            ``dataset_source`` contains the pentest scenarios in the
            ``multi_run_storyboard`` format with ``"runs"``.

        Raises:
            TimeoutError: If generation doesn't finish in time.
            AshrLabsError: If generation fails.

        Example:
            ```python
            dataset_id, source = client.generate_pentest_dataset(
                agent_name="Support Bot",
                agent_description="Answers customer support questions",
                system_prompt="You are a helpful support agent.",
                tools=[{"name": "lookup_order", "description": "Look up an order"}],
                categories=["prompt_injection", "data_exfiltration"],
                num_scenarios_per_category=5,
            )

            print(f"Dataset {dataset_id} with {len(source['runs'])} scenarios")

            # Run the dataset against your agent
            from ashr_labs import EvalRunner
            runner = EvalRunner(dataset_source=source)
            result = runner.run(my_agent)
            result.deploy(client, dataset_id=dataset_id)
            ```
        """
        req = self.request_pentest_dataset(
            agent_name=agent_name,
            agent_description=agent_description,
            system_prompt=system_prompt,
            tools=tools,
            categories=categories,
            num_scenarios_per_category=num_scenarios_per_category,
            difficulty=difficulty,
            include_multi_turn=include_multi_turn,
        )
        request_id = req["id"]

        self.wait_for_request(
            request_id, timeout=timeout, poll_interval=poll_interval
        )

        # Find the dataset created by this request
        resp = self.list_datasets(limit=10)
        datasets = resp.get("datasets", [])
        match = next(
            (d for d in datasets if d.get("request_id") == request_id),
            None,
        )
        if not match:
            raise AshrLabsError(
                f"No dataset found for request {request_id}"
            )

        dataset_id = match["id"]
        full_ds = self.get_dataset(dataset_id=dataset_id, include_signed_urls=False)
        source = full_ds.get("dataset_source", {})
        return dataset_id, source

    @staticmethod
    def _validate_config_structure(config: dict) -> None:
        """Lightweight pre-check that required gen_config sections are present.

        Raises ValueError with a clear message if ``agent`` or ``context``
        is missing or empty.  This catches obvious mistakes before the
        request hits the network.

        For pentest configs (identified by ``metadata.dataset_type == "pentest"``),
        the ``pentest_config`` section is also validated.
        """
        if not isinstance(config, dict):
            raise ValueError("request config must be a dict")

        # --- agent ---
        agent = config.get("agent")
        if not agent or not isinstance(agent, dict):
            raise ValueError(
                "request config must include an 'agent' section"
            )
        if not any(agent.get(k) for k in ("name", "description", "system_prompt")):
            raise ValueError(
                "agent must have at least one of: name, description, system_prompt"
            )

        # --- context ---
        context = config.get("context")
        if not context or not isinstance(context, dict):
            raise ValueError(
                "request config must include a 'context' section"
            )
        if not any(context.get(k) for k in ("domain", "use_case", "scenario_context")):
            raise ValueError(
                "context must have at least one of: domain, use_case, scenario_context"
            )

        # --- pentest_config (optional, for pentest datasets) ---
        metadata = config.get("metadata", {})
        if metadata.get("dataset_type") == "pentest":
            from .models import PENTEST_CATEGORIES

            pentest_config = config.get("pentest_config")
            if not pentest_config or not isinstance(pentest_config, dict):
                raise ValueError(
                    "pentest dataset config must include a 'pentest_config' section"
                )
            cats = pentest_config.get("categories", [])
            invalid = set(cats) - set(PENTEST_CATEGORIES)
            if invalid:
                raise ValueError(
                    f"Invalid pentest categories: {invalid}. "
                    f"Valid: {PENTEST_CATEGORIES}"
                )

    def __repr__(self) -> str:
        return f"AshrLabsClient(base_url='{self.base_url}', api_key='{self.api_key[:8]}...')"
