"""Stateful session management for Anthropic Responses API.

This module provides session storage and management for multi-turn
conversations using the Responses API. Sessions track message history,
tool definitions, and sampling parameters for conversation continuity.
"""

from __future__ import annotations

import threading
from dataclasses import dataclass, field
from typing import Iterable, Mapping, MutableMapping, Sequence

import anthropic

from ember._internal.exceptions import ProviderAPIError
from keiro.gateway.domain.model import StreamingResponse
from ember.models.providers.anthropic import AnthropicProvider
from ember.models.providers.anthropic_responses import (
    AnthropicParams,
    clone_messages,
    stream_events_from_anthropic,
)
from ember.models.providers.utils import ResponsesRequest

MessagePayload = Mapping[str, object]


@dataclass
class SamplingConfig:
    """Sampling parameters for Responses API calls."""

    temperature: float | None = None
    top_p: float | None = None
    stop_sequences: tuple[str, ...] = tuple()
    thinking: Mapping[str, object] | None = None
    output_config: Mapping[str, object] | None = None

    @classmethod
    def from_params(cls, params: AnthropicParams) -> "SamplingConfig":
        """Create a SamplingConfig from AnthropicParams."""
        return cls(
            temperature=params.temperature,
            top_p=params.top_p,
            stop_sequences=tuple(params.stop_sequences),
            thinking=None if params.thinking is None else dict(params.thinking),
            output_config=(
                None if params.output_config is None
                else dict(params.output_config)
            ),
        )

    def as_kwargs(self) -> dict[str, object]:
        """Convert to keyword arguments for API calls."""
        payload: dict[str, object] = {}
        if self.temperature is not None:
            payload["temperature"] = self.temperature
        if self.top_p is not None:
            payload["top_p"] = self.top_p
        if self.stop_sequences:
            payload["stop_sequences"] = list(self.stop_sequences)
        if self.thinking is not None:
            payload["thinking"] = dict(self.thinking)
        if self.output_config is not None:
            payload["output_config"] = dict(self.output_config)
        return payload


@dataclass
class ResponsesSession:
    """State for an active Responses API session.

    Tracks all information needed to resume a multi-turn conversation,
    including message history, tool definitions, and sampling parameters.
    """

    provider: str
    model: str
    messages: list[MessagePayload] = field(default_factory=list)
    tools: list[Mapping[str, object]] = field(default_factory=list)
    system: list[Mapping[str, str]] = field(default_factory=list)
    max_tokens: int = 0
    sampling: SamplingConfig = field(default_factory=SamplingConfig)
    api_key_override: str | None = None
    responses_request: ResponsesRequest | None = None
    streaming_response: StreamingResponse | None = None


# Thread-safe session storage
_SESSIONS: MutableMapping[str, ResponsesSession] = {}
_LOCK = threading.Lock()


def _coerce_sampling(value: SamplingConfig | Mapping[str, object]) -> SamplingConfig:
    """Convert a value to SamplingConfig."""
    if isinstance(value, SamplingConfig):
        return value
    temperature = _as_float(value.get("temperature"))
    top_p = _as_float(value.get("top_p"))
    stop_sequences_raw = value.get("stop_sequences")
    if stop_sequences_raw is None:
        stop_sequences: tuple[str, ...] = ()
    elif isinstance(stop_sequences_raw, Sequence):
        stop_sequences = tuple(
            str(item) for item in stop_sequences_raw if isinstance(item, str)
        )
    else:
        raise ProviderAPIError("stop_sequences must be an iterable of strings")
    thinking_raw = value.get("thinking")
    thinking = dict(thinking_raw) if isinstance(thinking_raw, Mapping) else None
    output_config_raw = value.get("output_config")
    output_config = dict(output_config_raw) if isinstance(output_config_raw, Mapping) else None
    return SamplingConfig(
        temperature=temperature,
        top_p=top_p,
        stop_sequences=stop_sequences,
        thinking=thinking,
        output_config=output_config,
    )


def _as_float(raw: object | None) -> float | None:
    """Coerce a value to float or None."""
    if raw is None:
        return None
    if isinstance(raw, (int, float)):
        return float(raw)
    try:
        return float(str(raw))
    except ValueError as exc:  # pragma: no cover
        raise ProviderAPIError("Sampling parameter must be numeric") from exc


def save_session(
    response_id: str,
    *,
    provider: str,
    model: str,
    messages: Sequence[MessagePayload],
    tools: Sequence[Mapping[str, object]] | None,
    system: Sequence[Mapping[str, str]] | None,
    max_tokens: int,
    sampling: SamplingConfig | Mapping[str, object],
    api_key_override: str | None,
    responses_request: ResponsesRequest | None,
    streaming_response: StreamingResponse | None,
) -> None:
    """Persist the state required to resume an Anthropic Responses run.

    Args:
        response_id: Unique identifier for this session.
        provider: Provider name (e.g., "anthropic").
        model: Model identifier.
        messages: Conversation history.
        tools: Tool definitions for this session.
        system: System prompts.
        max_tokens: Maximum tokens for responses.
        sampling: Sampling configuration.
        api_key_override: Optional API key override.
        responses_request: Original Responses API request.
        streaming_response: Associated streaming response.
    """
    session = ResponsesSession(
        provider=provider,
        model=model,
        messages=clone_messages(messages),
        tools=[dict(tool) for tool in (tools or ())],
        system=[dict(prompt) for prompt in (system or ())],
        max_tokens=max_tokens,
        sampling=_coerce_sampling(sampling),
        api_key_override=api_key_override,
        responses_request=responses_request,
        streaming_response=streaming_response,
    )
    with _LOCK:
        _SESSIONS[response_id] = session


def update_session_messages(
    response_id: str, messages: Sequence[MessagePayload]
) -> None:
    """Replace the stored message stack for ``response_id``.

    Args:
        response_id: Session identifier.
        messages: New message history.

    Raises:
        ProviderAPIError: If the session does not exist.
    """
    with _LOCK:
        session = _SESSIONS.get(response_id)
        if session is None:
            raise ProviderAPIError(f"Unknown response id '{response_id}'")
        session.messages = clone_messages(messages)


def get_session(response_id: str) -> ResponsesSession:
    """Return the saved session for ``response_id``.

    Args:
        response_id: Session identifier.

    Returns:
        The stored session.

    Raises:
        ProviderAPIError: If the session does not exist.
    """
    with _LOCK:
        session = _SESSIONS.get(response_id)
        if session is None:
            raise ProviderAPIError(f"Unknown response id '{response_id}'")
        return session


def drop_session(response_id: str) -> None:
    """Remove a stored session.

    Args:
        response_id: Session identifier to remove.
    """
    with _LOCK:
        _SESSIONS.pop(response_id, None)


def continue_with_tool_outputs(
    provider: AnthropicProvider,
    response_id: str,
    tool_outputs: Sequence[Mapping[str, object]],
) -> Iterable[str]:
    """Resume an Anthropic Responses run with tool results provided by the client.

    Args:
        provider: The Anthropic provider instance.
        response_id: Session identifier to resume.
        tool_outputs: Tool execution results from the client.

    Returns:
        An iterable of SSE event strings.

    Raises:
        ProviderAPIError: If the session does not exist or tool outputs are invalid.
    """
    session = get_session(response_id)
    messages = clone_messages(session.messages)

    tool_blocks = [
        {
            "type": "tool_result",
            "tool_use_id": str(
                output.get("tool_call_id")
                or output.get("call_id")
                or output.get("id")
                or output.get("tool_use_id")
            ),
            "content": output.get("output", output.get("result", "")),
        }
        for output in tool_outputs
    ]

    if any(not block["tool_use_id"] for block in tool_blocks):
        raise ProviderAPIError("Each tool output must include 'tool_call_id'")
    messages.append({"role": "user", "content": tool_blocks})
    update_session_messages(response_id, messages)

    args: dict[str, object] = {
        "model": provider._resolve_model_name(session.model),
        "messages": messages,
        "max_tokens": session.max_tokens,
        **session.sampling.as_kwargs(),
    }
    if session.system:
        args["system"] = session.system
    if session.tools:
        args["tools"] = session.tools

    def _generator() -> Iterable[str]:
        try:
            with provider.client.messages.stream(**args) as message_stream:
                yield from stream_events_from_anthropic(
                    stream=message_stream,
                    model=session.model,
                    response_id=response_id,
                )
        except anthropic.APIError as exc:
            raise ProviderAPIError(
                f"Anthropic API error while resuming stream: {exc}"
            ) from exc

    return _generator()


__all__ = [
    "ResponsesSession",
    "SamplingConfig",
    "continue_with_tool_outputs",
    "drop_session",
    "get_session",
    "save_session",
    "update_session_messages",
]
