"""
Data models for the Ashr Labs SDK.

These are TypedDict definitions for better IDE support and documentation.
"""

from typing import TypedDict, Any
from datetime import datetime


class User(TypedDict, total=False):
    """Represents a user."""
    id: int
    created_at: str
    email: str
    name: str | None
    tenant: int
    is_active: bool


class Tenant(TypedDict, total=False):
    """Represents a tenant."""
    id: int
    created_at: str
    tenant_name: str
    is_active: bool


class Session(TypedDict):
    """Response from init() containing session information."""
    status: str
    user: User
    tenant: Tenant


class APIKey(TypedDict, total=False):
    """Represents an API key."""
    id: int
    key: str  # Only present on creation
    key_prefix: str
    name: str
    scopes: list[str]
    user_id: int
    tenant_id: int
    created_at: str
    last_used_at: str | None
    expires_at: str | None
    is_active: bool


class Dataset(TypedDict, total=False):
    """Represents a dataset."""
    id: int
    created_at: str
    tenant: int
    creator: int
    name: str
    description: str | None
    dataset_source: dict[str, Any]


class Run(TypedDict, total=False):
    """Represents a test run."""
    id: int
    created_at: str
    dataset: int
    tenant: int
    runner: int
    result: dict[str, Any]


class Request(TypedDict, total=False):
    """Represents a request."""
    id: int
    created_at: str
    requestor_id: int
    requestor_tenant: int
    request_name: str
    request_status: str
    request_input_schema: dict[str, Any] | None
    request: dict[str, Any]


class ListResponse(TypedDict):
    """Generic list response structure."""
    status: str
    total: int


class DatasetsListResponse(ListResponse):
    """Response from list_datasets."""
    datasets: list[Dataset]


class RunsListResponse(ListResponse):
    """Response from list_runs."""
    runs: list[Run]


class RequestsListResponse(ListResponse):
    """Response from list_requests."""
    requests: list[Request]


class APIKeysListResponse(ListResponse):
    """Response from list_api_keys."""
    api_keys: list[APIKey]


# =========================================================================
# Dataset structure types (used by EvalRunner)
# =========================================================================

class ToolCall(TypedDict, total=False):
    """A tool call within an expected response."""
    name: str
    arguments_json: str


class ExpectedResponse(TypedDict, total=False):
    """The expected agent response for a dataset action."""
    tool_calls: list[ToolCall]
    text: str


class Action(TypedDict, total=False):
    """A single action within a dataset scenario."""
    actor: str  # "user" or "agent"
    content: str
    name: str
    expected_response: ExpectedResponse


class Scenario(TypedDict, total=False):
    """A test scenario containing a sequence of actions."""
    title: str
    actions: list[Action]


# =========================================================================
# Pentest dataset types
# =========================================================================

PENTEST_CATEGORIES = [
    "prompt_injection",
    "privilege_escalation",
    "data_exfiltration",
    "jailbreak",
    "tool_abuse",
    "denial_of_service",
    "social_engineering",
]


class PentestDatasetConfig(TypedDict, total=False):
    """Configuration for requesting a pentest dataset from the API."""
    agent_name: str
    agent_description: str
    system_prompt: str
    tools: list[dict[str, Any]]
    categories: list[str]
    num_scenarios_per_category: int
    difficulty: str  # "easy", "medium", "hard", "mixed"
    include_multi_turn: bool
