"""
Boring-OS SDK agent session.

Provides the AgentSession class for multi-turn conversations
with a Boring-OS hosted agent. Maintains conversation memory
across multiple `send()` calls via a server-side session.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, AsyncIterator

if TYPE_CHECKING:
    from boring_os.client import BoringOS

from boring_os.agent_models import AgentResponse, AgentStreamChunk


class AgentSession:
    """
    Multi-turn agent conversation session.

    Created via `boring.agents.create_session("agent_name")`.
    All messages sent through the session share conversation memory.

    Can be used as a context manager:

        session = boring.agents.create_session("billing_agent")
        with session:
            r1 = session.send("Why was I overcharged?")
            r2 = session.send("Can you break down roaming charges?")

    Or async context manager:

        async with boring.agents.create_session("billing_agent") as session:
            r1 = session.send("Why was I overcharged?")
    """

    def __init__(
        self,
        client: "BoringOS",
        session_id: str,
        agent_name: str,
    ):
        """
        Initialize an agent session.

        Args:
            client: The BoringOS client instance
            session_id: Server-assigned session identifier
            agent_name: Agent associated with this session
        """
        self._client = client
        self._session_id = session_id
        self._agent_name = agent_name
        self._total_messages = 0
        self._total_cost = 0.0
        self._closed = False

    # -------------------------------------------------------------------------
    # Properties
    # -------------------------------------------------------------------------

    @property
    def session_id(self) -> str:
        """Unique session identifier."""
        return self._session_id

    @property
    def agent_name(self) -> str:
        """Agent name for this session."""
        return self._agent_name

    @property
    def total_messages(self) -> int:
        """Number of messages exchanged in this session."""
        return self._total_messages

    @property
    def total_cost(self) -> float:
        """Cumulative cost of all messages in this session."""
        return self._total_cost

    # -------------------------------------------------------------------------
    # Messaging
    # -------------------------------------------------------------------------

    def send(self, message: str, timeout: int = 120) -> AgentResponse:
        """
        Send a message to the agent in this session.

        Args:
            message: User message text
            timeout: Max seconds to wait for the response

        Returns:
            AgentResponse with the agent's reply

        Raises:
            SessionError: If the session is closed or invalid
        """
        self._check_closed()

        response = self._client._request(
            "POST",
            f"/agents/sessions/{self._session_id}/messages",
            json={"message": message},
            timeout=timeout + 5,
        )

        result = AgentResponse.model_validate(response)
        self._total_messages += 2  # user + assistant
        self._total_cost += result.cost
        return result

    async def send_async(self, message: str, timeout: int = 120) -> AgentResponse:
        """
        Async version of send().

        Args:
            message: User message text
            timeout: Max seconds to wait

        Returns:
            AgentResponse with the agent's reply
        """
        self._check_closed()

        response = await self._client._request_async(
            "POST",
            f"/agents/sessions/{self._session_id}/messages",
            json={"message": message},
            timeout=timeout + 5,
        )

        result = AgentResponse.model_validate(response)
        self._total_messages += 2
        self._total_cost += result.cost
        return result

    async def stream(self, message: str) -> AsyncIterator[AgentStreamChunk]:
        """
        Stream a message to the agent and yield response chunks.

        Args:
            message: User message text

        Yields:
            AgentStreamChunk objects as the response is generated
        """
        self._check_closed()

        async for chunk in self._client._stream_request(
            "POST",
            f"/agents/sessions/{self._session_id}/messages",
            json={"message": message, "stream": True},
        ):
            parsed = AgentStreamChunk.model_validate(chunk)
            yield parsed
            if parsed.done and parsed.token_usage:
                self._total_messages += 2
                self._total_cost += parsed.token_usage.cost

    # -------------------------------------------------------------------------
    # Lifecycle
    # -------------------------------------------------------------------------

    def close(self) -> None:
        """
        Close the session.

        After closing, no more messages can be sent.
        """
        if not self._closed:
            try:
                self._client._request(
                    "DELETE",
                    f"/agents/sessions/{self._session_id}",
                )
            except Exception:
                pass  # Best-effort close
            self._closed = True

    async def aclose(self) -> None:
        """Async version of close()."""
        if not self._closed:
            try:
                await self._client._request_async(
                    "DELETE",
                    f"/agents/sessions/{self._session_id}",
                )
            except Exception:
                pass
            self._closed = True

    def _check_closed(self) -> None:
        """Raise an error if the session has been closed."""
        if self._closed:
            from boring_os.exceptions import SessionError
            raise SessionError("Session has been closed")

    # -------------------------------------------------------------------------
    # Context Manager Support
    # -------------------------------------------------------------------------

    def __enter__(self) -> "AgentSession":
        """Context manager entry."""
        return self

    def __exit__(self, *args: Any) -> None:
        """Context manager exit -- closes the session."""
        self.close()

    async def __aenter__(self) -> "AgentSession":
        """Async context manager entry."""
        return self

    async def __aexit__(self, *args: Any) -> None:
        """Async context manager exit -- closes the session."""
        await self.aclose()

    def __repr__(self) -> str:
        status = "closed" if self._closed else "open"
        return (
            f"AgentSession(session_id={self._session_id!r}, "
            f"agent={self._agent_name!r}, "
            f"messages={self._total_messages}, status={status})"
        )
