"""Data types for the ARBI SDK."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, TypeAlias

if TYPE_CHECKING:
    from arbi_client.models.agent_step_event import AgentStepEvent
    from arbi_client.models.artifact_event import ArtifactEvent
    from arbi_client.models.user_input_request_event import UserInputRequestEvent
    from arbi_client.models.user_message_event import UserMessageEvent

#: Union of all typed SSE event models (ARBI extension events only;
#: OpenAI standard events are handled as raw dicts).
SseEvent: TypeAlias = "AgentStepEvent | UserInputRequestEvent | ArtifactEvent | UserMessageEvent"


@dataclass
class UserIdentity:
    """Client-side user identity containing keys derived from password."""

    email: str
    external_id: str
    given_name: str
    family_name: str
    signing_private_key: bytes  # Ed25519 private key (64 bytes)
    signing_public_key_b64: str  # Ed25519 public key (base64)
    encryption_private_key: bytes  # X25519 private key (32 bytes)
    encryption_public_key_b64: str  # X25519 public key (base64)


@dataclass
class WorkspaceInfo:
    """Workspace metadata returned at login (for auto-caching headers)."""

    external_id: str
    wrapped_key: str | None = None  # X25519-wrapped workspace key (base64)


@dataclass
class SessionState:
    """Server session state after login."""

    access_token: str
    session_public_key_b64: str  # X25519 session public key (base64)
    wrapped_keys: dict[str, str] = field(default_factory=dict)  # workspace_ext_id -> X25519-wrapped key from DB (lazy)
    workspace_headers: dict[str, str] = field(
        default_factory=dict
    )  # workspace_ext_id -> SealedBox-encrypted header (computed on demand)


@dataclass
class StreamingEvent:
    """A single SSE event with its typed payload and timing.

    Attributes:
        event_type: The SSE event name (``"response.created"``, ``"arbi.agent_step"``, etc.).
        data: Parsed typed model for known events, raw dict/str for unknown.
        elapsed: Seconds since the HTTP stream started (client-side measurement).
    """

    event_type: str
    data: SseEvent | dict[str, Any] | str
    elapsed: float | None = None


@dataclass
class QueryResult:
    """Result of a streaming agent query.

    Attributes:
        content: Concatenated answer text from token events.
        events: All SSE events as :class:`StreamingEvent` instances.
        user_input_events: ``user_input_request`` events (typed).
        metadata: Terminal ``message_metadata`` payload, if received.
        assistant_message_ext_id: The message ID assigned by the server.
    """

    content: str
    events: list[StreamingEvent] = field(default_factory=list)
    user_input_events: list[UserInputRequestEvent] = field(default_factory=list)
    metadata: dict[str, Any] | None = None
    assistant_message_ext_id: str | None = None
    error: str | None = None
