"""
Conversation manager for tracking multi-turn chats and context usage.
"""

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

from .context_analyzer import ContextAnalyzer, ContextStatus
from .cost_calculator import CostCalculator
from .models import get_model_config
from .tokenizer import TokenCounter


class MessageRole(Enum):
    """Message roles in a conversation."""

    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    FUNCTION = "function"
    TOOL = "tool"


@dataclass
class Message:
    """A single message in a conversation.

    Attributes:
        role: Message role (system, user, assistant)
        content: Message content
        tokens: Number of tokens in the message
    """

    role: str
    content: str
    tokens: int


@dataclass
class ConversationTurn:
    """A single turn in a conversation."""

    user_message: str
    assistant_response: str
    user_tokens: int
    assistant_tokens: int
    total_tokens: int
    cost: float
    timestamp: datetime
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ConversationStats:
    """Statistics for a conversation."""

    total_turns: int
    total_tokens: int
    total_cost: float
    input_tokens: int
    output_tokens: int
    context_status: ContextStatus
    context_usage_percentage: float
    estimated_turns_remaining: int
    warnings: List[str]


class ConversationManager:
    """
    Manage multi-turn conversations with context tracking.
    """

    def __init__(
        self,
        model_name: str,
        system_message: Optional[str] = None,
        rag_context: Optional[str] = None,
        functions: Optional[List[Dict[str, Any]]] = None,
    ):
        """
        Initialize conversation manager.

        Args:
            model_name: Name of the model
            system_message: Optional system message
            rag_context: Optional RAG/knowledge base context
            functions: Optional function definitions
        """
        self.model_name = model_name
        self.model_config = get_model_config(model_name)
        self.token_counter = TokenCounter(model_name)
        self.context_analyzer = ContextAnalyzer(model_name)
        self.cost_calculator = CostCalculator(model_name)

        # Conversation state
        self.messages: List[Dict[str, Any]] = []
        self.turns: List[ConversationTurn] = []
        self.functions = functions or []
        self.rag_context = rag_context or ""

        # Add system message if provided
        if system_message:
            self.messages.append(
                {
                    "role": MessageRole.SYSTEM.value,
                    "content": system_message,
                }
            )

        # Calculate fixed overhead
        self.fixed_overhead = self._calculate_fixed_overhead()

    def _calculate_fixed_overhead(self) -> int:
        """Calculate fixed token overhead (system, RAG, functions)."""
        overhead = 0

        # System message tokens
        if self.messages:
            for msg in self.messages:
                if msg.get("role") == MessageRole.SYSTEM.value:
                    overhead += self.token_counter.count_tokens(msg.get("content", ""))

        # RAG context tokens
        if self.rag_context:
            overhead += self.token_counter.count_tokens(self.rag_context)

        # Function definition tokens
        if self.functions:
            overhead += self.token_counter.count_function_definitions(self.functions)

        return overhead

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

        Args:
            user_message: User's message
            assistant_response: Assistant's response
            metadata: Optional metadata for this turn

        Returns:
            ConversationTurn object
        """
        # Count tokens
        user_tokens = self.token_counter.count_tokens(user_message)
        assistant_tokens = self.token_counter.count_tokens(assistant_response)
        total_tokens = user_tokens + assistant_tokens

        # Calculate cost
        cost_breakdown = self.cost_calculator.calculate_cost(
            user_tokens,
            assistant_tokens,
        )

        # Create turn object
        turn = ConversationTurn(
            user_message=user_message,
            assistant_response=assistant_response,
            user_tokens=user_tokens,
            assistant_tokens=assistant_tokens,
            total_tokens=total_tokens,
            cost=cost_breakdown.total_cost,
            timestamp=datetime.now(),
            metadata=metadata or {},
        )

        # Add to messages
        self.messages.append(
            {
                "role": MessageRole.USER.value,
                "content": user_message,
            }
        )
        self.messages.append(
            {
                "role": MessageRole.ASSISTANT.value,
                "content": assistant_response,
            }
        )

        # Add to turns
        self.turns.append(turn)

        return turn

    def get_stats(self) -> ConversationStats:
        """
        Get current conversation statistics.

        Returns:
            ConversationStats object
        """
        # Analyze context
        analysis = self.context_analyzer.analyze_messages(
            self.messages,
            functions=self.functions if self.functions else None,
        )

        # Calculate totals
        total_tokens = sum(turn.total_tokens for turn in self.turns)
        total_cost = sum(turn.cost for turn in self.turns)
        input_tokens = sum(turn.user_tokens for turn in self.turns)
        output_tokens = sum(turn.assistant_tokens for turn in self.turns)

        # Estimate remaining turns
        if self.turns:
            avg_tokens_per_turn = total_tokens / len(self.turns)
            available_tokens = self.model_config.context_window - analysis.total_tokens
            estimated_remaining = (
                int(available_tokens / avg_tokens_per_turn) if avg_tokens_per_turn > 0 else 0
            )
        else:
            estimated_remaining = 0

        return ConversationStats(
            total_turns=len(self.turns),
            total_tokens=total_tokens,
            total_cost=total_cost,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            context_status=analysis.status,
            context_usage_percentage=analysis.usage_percentage,
            estimated_turns_remaining=max(0, estimated_remaining),
            warnings=analysis.warnings,
        )

    def can_add_turn(
        self,
        estimated_user_tokens: int,
        estimated_assistant_tokens: int,
    ) -> Dict[str, Any]:
        """
        Check if a new turn can be added without exceeding context.

        Args:
            estimated_user_tokens: Expected tokens in user message
            estimated_assistant_tokens: Expected tokens in assistant response

        Returns:
            Dictionary with feasibility analysis
        """
        # Get current token count
        current_analysis = self.context_analyzer.analyze_messages(
            self.messages,
            functions=self.functions if self.functions else None,
        )

        # Calculate tokens after adding turn
        tokens_after = (
            current_analysis.total_tokens + estimated_user_tokens + estimated_assistant_tokens
        )

        can_add = tokens_after <= self.model_config.max_input_tokens

        return {
            "can_add": can_add,
            "current_tokens": current_analysis.total_tokens,
            "tokens_after": tokens_after,
            "context_window": self.model_config.context_window,
            "max_input_tokens": self.model_config.max_input_tokens,
            "overflow": max(0, tokens_after - self.model_config.max_input_tokens),
            "recommendation": self._get_turn_recommendation(
                can_add, tokens_after, current_analysis.total_tokens
            ),
        }

    def _get_turn_recommendation(
        self,
        can_add: bool,
        tokens_after: int,
        current_tokens: int,
    ) -> str:
        """Get recommendation for adding a turn."""
        if not can_add:
            return "❌ Cannot add turn - would exceed context window. Use summarize_conversation() first."

        usage_after = (tokens_after / self.model_config.context_window) * 100

        if usage_after >= 90:
            return "⚠️ Can add but context is critical (>90%). Consider summarizing soon."
        elif usage_after >= 70:
            return "⚠️ Can add but context is high (>70%). Monitor closely."
        else:
            return "✅ Safe to add turn."

    def summarize_conversation(
        self,
        keep_recent_turns: int = 5,
        summary_template: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Summarize conversation to reduce context usage.

        Args:
            keep_recent_turns: Number of recent turns to keep in full
            summary_template: Optional template for summary

        Returns:
            Dictionary with summarization results
        """
        if len(self.turns) <= keep_recent_turns:
            return {
                "summarized": False,
                "reason": "Not enough turns to summarize",
                "total_turns": len(self.turns),
            }

        # Create summary of older turns
        older_turns = self.turns[:-keep_recent_turns]
        recent_turns = self.turns[-keep_recent_turns:]

        # Generate summary
        if summary_template:
            summary = summary_template
        else:
            summary = f"Previous conversation summary ({len(older_turns)} turns):\n"
            for i, turn in enumerate(older_turns, 1):
                summary += f"Turn {i}: User asked about something, assistant responded.\n"

        summary_tokens = self.token_counter.count_tokens(summary)

        # Calculate original tokens
        original_tokens = sum(turn.total_tokens for turn in older_turns)

        # Rebuild messages list
        new_messages = []

        # Keep system message
        for msg in self.messages:
            if msg.get("role") == MessageRole.SYSTEM.value:
                new_messages.append(msg)

        # Add summary
        new_messages.append(
            {
                "role": MessageRole.SYSTEM.value,
                "content": summary,
            }
        )

        # Add recent turns
        for turn in recent_turns:
            new_messages.append(
                {
                    "role": MessageRole.USER.value,
                    "content": turn.user_message,
                }
            )
            new_messages.append(
                {
                    "role": MessageRole.ASSISTANT.value,
                    "content": turn.assistant_response,
                }
            )

        # Update state
        old_message_count = len(self.messages)
        self.messages = new_messages

        return {
            "summarized": True,
            "original_turns": len(older_turns),
            "kept_turns": len(recent_turns),
            "original_tokens": original_tokens,
            "summary_tokens": summary_tokens,
            "tokens_saved": original_tokens - summary_tokens,
            "old_message_count": old_message_count,
            "new_message_count": len(self.messages),
        }

    def get_context_breakdown(self) -> Dict[str, Any]:
        """
        Get detailed breakdown of context usage.

        Returns:
            Dictionary with context breakdown
        """
        analysis = self.context_analyzer.analyze_messages(
            self.messages,
            functions=self.functions if self.functions else None,
        )

        return {
            "total_tokens": analysis.total_tokens,
            "context_window": analysis.context_window,
            "usage_percentage": analysis.usage_percentage,
            "status": analysis.status.value,
            "breakdown": analysis.input_breakdown,
            "available_for_output": analysis.available_for_output,
            "warnings": analysis.warnings,
            "recommendations": analysis.recommendations,
        }

    def export_conversation(self) -> Dict[str, Any]:
        """
        Export conversation for analysis or storage.

        Returns:
            Dictionary with full conversation data
        """
        return {
            "model_name": self.model_name,
            "messages": self.messages,
            "turns": [
                {
                    "user_message": turn.user_message,
                    "assistant_response": turn.assistant_response,
                    "user_tokens": turn.user_tokens,
                    "assistant_tokens": turn.assistant_tokens,
                    "total_tokens": turn.total_tokens,
                    "cost": turn.cost,
                    "timestamp": turn.timestamp.isoformat(),
                    "metadata": turn.metadata,
                }
                for turn in self.turns
            ],
            "stats": self.get_stats().__dict__,
            "fixed_overhead": self.fixed_overhead,
        }
