"""
Conversation Engine for the Tiny Conversational AI.
Handles conversation memory, context management, and natural dialogue flow.
"""

import re
import time
import uuid
from typing import Optional, List, Dict, Any, Tuple
from dataclasses import dataclass, field
from datetime import datetime
from collections import deque
from enum import Enum

from config import Config, get_config
from utils import (
    get_logger, count_tokens_approx, truncate_text, extract_entities,
    LRUCache, Timer
)

logger = get_logger(__name__)


class Intent(Enum):
    """User intent categories."""
    GREETING = "greeting"
    QUESTION = "question"
    REQUEST = "request"
    FOLLOWUP = "followup"
    CLARIFICATION = "clarification"
    FAREWELL = "farewell"
    STATEMENT = "statement"
    COMMAND = "command"
    UNKNOWN = "unknown"


class Sentiment(Enum):
    """User sentiment categories."""
    POSITIVE = "positive"
    NEGATIVE = "negative"
    NEUTRAL = "neutral"
    FRUSTRATED = "frustrated"
    CURIOUS = "curious"
    EXCITED = "excited"


@dataclass
class Message:
    """A single message in the conversation."""
    role: str  # "user", "assistant", or "system"
    content: str
    timestamp: datetime = field(default_factory=datetime.now)
    tokens: int = 0
    intent: Optional[Intent] = None
    sentiment: Optional[Sentiment] = None
    entities: Dict[str, List[str]] = field(default_factory=dict)
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def __post_init__(self):
        if self.tokens == 0:
            self.tokens = count_tokens_approx(self.content)


@dataclass
class ConversationContext:
    """Tracked context throughout the conversation."""
    # Entities mentioned
    names: List[str] = field(default_factory=list)
    topics: List[str] = field(default_factory=list)
    
    # User information
    user_name: Optional[str] = None
    user_preferences: Dict[str, Any] = field(default_factory=dict)
    
    # Current focus
    current_topic: Optional[str] = None
    last_entity_mentioned: Optional[str] = None
    
    # Conversation state
    turn_count: int = 0
    avg_user_message_length: float = 0.0
    dominant_sentiment: Sentiment = Sentiment.NEUTRAL


class Conversation:
    """
    Manages a single conversation session.
    Handles history, context tracking, and message formatting.
    """
    
    def __init__(
        self,
        session_id: Optional[str] = None,
        config: Optional[Config] = None
    ):
        self.session_id = session_id or str(uuid.uuid4())
        self.config = config or get_config()
        self.conv_config = self.config.conversation
        
        # Message history
        self.messages: deque[Message] = deque(maxlen=self.conv_config.max_history_turns * 2)
        
        # Context tracking
        self.context = ConversationContext()
        
        # System prompt
        self._system_message = Message(
            role="system",
            content=self.conv_config.system_prompt
        )
        
        # Timestamps
        self.created_at = datetime.now()
        self.last_activity = datetime.now()
    
    def add_user_message(self, content: str) -> Message:
        """Add a user message to the conversation."""
        # Analyze the message
        intent = self._detect_intent(content)
        sentiment = self._detect_sentiment(content)
        entities = extract_entities(content)
        
        # Create message
        message = Message(
            role="user",
            content=content,
            intent=intent,
            sentiment=sentiment,
            entities=entities,
        )
        
        self.messages.append(message)
        self._update_context(message)
        self.last_activity = datetime.now()
        
        logger.debug(f"User message - Intent: {intent}, Sentiment: {sentiment}")
        return message
    
    def add_assistant_message(self, content: str) -> Message:
        """Add an assistant message to the conversation."""
        message = Message(
            role="assistant",
            content=content,
        )
        
        self.messages.append(message)
        self.last_activity = datetime.now()
        
        return message
    
    def _detect_intent(self, text: str) -> Intent:
        """Detect the intent of a user message."""
        text_lower = text.lower().strip()
        
        # Greeting patterns
        greetings = ["hello", "hi", "hey", "good morning", "good afternoon", 
                     "good evening", "howdy", "greetings", "what's up", "sup"]
        if any(text_lower.startswith(g) or text_lower == g for g in greetings):
            return Intent.GREETING
        
        # Farewell patterns
        farewells = ["bye", "goodbye", "see you", "take care", "later", 
                     "farewell", "good night", "gotta go"]
        if any(f in text_lower for f in farewells):
            return Intent.FAREWELL
        
        # Question indicators
        question_words = ["what", "why", "how", "when", "where", "who", "which", "can", "could", "would", "is", "are", "do", "does"]
        if text_lower.endswith("?") or any(text_lower.startswith(w) for w in question_words):
            # Check if it's a follow-up question
            if self._is_followup(text):
                return Intent.FOLLOWUP
            return Intent.QUESTION
        
        # Command patterns
        command_words = ["tell me", "show me", "explain", "describe", "list", "give me", "help me"]
        if any(text_lower.startswith(c) for c in command_words):
            return Intent.REQUEST
        
        # Follow-up indicators
        followup_words = ["and", "also", "what about", "how about", "then", "so", "but"]
        if any(text_lower.startswith(f) for f in followup_words):
            return Intent.FOLLOWUP
        
        # Clarification
        clarify_words = ["i mean", "i meant", "no i", "actually", "sorry", "let me rephrase"]
        if any(c in text_lower for c in clarify_words):
            return Intent.CLARIFICATION
        
        # Default to statement
        return Intent.STATEMENT
    
    def _detect_sentiment(self, text: str) -> Sentiment:
        """Detect the sentiment of a user message."""
        text_lower = text.lower()
        
        # Positive indicators
        positive_words = ["thanks", "thank you", "great", "awesome", "perfect", 
                         "amazing", "love", "excellent", "wonderful", "happy",
                         "glad", "appreciate", ":)", "😊", "👍"]
        if any(p in text_lower for p in positive_words):
            return Sentiment.POSITIVE
        
        # Excited indicators
        if text.count("!") >= 2 or any(w in text_lower for w in ["wow", "omg", "amazing"]):
            return Sentiment.EXCITED
        
        # Frustrated indicators
        frustrated_words = ["frustrated", "annoying", "annoyed", "hate", "stupid",
                          "doesn't work", "not working", "broken", "useless", "ugh"]
        if any(f in text_lower for f in frustrated_words):
            return Sentiment.FRUSTRATED
        
        # Negative indicators
        negative_words = ["bad", "terrible", "awful", "wrong", "problem", "issue",
                        "error", "fail", "disappointing", ":(", "😢", "😞"]
        if any(n in text_lower for n in negative_words):
            return Sentiment.NEGATIVE
        
        # Curious indicators
        if "?" in text or any(w in text_lower for w in ["curious", "wondering", "interested"]):
            return Sentiment.CURIOUS
        
        return Sentiment.NEUTRAL
    
    def _is_followup(self, text: str) -> bool:
        """Check if the message is a follow-up to previous context."""
        text_lower = text.lower()
        
        # Pronoun references
        pronouns = ["it", "they", "them", "this", "that", "those", "these", "he", "she"]
        words = text_lower.split()
        has_pronoun_ref = any(w in pronouns for w in words[:3])  # Check first 3 words
        
        # Short messages with pronouns are likely follow-ups
        if has_pronoun_ref and len(words) < 10:
            return True
        
        # "And" or "Also" at start
        if text_lower.startswith(("and ", "also ", "what about ", "how about ")):
            return True
        
        # References to "the" something (the same, the other, etc.)
        if re.search(r"\bthe (same|other|previous|last|first)\b", text_lower):
            return True
        
        return False
    
    def _update_context(self, message: Message):
        """Update conversation context based on new message."""
        self.context.turn_count += 1
        
        # Extract and store entities
        if message.entities.get("names"):
            for name in message.entities["names"]:
                if name not in self.context.names:
                    self.context.names.append(name)
                # Check if it's a user's name introduction
                if self._is_name_introduction(message.content, name):
                    self.context.user_name = name
        
        # Update topic tracking
        # Simple: use nouns and key phrases as topics
        words = message.content.lower().split()
        if len(words) >= 3:
            # Use content words as potential topics
            topic_words = [w for w in words if len(w) > 4 and w.isalpha()]
            if topic_words:
                self.context.current_topic = topic_words[0]
        
        # Update sentiment tracking
        if message.sentiment:
            self.context.dominant_sentiment = message.sentiment
        
        # Track last mentioned entity for pronoun resolution
        if message.entities.get("names"):
            self.context.last_entity_mentioned = message.entities["names"][-1]
    
    def _is_name_introduction(self, text: str, name: str) -> bool:
        """Check if the user is introducing their name."""
        text_lower = text.lower()
        name_lower = name.lower()
        
        patterns = [
            f"i'm {name_lower}",
            f"i am {name_lower}",
            f"my name is {name_lower}",
            f"call me {name_lower}",
            f"this is {name_lower}",
            f"name's {name_lower}",
        ]
        
        return any(p in text_lower for p in patterns)
    
    def resolve_pronouns(self, text: str) -> str:
        """
        Resolve pronouns in text based on conversation context.
        Returns text with pronouns replaced for clarity (for internal use).
        """
        if not self.context.last_entity_mentioned:
            return text
        
        # Simple pronoun resolution
        text_resolved = text
        
        # Replace "it" with last mentioned entity if at start
        if text.lower().startswith("it ") and self.context.last_entity_mentioned:
            text_resolved = self.context.last_entity_mentioned + text[2:]
        
        return text_resolved
    
    def get_formatted_history(
        self,
        max_tokens: Optional[int] = None,
        include_system: bool = True
    ) -> str:
        """
        Get formatted conversation history for model input.
        Uses a chat format suitable for instruction-tuned models.
        """
        max_tokens = max_tokens or self.conv_config.max_history_tokens
        
        # Build formatted messages
        formatted_parts = []
        total_tokens = 0
        
        # Add system prompt first
        if include_system:
            system_formatted = f"<|system|>\n{self._system_message.content}</s>\n"
            formatted_parts.append(system_formatted)
            total_tokens += self._system_message.tokens
        
        # Add messages from history (most recent first, then reverse)
        messages_to_include = []
        
        for message in reversed(self.messages):
            msg_tokens = message.tokens + 10  # Account for formatting
            
            if total_tokens + msg_tokens > max_tokens:
                break
            
            total_tokens += msg_tokens
            messages_to_include.append(message)
        
        # Reverse to get chronological order
        messages_to_include.reverse()
        
        # Format messages using Zephyr/TinyLlama chat template
        for message in messages_to_include:
            if message.role == "user":
                formatted_parts.append(f"<|user|>\n{message.content}</s>\n")
            elif message.role == "assistant":
                formatted_parts.append(f"<|assistant|>\n{message.content}</s>\n")
        
        # Add assistant prompt for response
        formatted_parts.append("<|assistant|>\n")
        
        return "".join(formatted_parts)
    
    def get_chat_messages(self, thinking_mode: bool = False) -> List[Dict[str, str]]:
        """Get messages in chat format for API-style models.
        Limits history to fit within token budget for small models.
        Supports thinking mode for deeper reasoning.
        """
        max_tokens = self.conv_config.max_history_tokens
        
        system_content = self._system_message.content
        
        # Enhance system prompt for thinking mode
        if thinking_mode:
            system_content += (
                "\n\nTHINKING MODE IS ACTIVE: "
                "For this response, think deeply step by step. "
                "First show your reasoning inside <think>...</think> tags, "
                "then provide your clear final answer. "
                "Break down complex problems into smaller parts. "
                "Consider multiple angles before concluding."
            )
        
        messages = [
            {"role": "system", "content": system_content}
        ]
        total_tokens = count_tokens_approx(system_content)
        
        # Collect messages that fit within token budget (most recent first)
        msgs_to_include = []
        for message in reversed(list(self.messages)):
            msg_tokens = message.tokens + 10
            if total_tokens + msg_tokens > max_tokens:
                break
            total_tokens += msg_tokens
            msgs_to_include.append(message)
        
        msgs_to_include.reverse()
        
        for message in msgs_to_include:
            messages.append({
                "role": message.role,
                "content": message.content
            })
        
        return messages
    
    def get_context_summary(self) -> str:
        """Get a brief summary of the conversation context."""
        parts = []
        
        if self.context.user_name:
            parts.append(f"User's name: {self.context.user_name}")
        
        if self.context.current_topic:
            parts.append(f"Current topic: {self.context.current_topic}")
        
        if self.context.names:
            parts.append(f"Mentioned: {', '.join(self.context.names[-5:])}")
        
        if self.context.dominant_sentiment != Sentiment.NEUTRAL:
            parts.append(f"User mood: {self.context.dominant_sentiment.value}")
        
        return " | ".join(parts) if parts else "New conversation"
    
    def should_adapt_response(self) -> Dict[str, Any]:
        """
        Determine how to adapt the response based on context.
        Returns guidance for response generation.
        """
        adaptations = {
            "be_concise": False,
            "be_empathetic": False,
            "be_enthusiastic": False,
            "add_context": False,
            "tone": "friendly",
        }
        
        # If user sends short messages, keep responses concise
        if self.messages:
            recent_user_msgs = [m for m in list(self.messages)[-6:] if m.role == "user"]
            if recent_user_msgs:
                avg_length = sum(len(m.content) for m in recent_user_msgs) / len(recent_user_msgs)
                if avg_length < 30:
                    adaptations["be_concise"] = True
        
        # Adapt to sentiment
        if self.context.dominant_sentiment == Sentiment.FRUSTRATED:
            adaptations["be_empathetic"] = True
            adaptations["tone"] = "calm and helpful"
        elif self.context.dominant_sentiment == Sentiment.EXCITED:
            adaptations["be_enthusiastic"] = True
            adaptations["tone"] = "enthusiastic"
        elif self.context.dominant_sentiment == Sentiment.CURIOUS:
            adaptations["add_context"] = True
            adaptations["tone"] = "informative"
        
        return adaptations
    
    def clear(self):
        """Clear conversation history but keep context."""
        self.messages.clear()
        self.context.turn_count = 0
    
    def reset(self):
        """Fully reset the conversation."""
        self.messages.clear()
        self.context = ConversationContext()
        self.last_activity = datetime.now()
    
    def to_dict(self) -> Dict[str, Any]:
        """Serialize conversation to dictionary."""
        return {
            "session_id": self.session_id,
            "created_at": self.created_at.isoformat(),
            "last_activity": self.last_activity.isoformat(),
            "messages": [
                {
                    "role": m.role,
                    "content": m.content,
                    "timestamp": m.timestamp.isoformat(),
                    "intent": m.intent.value if m.intent else None,
                    "sentiment": m.sentiment.value if m.sentiment else None,
                }
                for m in self.messages
            ],
            "context": {
                "user_name": self.context.user_name,
                "turn_count": self.context.turn_count,
                "current_topic": self.context.current_topic,
                "names": self.context.names,
            },
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any], config: Optional[Config] = None) -> "Conversation":
        """Deserialize conversation from dictionary."""
        conv = cls(session_id=data["session_id"], config=config)
        conv.created_at = datetime.fromisoformat(data["created_at"])
        conv.last_activity = datetime.fromisoformat(data["last_activity"])
        
        for msg_data in data.get("messages", []):
            message = Message(
                role=msg_data["role"],
                content=msg_data["content"],
                timestamp=datetime.fromisoformat(msg_data["timestamp"]),
                intent=Intent(msg_data["intent"]) if msg_data.get("intent") else None,
                sentiment=Sentiment(msg_data["sentiment"]) if msg_data.get("sentiment") else None,
            )
            conv.messages.append(message)
        
        ctx = data.get("context", {})
        conv.context.user_name = ctx.get("user_name")
        conv.context.turn_count = ctx.get("turn_count", 0)
        conv.context.current_topic = ctx.get("current_topic")
        conv.context.names = ctx.get("names", [])
        
        return conv


class ConversationManager:
    """
    Manages multiple conversation sessions.
    Provides session creation, retrieval, and cleanup.
    """
    
    def __init__(self, config: Optional[Config] = None):
        self.config = config or get_config()
        self.conversations: Dict[str, Conversation] = {}
        self._cache = LRUCache(max_size=100, ttl_seconds=3600)
    
    def create(self, session_id: Optional[str] = None) -> Conversation:
        """Create a new conversation session."""
        conv = Conversation(session_id=session_id, config=self.config)
        self.conversations[conv.session_id] = conv
        logger.info(f"Created conversation: {conv.session_id}")
        return conv
    
    def get(self, session_id: str) -> Optional[Conversation]:
        """Get an existing conversation by session ID."""
        return self.conversations.get(session_id)
    
    def get_or_create(self, session_id: str) -> Conversation:
        """Get an existing conversation or create a new one."""
        if session_id in self.conversations:
            return self.conversations[session_id]
        return self.create(session_id)
    
    def delete(self, session_id: str) -> bool:
        """Delete a conversation session."""
        if session_id in self.conversations:
            del self.conversations[session_id]
            logger.info(f"Deleted conversation: {session_id}")
            return True
        return False
    
    def list_sessions(self) -> List[Dict[str, Any]]:
        """List all active conversation sessions."""
        sessions = []
        for session_id, conv in self.conversations.items():
            sessions.append({
                "session_id": session_id,
                "created_at": conv.created_at.isoformat(),
                "last_activity": conv.last_activity.isoformat(),
                "turn_count": conv.context.turn_count,
                "user_name": conv.context.user_name,
            })
        return sessions
    
    def cleanup_inactive(self, max_age_seconds: int = 3600) -> int:
        """Remove inactive conversations."""
        now = datetime.now()
        to_remove = []
        
        for session_id, conv in self.conversations.items():
            age = (now - conv.last_activity).total_seconds()
            if age > max_age_seconds:
                to_remove.append(session_id)
        
        for session_id in to_remove:
            self.delete(session_id)
        
        return len(to_remove)


# =============================================================================
# CONVENIENCE FUNCTIONS
# =============================================================================

_global_manager: Optional[ConversationManager] = None


def get_conversation_manager() -> ConversationManager:
    """Get the global conversation manager instance."""
    global _global_manager
    if _global_manager is None:
        _global_manager = ConversationManager()
    return _global_manager


def create_conversation(session_id: Optional[str] = None) -> Conversation:
    """Create a new conversation."""
    return get_conversation_manager().create(session_id)


def get_conversation(session_id: str) -> Optional[Conversation]:
    """Get an existing conversation."""
    return get_conversation_manager().get(session_id)


if __name__ == "__main__":
    # Test conversation engine
    print("Testing Conversation Engine\n")
    print("=" * 50)
    
    conv = Conversation()
    
    # Test messages
    test_messages = [
        "Hello! My name is Alex.",
        "I'm having trouble with my laptop",
        "It won't turn on",
        "I tried charging it but nothing happens",
        "What should I do?",
    ]
    
    for msg in test_messages:
        user_msg = conv.add_user_message(msg)
        print(f"\nUser: {msg}")
        print(f"  Intent: {user_msg.intent.value}")
        print(f"  Sentiment: {user_msg.sentiment.value}")
        
        # Simulate assistant response
        conv.add_assistant_message("I understand. Let me help you with that.")
    
    print("\n" + "=" * 50)
    print("\nContext Summary:")
    print(conv.get_context_summary())
    
    print("\nUser Name Detected:", conv.context.user_name)
    print("Current Topic:", conv.context.current_topic)
    print("Turn Count:", conv.context.turn_count)
    
    print("\n" + "=" * 50)
    print("\nFormatted History (for model):")
    print(conv.get_formatted_history(max_tokens=500))