"""Memory management for conversation history.

This module provides interfaces and implementations for storing
and retrieving conversation history.

Includes:
- MemoryStore: Abstract interface for storage backends
- InMemoryStore: Simple in-memory storage (for development)
- DynamoDBMemoryStore: Persistent storage using AWS DynamoDB (for production)
"""

import json
import os
import time
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional

from agent_factory.models.base import Message

# DynamoDB client (lazy initialized)
_dynamodb_client = None


def _get_dynamodb_client():
    """Get or create DynamoDB client."""
    global _dynamodb_client
    if _dynamodb_client is None:
        import boto3

        _dynamodb_client = boto3.client("dynamodb")
    return _dynamodb_client


class MemoryStore(ABC):
    """Abstract interface for memory storage backends."""

    @abstractmethod
    def save_turn(
        self,
        session_id: str,
        user_input: str,
        agent_response: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a conversation turn."""
        pass

    @abstractmethod
    def get_messages(
        self,
        session_id: str,
        limit: Optional[int] = None,
    ) -> List[Message]:
        """Get messages for a session."""
        pass

    @abstractmethod
    def clear(self, session_id: str) -> None:
        """Clear messages for a session."""
        pass


class InMemoryStore(MemoryStore):
    """Simple in-memory storage for development/testing."""

    def __init__(self):
        self._sessions: Dict[str, List[Message]] = {}

    def save_turn(
        self,
        session_id: str,
        user_input: str,
        agent_response: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        if session_id not in self._sessions:
            self._sessions[session_id] = []

        self._sessions[session_id].extend(
            [
                Message(role="user", content=user_input, metadata=metadata or {}),
                Message(role="assistant", content=agent_response, metadata=metadata or {}),
            ]
        )

    def get_messages(
        self,
        session_id: str,
        limit: Optional[int] = None,
    ) -> List[Message]:
        messages = self._sessions.get(session_id, [])
        if limit:
            return messages[-limit:]
        return messages

    def clear(self, session_id: str) -> None:
        self._sessions.pop(session_id, None)


class DynamoDBMemoryStore(MemoryStore):
    """DynamoDB-based storage for conversation memory.

    Persists conversation history across Lambda invocations using DynamoDB.
    Supports TTL for automatic cleanup of old conversations.

    Table schema:
        - pk (String): Partition key - session_id
        - sk (String): Sort key - "turn#<timestamp>" or "meta"
        - messages (String): JSON-encoded messages for the turn
        - ttl (Number): TTL timestamp for auto-deletion

    Example:
        store = DynamoDBMemoryStore(table_name="my-conversations")
        # or use env var
        store = DynamoDBMemoryStore()  # uses MEMORY_TABLE env var

        memory = Memory(session_id="user-123", store=store)
    """

    def __init__(
        self,
        table_name: Optional[str] = None,
        ttl_days: int = 7,
    ):
        """Initialize DynamoDB memory store.

        Args:
            table_name: DynamoDB table name. Defaults to MEMORY_TABLE env var.
            ttl_days: Days to keep conversations before auto-deletion.
        """
        self.table_name = table_name or os.environ.get("MEMORY_TABLE", "")
        self.ttl_seconds = ttl_days * 24 * 60 * 60

        if not self.table_name:
            raise ValueError("table_name must be provided or MEMORY_TABLE env var must be set")

    def save_turn(
        self,
        session_id: str,
        user_input: str,
        agent_response: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save a conversation turn to DynamoDB."""
        client = _get_dynamodb_client()
        timestamp = int(time.time() * 1000)
        ttl = int(time.time()) + self.ttl_seconds

        # Create messages for this turn
        messages = [
            {"role": "user", "content": user_input},
            {"role": "assistant", "content": agent_response},
        ]

        item = {
            "pk": {"S": session_id},
            "sk": {"S": f"turn#{timestamp}"},
            "messages": {"S": json.dumps(messages)},
            "ttl": {"N": str(ttl)},
            "created_at": {"N": str(timestamp)},
        }

        if metadata:
            item["metadata"] = {"S": json.dumps(metadata)}

        client.put_item(TableName=self.table_name, Item=item)

    def get_messages(
        self,
        session_id: str,
        limit: Optional[int] = None,
    ) -> List[Message]:
        """Get messages for a session from DynamoDB."""
        client = _get_dynamodb_client()

        # Query all turns for this session, sorted by timestamp
        response = client.query(
            TableName=self.table_name,
            KeyConditionExpression="pk = :pk AND begins_with(sk, :prefix)",
            ExpressionAttributeValues={
                ":pk": {"S": session_id},
                ":prefix": {"S": "turn#"},
            },
            ScanIndexForward=True,  # Oldest first
        )

        all_messages: List[Message] = []

        for item in response.get("Items", []):
            messages_json = item.get("messages", {}).get("S", "[]")
            turn_messages = json.loads(messages_json)

            for msg in turn_messages:
                all_messages.append(
                    Message(
                        role=msg.get("role", "user"),
                        content=msg.get("content", ""),
                    )
                )

        # Apply limit if specified
        if limit and len(all_messages) > limit:
            return all_messages[-limit:]

        return all_messages

    def clear(self, session_id: str) -> None:
        """Clear all messages for a session."""
        client = _get_dynamodb_client()

        # Query all items for this session
        response = client.query(
            TableName=self.table_name,
            KeyConditionExpression="pk = :pk",
            ExpressionAttributeValues={":pk": {"S": session_id}},
            ProjectionExpression="pk, sk",
        )

        # Delete each item
        for item in response.get("Items", []):
            client.delete_item(
                TableName=self.table_name,
                Key={
                    "pk": item["pk"],
                    "sk": item["sk"],
                },
            )


class Memory:
    """High-level memory manager for conversations.

    Provides a convenient interface for managing conversation history
    with pluggable storage backends.

    Example:
        # Using in-memory storage (default)
        memory = Memory(session_id="user-123")

        # Using custom storage
        memory = Memory(session_id="user-123", store=DynamoDBStore())

        # Add a conversation turn
        memory.add_turn("Hello!", "Hi there!")

        # Get formatted messages for AI
        messages = memory.get_messages()
    """

    def __init__(
        self,
        session_id: str,
        store: Optional[MemoryStore] = None,
        max_turns: int = 20,
    ):
        """Initialize memory manager.

        Args:
            session_id: Unique identifier for this conversation
            store: Storage backend (defaults to InMemoryStore)
            max_turns: Maximum turns to include in context
        """
        self.session_id = session_id
        self._store = store or InMemoryStore()
        self._max_turns = max_turns
        self._messages: List[Message] = []

    def add_turn(
        self,
        user_input: str,
        agent_response: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Add a conversation turn.

        Args:
            user_input: User's message
            agent_response: Agent's response
            metadata: Optional metadata for the turn
        """
        self._messages.append(Message(role="user", content=user_input))
        self._messages.append(Message(role="assistant", content=agent_response))

        # Persist to store
        self._store.save_turn(
            self.session_id,
            user_input,
            agent_response,
            metadata,
        )

    def add_message(self, role: str, content: str) -> None:
        """Add a single message.

        Args:
            role: Message role (user, assistant, system)
            content: Message content
        """
        self._messages.append(Message(role=role, content=content))

    def get_messages(self, limit: Optional[int] = None) -> List[Message]:
        """Get conversation messages.

        Args:
            limit: Maximum messages to return (defaults to max_turns * 2)

        Returns:
            List of Message objects
        """
        max_messages = limit or (self._max_turns * 2)

        # Try to get from store first
        stored = self._store.get_messages(self.session_id, max_messages)
        if stored:
            return stored

        # Fall back to local buffer
        return self._messages[-max_messages:] if max_messages else self._messages

    def get_formatted_messages(self) -> List[Dict[str, Any]]:
        """Get messages formatted for AI model.

        Returns:
            Messages in Bedrock Converse API format
        """
        return [msg.to_bedrock_format() for msg in self.get_messages()]

    def clear(self) -> None:
        """Clear all messages for this session."""
        self._messages = []
        self._store.clear(self.session_id)

    def __len__(self) -> int:
        """Get number of messages."""
        return len(self._messages)
