"""Rolling history summarization for Oclawma.

This module provides intelligent summarization of conversation history
to stay within the 8K context budget. It keeps recent exchanges verbatim
while summarizing older exchanges using LLM.
"""

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Optional

if TYPE_CHECKING:
    from oclawma.providers.base import BaseProvider
    from oclawma.session import ConversationHistory, SessionMessage


@dataclass
class SummaryConfig:
    """Configuration for rolling history summarization.

    Attributes:
        keep_exchanges: Number of recent exchanges to keep verbatim (default: 3)
        trigger_threshold: Token count at which to trigger summarization (default: 6144 - 75%)
        max_summary_tokens: Maximum tokens to use for LLM summarization (default: 500)
        extract_facts: Whether to extract key facts during summarization (default: True)
        preserve_critical: Whether to preserve critical information (default: True)
        min_messages_to_summarize: Minimum messages before summarization triggers (default: 6)
    """

    keep_exchanges: int = 3
    trigger_threshold: int = 6144  # 75% of 8K
    max_summary_tokens: int = 500
    extract_facts: bool = True
    preserve_critical: bool = True
    min_messages_to_summarize: int = 6

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "SummaryConfig":
        """Create config from dictionary."""
        return cls(
            keep_exchanges=data.get("keep_exchanges", 3),
            trigger_threshold=data.get("trigger_threshold", 6144),
            max_summary_tokens=data.get("max_summary_tokens", 500),
            extract_facts=data.get("extract_facts", True),
            preserve_critical=data.get("preserve_critical", True),
            min_messages_to_summarize=data.get("min_messages_to_summarize", 6),
        )

    def to_dict(self) -> dict[str, Any]:
        """Convert config to dictionary."""
        return {
            "keep_exchanges": self.keep_exchanges,
            "trigger_threshold": self.trigger_threshold,
            "max_summary_tokens": self.max_summary_tokens,
            "extract_facts": self.extract_facts,
            "preserve_critical": self.preserve_critical,
            "min_messages_to_summarize": self.min_messages_to_summarize,
        }


@dataclass
class SummaryResult:
    """Result of a summarization operation.

    Attributes:
        summary: The generated summary text
        facts: List of extracted key facts
        messages_summarized: Number of messages summarized
        tokens_saved: Estimated tokens saved
        original_tokens: Original token count
        new_tokens: New token count after summarization
    """

    summary: str
    facts: list[str] = field(default_factory=list)
    messages_summarized: int = 0
    tokens_saved: int = 0
    original_tokens: int = 0
    new_tokens: int = 0

    @property
    def compression_ratio(self) -> float:
        """Calculate compression ratio."""
        if self.original_tokens == 0:
            return 1.0
        return self.new_tokens / self.original_tokens


@dataclass
class FactStore:
    """Store for extracted facts from conversation history.

    This ensures critical information is never lost even after summarization.
    """

    facts: list[dict[str, Any]] = field(default_factory=list)
    max_facts: int = 50

    def add_fact(
        self,
        fact: str,
        source: str = "summary",
        importance: str = "normal",
        metadata: Optional[dict[str, Any]] = None,
    ) -> None:
        """Add a fact to the store.

        Args:
            fact: The fact text
            source: Source of the fact (e.g., "summary", "user", "tool")
            importance: Importance level ("critical", "high", "normal", "low")
            metadata: Optional metadata about the fact
        """
        self.facts.append(
            {
                "fact": fact,
                "source": source,
                "importance": importance,
                "metadata": metadata or {},
            }
        )

        # Trim if exceeding max
        if len(self.facts) > self.max_facts:
            # Keep critical facts, remove oldest normal facts
            critical = [f for f in self.facts if f["importance"] == "critical"]
            high = [f for f in self.facts if f["importance"] == "high"]
            normal = [f for f in self.facts if f["importance"] == "normal"]
            low = [f for f in self.facts if f["importance"] == "low"]

            # Keep critical and high, trim normal and low
            keep_count = self.max_facts - len(critical) - len(high)
            if keep_count > 0:
                normal = normal[-(keep_count // 2) :] if len(normal) > keep_count // 2 else normal
                low = (
                    low[-(keep_count - len(normal)) :]
                    if len(low) > keep_count - len(normal)
                    else low
                )

            self.facts = critical + high + normal + low

    def get_facts(
        self, importance: Optional[str] = None, source: Optional[str] = None
    ) -> list[dict[str, Any]]:
        """Get facts with optional filtering.

        Args:
            importance: Filter by importance level
            source: Filter by source

        Returns:
            List of matching facts
        """
        result = self.facts

        if importance:
            result = [f for f in result if f["importance"] == importance]

        if source:
            result = [f for f in result if f["source"] == source]

        return result

    def get_critical_facts(self) -> list[str]:
        """Get all critical facts as strings."""
        return [f["fact"] for f in self.facts if f["importance"] == "critical"]

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

    def to_text(self, include_importance: bool = True) -> str:
        """Convert facts to text format for LLM context.

        Args:
            include_importance: Whether to include importance markers

        Returns:
            Formatted text of all facts
        """
        if not self.facts:
            return ""

        lines = ["Key facts from previous conversation:"]
        for fact in self.facts:
            prefix = ""
            if include_importance:
                if fact["importance"] == "critical":
                    prefix = "[CRITICAL] "
                elif fact["importance"] == "high":
                    prefix = "[IMPORTANT] "
            lines.append(f"- {prefix}{fact['fact']}")

        return "\n".join(lines)


class RollingHistorySummarizer:
    """Manages rolling summarization of conversation history.

    This class keeps recent exchanges verbatim while intelligently summarizing
    older exchanges to stay within the context budget. It extracts and preserves
    critical information to ensure nothing important is lost.
    """

    def __init__(
        self,
        provider: Optional["BaseProvider"] = None,
        model: Optional[str] = None,
        config: Optional[SummaryConfig] = None,
        on_summarize: Optional[Callable[[SummaryResult], None]] = None,
    ):
        """Initialize the rolling history summarizer.

        Args:
            provider: LLM provider for summarization (uses session provider if None)
            model: Model name to use for summarization
            config: Summarization configuration
            on_summarize: Optional callback when summarization occurs
        """
        self.provider = provider
        self.model = model
        self.config = config or SummaryConfig()
        self.on_summarize = on_summarize
        self.fact_store = FactStore()
        self._summarization_count = 0

    def should_summarize(self, history: "ConversationHistory") -> bool:
        """Check if history should be summarized based on config triggers.

        Args:
            history: The conversation history

        Returns:
            True if summarization should be triggered
        """
        messages = history.get_messages()

        # Don't summarize if too few messages
        if len(messages) < self.config.min_messages_to_summarize:
            return False

        # Calculate messages to keep (user+assistant pairs = exchanges)
        # Each exchange = 2 messages (user + assistant)
        messages_to_keep = self.config.keep_exchanges * 2

        # Don't summarize if there's nothing to summarize
        if len(messages) <= messages_to_keep + 1:  # +1 for system message
            return False

        # Check token threshold
        estimated_tokens = history.estimate_tokens()
        return estimated_tokens >= self.config.trigger_threshold

    async def summarize(
        self,
        history: "ConversationHistory",
        provider: Optional["BaseProvider"] = None,
        model: Optional[str] = None,
    ) -> SummaryResult:
        """Summarize conversation history, keeping recent exchanges verbatim.

        Args:
            history: The conversation history to summarize
            provider: LLM provider (uses self.provider if None)
            model: Model name (uses self.model if None)

        Returns:
            SummaryResult with details about the summarization
        """
        prov = provider or self.provider
        mdl = model or self.model or "qwen2.5:3b"

        messages = history.get_messages()
        original_tokens = history.estimate_tokens()

        # Identify messages to keep and summarize
        # Keep system message + last N exchanges
        messages_to_keep = self._get_messages_to_keep(messages)
        messages_to_summarize = self._get_messages_to_summarize(messages)

        if not messages_to_summarize:
            return SummaryResult(
                summary="",
                messages_summarized=0,
                tokens_saved=0,
                original_tokens=original_tokens,
                new_tokens=original_tokens,
            )

        # Generate summary
        summary_text = ""
        extracted_facts: list[str] = []

        if prov:
            summary_text, extracted_facts = await self._generate_summary_with_llm(
                messages_to_summarize, prov, mdl
            )
        else:
            summary_text = self._generate_simple_summary(messages_to_summarize)

        # Add facts to fact store
        for fact in extracted_facts:
            self.fact_store.add_fact(fact, source="summary", importance="normal")

        # Preserve critical information
        if self.config.preserve_critical:
            critical_facts = self._extract_critical_info(messages_to_summarize)
            for fact in critical_facts:
                self.fact_store.add_fact(fact, source="conversation", importance="critical")

        # Rebuild history
        history.clear()

        # Add system message if it was there
        system_msg = next((m for m in messages if m.role == "system"), None)
        if system_msg:
            history.add_message(
                role="system",
                content=system_msg.content,
                metadata=system_msg.metadata,
            )

        # Add summary as context
        if summary_text:
            context_parts = [f"[Previous conversation summarized]: {summary_text}"]

            # Add fact store info if we have facts
            facts_text = self.fact_store.to_text(include_importance=True)
            if facts_text:
                context_parts.append(f"\n{facts_text}")

            history.add_message(
                role="assistant",
                content="\n\n".join(context_parts),
                metadata={"type": "summary", "summarization_num": self._summarization_count + 1},
            )

        # Add back kept messages (non-system)
        for msg in messages_to_keep:
            if msg.role != "system":
                history.add_message(
                    role=msg.role,
                    content=msg.content,
                    metadata=msg.metadata,
                )

        new_tokens = history.estimate_tokens()
        tokens_saved = original_tokens - new_tokens

        self._summarization_count += 1

        result = SummaryResult(
            summary=summary_text,
            facts=extracted_facts,
            messages_summarized=len(messages_to_summarize),
            tokens_saved=tokens_saved,
            original_tokens=original_tokens,
            new_tokens=new_tokens,
        )

        if self.on_summarize:
            self.on_summarize(result)

        return result

    def _get_messages_to_keep(self, messages: list["SessionMessage"]) -> list["SessionMessage"]:
        """Get the messages to keep verbatim.

        Args:
            messages: All messages

        Returns:
            List of messages to keep
        """
        # Count exchanges (user-assistant pairs)
        exchanges = []
        current_exchange = []

        for msg in messages:
            if msg.role == "system":
                continue

            current_exchange.append(msg)

            # Complete exchange when we have assistant response
            if msg.role == "assistant":
                exchanges.append(current_exchange)
                current_exchange = []

        # Keep last N complete exchanges
        keep_exchanges = (
            exchanges[-self.config.keep_exchanges :]
            if len(exchanges) >= self.config.keep_exchanges
            else exchanges
        )

        # Flatten
        kept = []
        for exchange in keep_exchanges:
            kept.extend(exchange)

        # Also include incomplete exchange if present
        if current_exchange:
            kept.extend(current_exchange)

        return kept

    def _get_messages_to_summarize(
        self, messages: list["SessionMessage"]
    ) -> list["SessionMessage"]:
        """Get messages that should be summarized.

        Args:
            messages: All messages

        Returns:
            List of messages to summarize
        """
        kept = self._get_messages_to_keep(messages)
        kept_ids = {id(m) for m in kept}

        # Return non-system messages not in kept list
        return [m for m in messages if m.role != "system" and id(m) not in kept_ids]

    async def _generate_summary_with_llm(
        self,
        messages: list["SessionMessage"],
        provider: "BaseProvider",
        model: str,
    ) -> tuple[str, list[str]]:
        """Generate summary using LLM.

        Args:
            messages: Messages to summarize
            provider: LLM provider
            model: Model name

        Returns:
            Tuple of (summary text, extracted facts)
        """
        from oclawma.providers.base import CompletionRequest, Message

        # Format conversation for summarization
        conversation_text = self._format_messages_for_summary(messages)

        # Build prompt based on config
        prompt_parts = [
            "Summarize the following conversation concisely.",
            "Focus on: key facts, decisions made, and important context.",
        ]

        if self.config.extract_facts:
            prompt_parts.append(
                "After the summary, list 3-5 key facts extracted from the conversation, "
                "one per line, prefixed with 'FACT:'."
            )

        prompt_parts.append(f"\n\nConversation:\n{conversation_text}\n\nSummary:")

        prompt = "\n".join(prompt_parts)

        try:
            request = CompletionRequest(
                messages=[Message(role="user", content=prompt)],
                model=model,
                max_tokens=self.config.max_summary_tokens,
                temperature=0.3,
            )

            response = await provider.complete(request)
            content = response.content.strip()

            # Parse summary and facts
            summary_text = content
            facts: list[str] = []

            # Extract FACT: lines - handle both same-line and separate-line cases
            if "FACT:" in content:
                # First, split by FACT: to separate summary from facts
                parts = content.split("FACT:", 1)
                summary_text = parts[0].strip()

                # Process the facts section
                if len(parts) > 1:
                    facts_section = "FACT:" + parts[1]
                    lines = facts_section.split("\n")

                    for line in lines:
                        stripped = line.strip()
                        if stripped.startswith("FACT:"):
                            fact = stripped[5:].strip()
                            if fact:
                                facts.append(fact)
                        elif (
                            stripped
                            and not stripped.startswith("-")
                            and not stripped.startswith("*")
                        ):
                            # Continue collecting facts if they don't have bullet prefix
                            facts.append(stripped)

            return summary_text, facts

        except Exception:
            # Fall back to simple summary on error
            return self._generate_simple_summary(messages), []

    def _generate_simple_summary(self, messages: list["SessionMessage"]) -> str:
        """Generate a simple summary without LLM.

        Args:
            messages: Messages to summarize

        Returns:
            Simple summary text
        """
        # Count messages by role
        user_msgs = [m for m in messages if m.role == "user"]
        assistant_msgs = [m for m in messages if m.role == "assistant"]
        tool_msgs = [m for m in messages if m.metadata.get("type") == "tool_result"]

        parts = [
            f"Conversation with {len(user_msgs)} user messages, "
            f"{len(assistant_msgs)} assistant responses",
        ]

        if tool_msgs:
            parts.append(f", and {len(tool_msgs)} tool executions")

        # Include first user message as context
        if user_msgs:
            first = user_msgs[0].content[:100]
            parts.append(
                f". Started with: \"{first}{'...' if len(user_msgs[0].content) > 100 else ''}\""
            )

        return "".join(parts)

    def _format_messages_for_summary(self, messages: list["SessionMessage"]) -> str:
        """Format messages for summarization prompt.

        Args:
            messages: Messages to format

        Returns:
            Formatted text
        """
        lines = []
        for msg in messages:
            if msg.metadata.get("type") == "tool_result":
                lines.append(f"TOOL ({msg.metadata.get('tool', 'unknown')}): {msg.content[:200]}")
            else:
                content = msg.content[:500]  # Limit length
                if len(msg.content) > 500:
                    content += "..."
                lines.append(f"{msg.role.upper()}: {content}")
        return "\n\n".join(lines)

    def _extract_critical_info(self, messages: list["SessionMessage"]) -> list[str]:
        """Extract critical information that must be preserved.

        Args:
            messages: Messages to analyze

        Returns:
            List of critical facts
        """
        critical_facts = []

        # Patterns that indicate critical information
        # NOTE: Exclude sensitive keywords (password, token, key, secret)
        # to avoid persisting credentials in the fact store
        critical_patterns = [
            "error",
            "failed",
            "failure",
            "success",
            "completed",
            "important",
            "critical",
            "must",
            "remember",
            "note:",
            "warning",
        ]

        for msg in messages:
            content_lower = msg.content.lower()

            # Check for critical patterns
            for pattern in critical_patterns:
                if pattern in content_lower:
                    # Extract the sentence or relevant part
                    sentences = msg.content.split(".")
                    for sentence in sentences:
                        if pattern in sentence.lower():
                            fact = sentence.strip()
                            if len(fact) > 10 and fact not in critical_facts:
                                critical_facts.append(fact)
                            break

            # Check for tool results with errors
            if msg.metadata.get("type") == "tool_result" and (
                "error" in content_lower or "failed" in content_lower
            ):
                tool_name = msg.metadata.get("tool", "unknown")
                critical_facts.append(f"Tool '{tool_name}' failed: {msg.content[:100]}")

        # Limit to unique, non-empty facts
        return list(dict.fromkeys(f for f in critical_facts if f))[:10]

    def get_summary_stats(self) -> dict[str, Any]:
        """Get statistics about summarization.

        Returns:
            Dictionary with summarization statistics
        """
        return {
            "summarization_count": self._summarization_count,
            "facts_stored": len(self.fact_store.facts),
            "critical_facts": len(self.fact_store.get_facts(importance="critical")),
            "config": self.config.to_dict(),
        }

    def reset(self) -> None:
        """Reset the summarizer state."""
        self.fact_store.clear()
        self._summarization_count = 0
