"""
Context Models for Breakpilot Agents

Data classes for conversation messages and context management.
"""

from typing import Dict, Any, List, Optional
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
import logging

logger = logging.getLogger(__name__)


class MessageRole(Enum):
    """Message roles in a conversation"""
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"


@dataclass
class Message:
    """Represents a message in a conversation"""
    role: MessageRole
    content: str
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "role": self.role.value,
            "content": self.content,
            "timestamp": self.timestamp.isoformat(),
            "metadata": self.metadata
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Message":
        return cls(
            role=MessageRole(data["role"]),
            content=data["content"],
            timestamp=datetime.fromisoformat(data["timestamp"]) if "timestamp" in data else datetime.now(timezone.utc),
            metadata=data.get("metadata", {})
        )


@dataclass
class ConversationContext:
    """
    Context for a running conversation.

    Maintains:
    - Message history with automatic compression
    - Extracted entities
    - Intent history
    - Conversation summary
    """
    messages: List[Message] = field(default_factory=list)
    entities: Dict[str, Any] = field(default_factory=dict)
    intent_history: List[str] = field(default_factory=list)
    summary: Optional[str] = None
    max_messages: int = 50
    system_prompt: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def add_message(
        self,
        role: MessageRole,
        content: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> Message:
        """
        Adds a message to the conversation.

        Args:
            role: Message role
            content: Message content
            metadata: Optional message metadata

        Returns:
            The created Message
        """
        message = Message(
            role=role,
            content=content,
            metadata=metadata or {}
        )
        self.messages.append(message)

        # Compress if needed
        if len(self.messages) > self.max_messages:
            self._compress_history()

        return message

    def add_user_message(
        self,
        content: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> Message:
        """Convenience method to add a user message"""
        return self.add_message(MessageRole.USER, content, metadata)

    def add_assistant_message(
        self,
        content: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> Message:
        """Convenience method to add an assistant message"""
        return self.add_message(MessageRole.ASSISTANT, content, metadata)

    def add_system_message(
        self,
        content: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> Message:
        """Convenience method to add a system message"""
        return self.add_message(MessageRole.SYSTEM, content, metadata)

    def add_intent(self, intent: str) -> None:
        """
        Records an intent in the history.

        Args:
            intent: The detected intent
        """
        self.intent_history.append(intent)
        # Keep last 20 intents
        if len(self.intent_history) > 20:
            self.intent_history = self.intent_history[-20:]

    def set_entity(self, name: str, value: Any) -> None:
        """
        Sets an entity value.

        Args:
            name: Entity name
            value: Entity value
        """
        self.entities[name] = value

    def get_entity(self, name: str, default: Any = None) -> Any:
        """
        Gets an entity value.

        Args:
            name: Entity name
            default: Default value if not found

        Returns:
            Entity value or default
        """
        return self.entities.get(name, default)

    def get_last_message(self, role: Optional[MessageRole] = None) -> Optional[Message]:
        """
        Gets the last message, optionally filtered by role.

        Args:
            role: Optional role filter

        Returns:
            The last matching message or None
        """
        if not self.messages:
            return None

        if role is None:
            return self.messages[-1]

        for msg in reversed(self.messages):
            if msg.role == role:
                return msg

        return None

    def get_messages_for_llm(self) -> List[Dict[str, str]]:
        """
        Gets messages formatted for LLM API calls.

        Returns:
            List of message dicts with role and content
        """
        result = []

        # Add system prompt first
        if self.system_prompt:
            result.append({
                "role": "system",
                "content": self.system_prompt
            })

        # Add summary if we have one and history was compressed
        if self.summary:
            result.append({
                "role": "system",
                "content": f"Previous conversation summary: {self.summary}"
            })

        # Add recent messages
        for msg in self.messages:
            result.append({
                "role": msg.role.value,
                "content": msg.content
            })

        return result

    def _compress_history(self) -> None:
        """
        Compresses older messages to save context window space.

        Keeps:
        - System messages
        - Last 20 messages
        - Creates summary of compressed middle messages
        """
        # Keep system messages
        system_msgs = [m for m in self.messages if m.role == MessageRole.SYSTEM]

        # Keep last 20 messages
        recent_msgs = self.messages[-20:]

        # Middle messages to summarize
        middle_start = len(system_msgs)
        middle_end = len(self.messages) - 20
        middle_msgs = self.messages[middle_start:middle_end]

        if middle_msgs:
            # Create a basic summary (can be enhanced with LLM-based summarization)
            self.summary = self._create_summary(middle_msgs)

        # Combine
        self.messages = system_msgs + recent_msgs

        logger.debug(
            f"Compressed conversation: {middle_end - middle_start} messages summarized"
        )

    def _create_summary(self, messages: List[Message]) -> str:
        """
        Creates a summary of messages.

        This is a basic implementation - can be enhanced with LLM-based summarization.

        Args:
            messages: Messages to summarize

        Returns:
            Summary string
        """
        # Count message types
        user_count = sum(1 for m in messages if m.role == MessageRole.USER)
        assistant_count = sum(1 for m in messages if m.role == MessageRole.ASSISTANT)

        # Extract key topics (simplified - could use NLP)
        topics = set()
        for msg in messages:
            # Simple keyword extraction
            words = msg.content.lower().split()
            # Filter common words
            keywords = [w for w in words if len(w) > 5][:3]
            topics.update(keywords)

        topics_str = ", ".join(list(topics)[:5])

        return (
            f"Earlier conversation: {user_count} user messages, "
            f"{assistant_count} assistant responses. "
            f"Topics discussed: {topics_str}"
        )

    def clear(self) -> None:
        """Clears all context"""
        self.messages.clear()
        self.entities.clear()
        self.intent_history.clear()
        self.summary = None

    def to_dict(self) -> Dict[str, Any]:
        """Serializes context to dict"""
        return {
            "messages": [m.to_dict() for m in self.messages],
            "entities": self.entities,
            "intent_history": self.intent_history,
            "summary": self.summary,
            "max_messages": self.max_messages,
            "system_prompt": self.system_prompt,
            "metadata": self.metadata
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ConversationContext":
        """Deserializes context from dict"""
        ctx = cls(
            messages=[Message.from_dict(m) for m in data.get("messages", [])],
            entities=data.get("entities", {}),
            intent_history=data.get("intent_history", []),
            summary=data.get("summary"),
            max_messages=data.get("max_messages", 50),
            system_prompt=data.get("system_prompt"),
            metadata=data.get("metadata", {})
        )
        return ctx