"""
Request/response models for AMI SDK.

Defines action types, invoke options, and request/response structures.
"""
from dataclasses import dataclass
from datetime import datetime
from typing import Optional, List, Dict, Any, Literal, Union
from memrail.atoms import EventCorrelation, StateAtom, TagAtom, EventAtom


# ============================================================================
# Action Types
# ============================================================================

@dataclass
class ActionToolCall:
    """
    Tool call action - invoke an external function/tool.

    Attributes:
        type: Action type discriminator ("tool_call")
        tool: Tool provider (e.g., "mcp", "langchain")
        name: Function/tool name to invoke
        args: Arguments to pass to the tool
    """
    type: Literal["tool_call"]
    tool: str
    name: str
    args: Dict[str, Any]


@dataclass
class ActionDecisionPrompt:
    """
    Decision prompt action - send prompt to LLM for decision.

    Attributes:
        type: Action type discriminator ("decision_prompt")
        prompt: Prompt text to send to LLM
        model: Optional model identifier (e.g., "gpt-4", "claude-3")
        params: Optional model parameters (temperature, max_tokens, etc.)
    """
    type: Literal["decision_prompt"]
    prompt: str
    model: Optional[str] = None
    params: Optional[Dict[str, Any]] = None


@dataclass
class ActionRoute:
    """
    Route action - direct to a workflow/queue.

    Attributes:
        type: Action type discriminator ("route")
        route: Destination workflow/queue identifier
        metadata: Optional routing metadata
    """
    type: Literal["route"]
    route: str
    metadata: Optional[Dict[str, Any]] = None


@dataclass
class ActionContextDirective:
    """
    Context directive action - manipulate execution context.

    Attributes:
        type: Action type discriminator ("context_directive")
        directive: Directive text/instruction
        metadata: Optional directive metadata
    """
    type: Literal["context_directive"]
    directive: str
    metadata: Optional[Dict[str, Any]] = None


# Union type for all action variants
Action = Union[ActionToolCall, ActionDecisionPrompt, ActionRoute, ActionContextDirective]


# ============================================================================
# Invoke Options
# ============================================================================

@dataclass
class InvokeOptions:
    """
    Options for invoke request.

    Attributes:
        dry_run: If True, evaluate but don't execute actions
        top_k: Number of top-scoring EMUs to return (default: None - no limit)
        store_read_cutoff_ts: Only read state/events before this timestamp
        pin_registry_version: Pin to specific EMU registry version
        pin_policy_version: Pin to specific arbitration policy version
    """
    dry_run: bool = False
    top_k: Optional[int] = None
    store_read_cutoff_ts: Optional[datetime] = None
    pin_registry_version: Optional[str] = None
    pin_policy_version: Optional[str] = None


@dataclass
class TraceOptions:
    """
    Trace options for invoke request.

    Attributes:
        enable: If True, return detailed trace data in response
    """
    enable: bool = True


# ============================================================================
# Invoke Request
# ============================================================================

@dataclass
class InvokeRequest:
    """
    Request payload for decide/invoke endpoint.

    Attributes:
        context_ts: Timestamp for context evaluation (defaults to now)
        correlation: Correlation metadata for tracing
        context_atoms: Atoms providing context for EMU evaluation
        options: Invoke options (dry_run, top_k, etc.)
        trace: Trace options
        decision_point: Named decision point for analytics and topology mapping
    """
    context_ts: Optional[datetime] = None
    correlation: Optional[EventCorrelation] = None
    context_atoms: Optional[List[Union[StateAtom, TagAtom, EventAtom]]] = None
    options: InvokeOptions = None
    trace: TraceOptions = None
    decision_point: Optional[str] = None

    def __post_init__(self):
        # Ensure options and trace are initialized
        if self.options is None:
            self.options = InvokeOptions()
        if self.trace is None:
            self.trace = TraceOptions()


# ============================================================================
# Invoke Response
# ============================================================================

@dataclass
class SelectedItem:
    """
    A selected EMU from invoke response.

    Attributes:
        id: EMU ID (ULID)
        emu_key: EMU key (slug identifier)
        emu_version: EMU version number
        activation_id: Activation instance ID
        action: Action to execute (if EMU has action)
        score: EMU score from arbitration (0-1)
        resolved_idempotency_key: Resolved idempotency key for this activation
    """
    id: str
    emu_key: str
    emu_version: Optional[int] = None
    activation_id: Optional[str] = None
    action: Optional[Action] = None
    score: Optional[float] = None
    resolved_idempotency_key: Optional[str] = None


@dataclass
class IdempotencyMeta:
    """
    Idempotency metadata in invoke response.

    Attributes:
        applied: True if idempotency cache was applied (response from cache)
        invoke_key: Idempotency key used for this invocation
    """
    applied: bool
    invoke_key: Optional[str]


@dataclass
class InvokeResponse:
    """
    Response from invoke endpoint.

    Attributes:
        invocation_id: Unique invocation identifier
        idempotency: Idempotency metadata (if idempotency key was provided)
        selected: List of selected EMUs (length determined by top_k)
        trace: Trace data (if trace was enabled)
    """
    invocation_id: str
    idempotency: Optional[IdempotencyMeta]
    selected: List[SelectedItem]
    trace: Optional[Dict[str, Any]]
