"""Memory service for managing conversation history with LangChain.

This service provides memory management for AI conversations using
LangChain's message history abstractions (LangChain 1.0+ API).

Provides ConversationBufferMemory-style interface for managing conversation
history with support for multiple storage backends.

EntityMemory support enables tracking of named entities (projects, plans,
contexts, changes) across conversation sessions.
"""

from collections.abc import Callable, Sequence
from datetime import UTC, datetime
from enum import StrEnum
from typing import Any, cast

from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.chat_history import (
    BaseChatMessageHistory,
    InMemoryChatMessageHistory,
)
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from pydantic import BaseModel, Field


class EntityType(StrEnum):
    """Types of entities that can be tracked in memory."""

    PROJECT = "project"
    PLAN = "plan"
    CONTEXT = "context"
    CHANGE = "change"
    FILE = "file"
    FUNCTION = "function"
    CLASS = "class"
    CUSTOM = "custom"


class Entity(BaseModel):
    """Represents a tracked entity in memory.

    Attributes:
        name: The entity's name/identifier
        entity_type: The type of entity
        metadata: Additional metadata about the entity
        first_seen: When the entity was first tracked
        last_seen: When the entity was last accessed
        mention_count: Number of times entity has been referenced
    """

    model_config = {"arbitrary_types_allowed": True}

    name: str
    entity_type: EntityType
    metadata: dict[str, Any] = Field(default_factory=dict)
    first_seen: datetime = Field(default_factory=lambda: datetime.now(UTC))
    last_seen: datetime = Field(default_factory=lambda: datetime.now(UTC))
    mention_count: int = 1

    def touch(self) -> None:
        """Update the last_seen timestamp and increment mention count."""
        self.last_seen = datetime.now(UTC)
        self.mention_count += 1

    def to_dict(self) -> dict[str, Any]:
        """Convert entity to dictionary for serialization."""
        return {
            "name": self.name,
            "entity_type": self.entity_type.value,
            "metadata": self.metadata,
            "first_seen": self.first_seen.isoformat(),
            "last_seen": self.last_seen.isoformat(),
            "mention_count": self.mention_count,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Entity":
        """Create an entity from a dictionary."""
        return cls(
            name=data["name"],
            entity_type=EntityType(data["entity_type"]),
            metadata=data.get("metadata", {}),
            first_seen=datetime.fromisoformat(data["first_seen"]),
            last_seen=datetime.fromisoformat(data["last_seen"]),
            mention_count=data.get("mention_count", 1),
        )


class EntityStore:
    """Storage for tracked entities with optional SQL persistence.

    Provides in-memory entity storage with optional SQL-based persistence
    for cross-session entity recall.
    """

    def __init__(self, session_id: str, connection_string: str | None = None) -> None:
        """Initialize entity store.

        Args:
            session_id: Unique identifier for the session
            connection_string: Optional database connection for persistence
        """
        self.session_id = session_id
        self.connection_string = connection_string
        self._entities: dict[str, Entity] = {}
        self._dirty = False

        # Load entities from persistence if available
        if connection_string:
            self._load_from_persistence()

    def track(
        self,
        name: str,
        entity_type: EntityType,
        metadata: dict[str, Any] | None = None,
    ) -> Entity:
        """Track a new entity or update an existing one.

        Args:
            name: The entity's name/identifier
            entity_type: The type of entity
            metadata: Optional metadata about the entity

        Returns:
            The tracked entity
        """
        key = f"{entity_type.value}:{name}"

        if key in self._entities:
            entity = self._entities[key]
            entity.touch()
            if metadata:
                entity.metadata.update(metadata)
        else:
            entity = Entity(
                name=name,
                entity_type=entity_type,
                metadata=metadata or {},
            )
            self._entities[key] = entity

        self._dirty = True
        self._persist_if_needed()
        return entity

    def get(self, name: str, entity_type: EntityType) -> Entity | None:
        """Get a specific entity by name and type.

        Args:
            name: The entity's name/identifier
            entity_type: The type of entity

        Returns:
            The entity if found, None otherwise
        """
        key = f"{entity_type.value}:{name}"
        return self._entities.get(key)

    def get_by_type(self, entity_type: EntityType) -> list[Entity]:
        """Get all entities of a specific type.

        Args:
            entity_type: The type of entities to retrieve

        Returns:
            List of matching entities
        """
        return [
            entity
            for entity in self._entities.values()
            if entity.entity_type == entity_type
        ]

    def get_all(self) -> list[Entity]:
        """Get all tracked entities.

        Returns:
            List of all entities
        """
        return list(self._entities.values())

    def get_recent(self, count: int = 10) -> list[Entity]:
        """Get the most recently accessed entities.

        Args:
            count: Maximum number of entities to return

        Returns:
            List of recently accessed entities
        """
        sorted_entities = sorted(
            self._entities.values(), key=lambda e: e.last_seen, reverse=True
        )
        return sorted_entities[:count]

    def remove(self, name: str, entity_type: EntityType) -> bool:
        """Remove a specific entity.

        Args:
            name: The entity's name/identifier
            entity_type: The type of entity

        Returns:
            True if entity was removed, False if not found
        """
        key = f"{entity_type.value}:{name}"
        if key in self._entities:
            del self._entities[key]
            self._dirty = True
            self._persist_if_needed()
            return True
        return False

    def clear(self, entity_type: EntityType | None = None) -> int:
        """Clear entities from the store.

        Args:
            entity_type: If provided, only clear entities of this type.
                        If None, clear all entities.

        Returns:
            Number of entities removed
        """
        if entity_type is None:
            count = len(self._entities)
            self._entities.clear()
        else:
            keys_to_remove = [
                key
                for key, entity in self._entities.items()
                if entity.entity_type == entity_type
            ]
            count = len(keys_to_remove)
            for key in keys_to_remove:
                del self._entities[key]

        if count > 0:
            self._dirty = True
            self._persist_if_needed()
        return count

    def search(self, query: str) -> list[Entity]:
        """Search for entities by name (case-insensitive).

        Args:
            query: Search query string

        Returns:
            List of matching entities
        """
        query_lower = query.lower()
        return [
            entity
            for entity in self._entities.values()
            if query_lower in entity.name.lower()
        ]

    def _load_from_persistence(self) -> None:
        """Load entities from SQL persistence."""
        # For now, we use a simple in-memory approach
        # Future: Implement actual SQL persistence
        pass

    def _persist_if_needed(self) -> None:
        """Persist entities to SQL if dirty and connection available."""
        if not self._dirty or not self.connection_string:
            return
        # For now, mark as clean without actual persistence
        # Future: Implement actual SQL persistence
        self._dirty = False


def _buffer_to_string(messages: Sequence[BaseMessage]) -> str:
    """Convert a sequence of messages into a printable buffer string."""
    parts: list[str] = []
    for message in messages:
        role = "AI" if isinstance(message, AIMessage) else "Human"
        parts.append(f"{role}: {message.content!s}")
    return "\n".join(parts)


class ConversationBufferMemoryAdapter:
    """Lightweight adapter that mirrors LangChain's ConversationBufferMemory.

    This adapter provides a familiar interface for prompt templates and
    runnables that expect a conversational memory object while leveraging the
    configured `BaseChatMessageHistory` implementation under the hood.
    """

    def __init__(
        self,
        message_history: BaseChatMessageHistory,
        *,
        memory_key: str = "history",
        input_key: str = "input",
        output_key: str = "output",
        return_messages: bool = True,
        prune_callback: Callable[[], None] | None = None,
    ) -> None:
        self._message_history = message_history
        self.memory_key = memory_key
        self.input_key = input_key
        self.output_key = output_key
        self.return_messages = return_messages
        self._prune_callback = prune_callback

    @property
    def memory_variables(self) -> list[str]:
        """Return the keys that this memory provides."""

        return [self.memory_key]

    def load_memory_variables(
        self, _inputs: dict[str, Any] | None = None
    ) -> dict[str, Any]:
        """Return the current memory payload for prompt formatting."""

        messages = list(self._message_history.messages)
        if self.return_messages:
            value: Any = messages
        else:
            value = _buffer_to_string(messages)
        return {self.memory_key: value}

    async def aload_memory_variables(
        self, inputs: dict[str, Any] | None = None
    ) -> dict[str, Any]:
        """Async variant mirroring the sync load method."""

        return self.load_memory_variables(inputs)

    def save_context(self, inputs: dict[str, Any], outputs: dict[str, Any]) -> None:
        """Persist a complete interaction into the underlying history."""

        def _store_user(value: BaseMessage | Sequence[Any] | str | None) -> None:
            if value is None:
                return
            if isinstance(value, BaseMessage):
                self._message_history.add_message(value)
                return
            if isinstance(value, str | bytes | bytearray):
                self._message_history.add_user_message(str(value))
                return
            if not isinstance(value, Sequence):
                raise TypeError(
                    f"Expected Sequence for user message, got {type(value).__name__}"
                )
            for item in value:
                _store_user(item)

        def _store_ai(value: BaseMessage | Sequence[Any] | str | None) -> None:
            if value is None:
                return
            if isinstance(value, BaseMessage):
                self._message_history.add_message(value)
                return
            if isinstance(value, str | bytes | bytearray):
                self._message_history.add_ai_message(str(value))
                return
            if not isinstance(value, Sequence):
                raise TypeError(
                    f"Expected Sequence for AI message, got {type(value).__name__}"
                )
            for item in value:
                _store_ai(item)

        _store_user(inputs.get(self.input_key))
        _store_ai(outputs.get(self.output_key))

        if self._prune_callback:
            self._prune_callback()

    async def asave_context(
        self, inputs: dict[str, Any], outputs: dict[str, Any]
    ) -> None:
        """Async variant mirroring the sync save method."""

        self.save_context(inputs, outputs)

    def clear(self) -> None:
        """Clear the stored conversation history."""

        self._message_history.clear()
        if self._prune_callback:
            self._prune_callback()

    async def aclear(self) -> None:
        """Async variant of ``clear``."""

        self.clear()


class MemoryService:
    """Service for managing conversation memory using LangChain 1.0+.

    Provides different memory strategies:
    - In-memory: Keeps conversation history in memory
    - SQL persistence: Stores messages in database
    - Summarization: Maintains summarized conversation (via custom logic)
    - EntityMemory: Tracks named entities across sessions
    """

    def __init__(
        self,
        session_id: str,
        connection_string: str | None = None,
        max_messages: int | None = None,
    ):
        """Initialize memory service.

        Args:
            session_id: Unique identifier for the conversation session
            connection_string: Database connection for SQL persistence (optional)
            max_messages: Maximum number of messages to keep in memory (optional)
        """
        self.session_id = session_id
        self.connection_string = connection_string
        self.max_messages = max_messages

        # Initialize message history
        history: BaseChatMessageHistory
        if connection_string:
            # Use SQL persistence if connection string provided
            history = cast(
                BaseChatMessageHistory,
                SQLChatMessageHistory(
                    session_id=session_id,
                    connection_string=connection_string,
                    table_name="message_history",
                ),
            )
        else:
            # Use in-memory storage as fallback
            history = InMemoryChatMessageHistory()
        self.message_history = history

        # Default conversation-memory adapter mirroring LangChain's API.
        self._conversation_memory = ConversationBufferMemoryAdapter(
            self.message_history,
            prune_callback=self._enforce_max_messages,
        )

        # Initialize entity store for entity tracking
        self._entity_store = EntityStore(
            session_id=session_id,
            connection_string=connection_string,
        )

    def add_user_message(self, message: str) -> None:
        """Add a user message to memory.

        Args:
            message: The user's message
        """
        self.message_history.add_message(HumanMessage(content=message))
        self._enforce_max_messages()

    def add_ai_message(self, message: str) -> None:
        """Add an AI response to memory.

        Args:
            message: The AI's response
        """
        self.message_history.add_message(AIMessage(content=message))
        self._enforce_max_messages()

    def add_message(self, message: BaseMessage) -> None:
        """Add a message object to memory.

        Args:
            message: The message object to add
        """
        self.message_history.add_message(message)
        self._enforce_max_messages()

    def save_context(self, inputs: dict[str, Any], outputs: dict[str, Any]) -> None:
        """Save a complete interaction to memory.

        Args:
            inputs: The input variables (expects 'input' key for user message)
            outputs: The output variables (expects 'output' key for AI message)
        """
        user_message = inputs.get("input", "")
        ai_message = outputs.get("output", "")

        if user_message:
            self.add_user_message(user_message)
        if ai_message:
            self.add_ai_message(ai_message)

    def get_messages(self) -> list[BaseMessage]:
        """Get full conversation history.

        Returns:
            List of message objects
        """
        return self.message_history.messages

    def get_recent_messages(self, count: int) -> list[BaseMessage]:
        """Get the most recent N messages.

        Args:
            count: Number of recent messages to retrieve

        Returns:
            List of most recent message objects
        """
        messages = self.message_history.messages
        return messages[-count:] if count < len(messages) else messages

    def get_summary(self, max_chars: int = 500) -> str:
        """Get a simple summary of the conversation.

        Args:
            max_chars: Maximum characters in summary

        Returns:
            Summary string
        """
        messages = self.get_messages()
        if not messages:
            return ""

        # Simple summarization: concatenate recent messages
        summary_parts: list[str] = []
        char_count = 0

        for msg in reversed(messages):
            content_str = str(msg.content)
            prefix = "User: " if isinstance(msg, HumanMessage) else "AI: "
            msg_text = f"{prefix}{content_str}"

            if char_count + len(msg_text) > max_chars:
                break

            summary_parts.insert(0, msg_text)
            char_count += len(msg_text)

        return "\n".join(summary_parts)

    def get_memory_variables(self) -> dict[str, Any]:
        """Get all memory variables for use in prompts."""
        messages = self.get_messages()
        return {
            "chat_history": messages,
            "chat_summary": self.get_summary(),
            "message_count": len(messages),
        }

    @property
    def conversation_memory(self) -> ConversationBufferMemoryAdapter:
        """Return the default conversation buffer memory adapter."""

        return self._conversation_memory

    def create_conversation_memory(
        self,
        *,
        memory_key: str = "history",
        input_key: str = "input",
        output_key: str = "output",
        return_messages: bool = True,
    ) -> ConversationBufferMemoryAdapter:
        """Create a new conversation memory adapter with custom settings."""

        return ConversationBufferMemoryAdapter(
            self.message_history,
            memory_key=memory_key,
            input_key=input_key,
            output_key=output_key,
            return_messages=return_messages,
            prune_callback=self._enforce_max_messages,
        )

    def set_max_messages(self, max_messages: int | None) -> None:
        """Update the max message window and prune history if needed."""

        self.max_messages = max_messages
        self._enforce_max_messages()

    def clear(self) -> None:
        """Clear all memory."""
        self.message_history.clear()

    def get_token_count(self) -> int:
        """Get approximate token count of conversation.

        Returns:
            Estimated token count
        """
        # Simple approximation: 1 token ~= 4 characters
        messages = self.get_messages()
        total_chars = sum(len(str(msg.content)) for msg in messages)
        return total_chars // 4

    def _enforce_max_messages(self) -> None:
        """Enforce maximum message limit by removing oldest messages."""
        if self.max_messages is None:
            return

        messages = list(self.message_history.messages)
        if len(messages) > self.max_messages:
            # Remove oldest messages
            messages_to_remove = len(messages) - self.max_messages
            remaining_messages = messages[messages_to_remove:]
            self.message_history.clear()
            for msg in remaining_messages:
                self.message_history.add_message(msg)

    # ==================== EntityMemory Methods ====================

    def _resolve_entity_type(self, entity_type: EntityType | str) -> EntityType:
        """Convert string entity type to EntityType enum if needed."""
        if not isinstance(entity_type, EntityType):
            return EntityType(entity_type)
        return entity_type

    def track_entity(
        self,
        name: str,
        entity_type: EntityType | str,
        metadata: dict[str, Any] | None = None,
    ) -> Entity:
        """Track a new entity or update an existing one.

        This enables cross-session entity recall for projects, plans,
        contexts, and other domain objects.

        Args:
            name: The entity's name/identifier
            entity_type: The type of entity (EntityType enum or string)
            metadata: Optional metadata about the entity

        Returns:
            The tracked entity
        """
        resolved_type = self._resolve_entity_type(entity_type)
        return self._entity_store.track(name, resolved_type, metadata)

    def get_entity(self, name: str, entity_type: EntityType | str) -> Entity | None:
        """Get a specific entity by name and type.

        Args:
            name: The entity's name/identifier
            entity_type: The type of entity

        Returns:
            The entity if found, None otherwise
        """
        resolved_type = self._resolve_entity_type(entity_type)
        return self._entity_store.get(name, resolved_type)

    def get_entities(self, entity_type: EntityType | str | None = None) -> list[Entity]:
        """Get tracked entities, optionally filtered by type.

        Args:
            entity_type: If provided, only return entities of this type.
                        If None, return all entities.

        Returns:
            List of matching entities
        """
        if entity_type is None:
            return self._entity_store.get_all()
        resolved_type = self._resolve_entity_type(entity_type)
        return self._entity_store.get_by_type(resolved_type)

    def get_recent_entities(self, count: int = 10) -> list[Entity]:
        """Get the most recently accessed entities.

        Args:
            count: Maximum number of entities to return

        Returns:
            List of recently accessed entities
        """
        return self._entity_store.get_recent(count)

    def search_entities(self, query: str) -> list[Entity]:
        """Search for entities by name (case-insensitive).

        Args:
            query: Search query string

        Returns:
            List of matching entities
        """
        return self._entity_store.search(query)

    def clear_entities(self, entity_type: EntityType | str | None = None) -> int:
        """Clear tracked entities.

        Args:
            entity_type: If provided, only clear entities of this type.
                        If None, clear all entities.

        Returns:
            Number of entities removed
        """
        if entity_type is None:
            return self._entity_store.clear()
        resolved_type = self._resolve_entity_type(entity_type)
        return self._entity_store.clear(resolved_type)

    def remove_entity(self, name: str, entity_type: EntityType | str) -> bool:
        """Remove a specific entity.

        Args:
            name: The entity's name/identifier
            entity_type: The type of entity

        Returns:
            True if entity was removed, False if not found
        """
        resolved_type = self._resolve_entity_type(entity_type)
        return self._entity_store.remove(name, resolved_type)

    def get_entity_summary(self) -> dict[str, Any]:
        """Get a summary of all tracked entities.

        Returns:
            Dictionary with entity counts by type and recent entities
        """
        all_entities = self._entity_store.get_all()

        # Count by type
        counts: dict[str, int] = {}
        for entity in all_entities:
            type_name = entity.entity_type.value
            counts[type_name] = counts.get(type_name, 0) + 1

        return {
            "total_entities": len(all_entities),
            "counts_by_type": counts,
            "recent_entities": [e.to_dict() for e in self._entity_store.get_recent(5)],
        }

    @property
    def entity_store(self) -> EntityStore:
        """Return the entity store for direct access."""
        return self._entity_store
