"""AgentCore Memory client wrapping STM and LTM operations.

Provides :class:`MemoryClient` — a thin wrapper around the ``bedrock-agentcore``
Memory SDK that isolates SDK specifics and enables graceful degradation when the
service is unreachable or the memory resource is not active.

SDK version note: bedrock-agentcore split the old ``AgentCoreMemory`` class into
two separate clients:
  - ``bedrock_agentcore.memory.MemoryClient`` — control plane (create/status)
  - ``bedrock_agentcore.memory.MemorySessionManager`` — data plane (events/LTM)
"""

from __future__ import annotations

import contextlib
import io
import logging
from typing import TYPE_CHECKING, Any

from alice_cli.agentcore_deps import require_agentcore
from alice_cli.console import spinner
from alice_cli.errors import MemoryError

if TYPE_CHECKING:
    import boto3

logger = logging.getLogger(__name__)


class MemoryClient:
    """Wraps AgentCore Memory SDK for STM and LTM operations.

    Parameters
    ----------
    memory_name:
        Name of the AgentCore Memory resource (e.g. ``alice_memory``).
    region:
        AWS region for the AgentCore Memory service.
    actor_id:
        Actor identifier used for session-scoped operations.  Defaults to
        ``"default"``; pass the user's JHED when available for per-user isolation.
    boto3_session:
        Optional boto3 Session carrying AWS credentials (profile, API keys).
        When omitted the SDK uses boto3's default credential chain.
    """

    def __init__(
        self,
        memory_name: str,
        region: str,
        actor_id: str = "default",
        boto3_session: boto3.Session | None = None,
        quiet: bool = False,
    ) -> None:
        require_agentcore()
        self._memory_name = memory_name
        self._region = region
        self._actor_id = actor_id
        self._boto3_session = boto3_session
        self._quiet = quiet
        self._memory_id, self._control, self._session_mgr = self._create_clients()

    def _create_clients(self) -> tuple[str, Any, Any]:
        """Resolve memory_id from name, then create control + data plane clients."""
        try:
            from bedrock_agentcore.memory import (
                MemoryClient as SDKMemoryClient,
            )
            from bedrock_agentcore.memory import (
                MemorySessionManager,
            )

            control = SDKMemoryClient(region_name=self._region, boto3_session=self._boto3_session)
            spinner_msg = (
                f"Connecting to memory '{self._memory_name}'"
                " (first-time provisioning may take ~2 min)…"
            )
            with (
                spinner(spinner_msg, self._quiet),
                contextlib.redirect_stdout(io.StringIO()),
            ):
                memory = control.create_or_get_memory(name=self._memory_name)
            memory_id: str = memory["id"]
            session_mgr = MemorySessionManager(
                memory_id=memory_id,
                region_name=self._region,
                boto3_session=self._boto3_session,
            )
            return memory_id, control, session_mgr
        except Exception as exc:
            raise MemoryError(
                f"Failed to create AgentCore Memory client for '{self._memory_name}': {exc}"
            ) from exc

    # ------------------------------------------------------------------
    # Status
    # ------------------------------------------------------------------

    def check_status(self) -> str:
        """Check the status of the Memory resource.

        Returns the status string (e.g. ``"ACTIVE"``).

        Raises :class:`MemoryError` if the resource is not ACTIVE or
        unreachable, with instructions to create it.
        """
        try:
            status = self._control.get_memory_status(self._memory_id)
            if status != "ACTIVE":
                raise MemoryError(
                    f"Memory resource '{self._memory_name}' is in state '{status}'. "
                    "It must be ACTIVE before use. "
                    "Create it with: agentcore memory create"
                )
            return str(status)
        except MemoryError:
            raise
        except Exception as exc:
            raise MemoryError(
                f"Memory resource '{self._memory_name}' is not reachable: {exc}. "
                "Create it with: agentcore memory create"
            ) from exc

    # ------------------------------------------------------------------
    # STM operations
    # ------------------------------------------------------------------

    def store_event(self, session_id: str, role: str, content: str) -> None:
        """Store a conversation event (user or assistant message) in STM.

        Parameters
        ----------
        session_id:
            UUID4 identifying the conversation session.
        role:
            Message role — ``"user"`` or ``"assistant"``.
        content:
            The message text.
        """
        try:
            from bedrock_agentcore.memory.constants import (
                ConversationalMessage,
                MessageRole,
            )

            msg_role = MessageRole.USER if role.lower() == "user" else MessageRole.ASSISTANT
            self._session_mgr.add_turns(
                actor_id=self._actor_id,
                session_id=session_id,
                messages=[ConversationalMessage(text=content, role=msg_role)],
            )
        except Exception as exc:
            raise MemoryError(f"Failed to store STM event: {exc}") from exc

    def get_events(self, session_id: str) -> list[dict[str, Any]]:
        """Retrieve all STM events for a session, ordered chronologically.

        Returns an empty list if the session has no events or the service is
        unreachable.
        """
        try:
            events = self._session_mgr.list_events(actor_id=self._actor_id, session_id=session_id)
            result: list[dict[str, Any]] = []
            for ev in events:
                payload = ev.get("payload") or []
                for item in payload:
                    if "conversational" in item:
                        conv = item["conversational"]
                        result.append(
                            {
                                "role": conv.get("role", "").lower(),
                                "content": conv.get("content", {}).get("text", ""),
                            }
                        )
            return result
        except Exception as exc:
            logger.warning("Failed to retrieve STM events for session %s: %s", session_id, exc)
            return []

    def delete_events(self, session_id: str) -> None:
        """Delete all STM events for a session."""
        try:
            events = self._session_mgr.list_events(actor_id=self._actor_id, session_id=session_id)
            for ev in events:
                event_id = ev.get("eventId")
                if event_id:
                    self._session_mgr.delete_event(
                        actor_id=self._actor_id,
                        session_id=session_id,
                        event_id=event_id,
                    )
        except Exception as exc:
            raise MemoryError(
                f"Failed to delete STM events for session {session_id}: {exc}"
            ) from exc

    def list_sessions(self) -> list[dict[str, Any]]:
        """List all sessions for this client's actor.

        Returns an empty list if no sessions exist or the service is unreachable.
        """
        try:
            sessions = self._session_mgr.list_actor_sessions(actor_id=self._actor_id)
            return [
                {
                    "session_id": s.get("sessionId", ""),
                    "created_at": s.get("createdAt", ""),
                    "message_count": s.get("eventCount", 0),
                }
                for s in sessions
            ]
        except Exception as exc:
            logger.warning("Failed to list sessions: %s", exc)
            return []

    # ------------------------------------------------------------------
    # LTM operations
    # ------------------------------------------------------------------

    def store_fact(self, text: str) -> None:
        """Store a fact or knowledge item in long-term memory.

        In the current SDK, LTM records are derived from events via configured
        memory strategies.  This method stores the fact as a user event in a
        dedicated session so it is eligible for LTM consolidation.

        Raises :class:`MemoryError` if the memory resource lacks LTM strategies.
        """
        import uuid

        try:
            from bedrock_agentcore.memory.constants import (
                ConversationalMessage,
                MessageRole,
            )

            # Store in a dedicated single-event session for LTM consolidation.
            fact_session_id = str(uuid.uuid4())
            self._session_mgr.add_turns(
                actor_id=self._actor_id,
                session_id=fact_session_id,
                messages=[ConversationalMessage(text=text, role=MessageRole.USER)],
            )
        except Exception as exc:
            msg = str(exc)
            if "strategy" in msg.lower() or "ltm" in msg.lower():
                raise MemoryError(
                    "LTM requires memory strategies to be configured on the memory resource. "
                    "Set up LTM strategies before using 'thread remember'."
                ) from exc
            raise MemoryError(f"Failed to store LTM fact: {exc}") from exc

    def search_facts(self, query: str, limit: int = 10) -> list[dict[str, Any]]:
        """Perform semantic search over stored LTM facts.

        Returns results ordered by relevance score (descending). Returns an
        empty list if no facts match or the service is unreachable.
        """
        try:
            results = self._session_mgr.search_long_term_memories(
                query=query,
                namespace_prefix=self._memory_name,
                top_k=limit,
            )
            return [
                {
                    "text": r.get("content", {}).get("text", str(r)),
                    "score": r.get("score", r.get("relevanceScore", "")),
                }
                for r in results
            ]
        except Exception as exc:
            logger.warning("Failed to search LTM facts: %s", exc)
            return []

    def clear_facts(self) -> None:
        """Clear all LTM facts from the memory resource."""
        try:
            self._session_mgr.delete_all_long_term_memories_in_namespace(
                namespace=self._memory_name
            )
        except Exception as exc:
            raise MemoryError(f"Failed to clear LTM facts: {exc}") from exc
