"""Conversation manager with automatic summarization for unlimited context."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any
from datetime import datetime


@dataclass
class ConversationMessage:
    """A message in the conversation."""
    role: str  # "user" or "assistant"
    content: str
    timestamp: datetime = field(default_factory=datetime.now)
    task: str | None = None  # Original task if this was a user request
    files_modified: list[str] = field(default_factory=list)
    tools_used: list[str] = field(default_factory=list)
    outcome: str | None = None  # "success", "partial", "failure"

    def estimate_tokens(self) -> int:
        """Rough token estimate (4 chars per token)."""
        return len(self.content) // 4


@dataclass
class ConversationSummary:
    """A summary of older conversation history."""
    content: str
    messages_summarized: int
    original_token_count: int
    created_at: datetime = field(default_factory=datetime.now)

    def estimate_tokens(self) -> int:
        return len(self.content) // 4


class ConversationManager:
    """
    Manages conversation history with automatic summarization.

    Instead of keeping a fixed number of messages, this uses Claude to
    summarize older parts of the conversation when context grows too long.
    This provides effectively unlimited context through compression.

    Structure:
    - summary: Compressed history of older exchanges
    - recent_messages: Last N messages kept verbatim

    When recent_messages grows too large, older messages are summarized
    and merged into the summary.
    """

    # Token thresholds
    MAX_RECENT_TOKENS = 8000  # Keep recent messages under this
    TARGET_SUMMARY_TOKENS = 2000  # Target size for summary
    MIN_MESSAGES_TO_KEEP = 4  # Always keep at least this many recent messages

    def __init__(self, claude_client=None):
        """
        Initialize the conversation manager.

        Args:
            claude_client: Optional Claude client for summarization.
                          If not provided, uses simple truncation.
        """
        self._claude = claude_client
        self._summary: ConversationSummary | None = None
        self._recent_messages: list[ConversationMessage] = []
        self._total_messages: int = 0

    def add_user_message(self, content: str, task: str | None = None) -> None:
        """Add a user message to the conversation."""
        msg = ConversationMessage(
            role="user",
            content=content,
            task=task or content,
        )
        self._recent_messages.append(msg)
        self._total_messages += 1

    def add_assistant_message(
        self,
        content: str,
        files_modified: list[str] | None = None,
        tools_used: list[str] | None = None,
        outcome: str | None = None,
    ) -> None:
        """Add an assistant message to the conversation."""
        msg = ConversationMessage(
            role="assistant",
            content=content,
            files_modified=files_modified or [],
            tools_used=tools_used or [],
            outcome=outcome,
        )
        self._recent_messages.append(msg)
        self._total_messages += 1

    def _estimate_recent_tokens(self) -> int:
        """Estimate tokens in recent messages."""
        return sum(m.estimate_tokens() for m in self._recent_messages)

    async def maybe_summarize(self) -> bool:
        """
        Check if summarization is needed and perform it.

        Returns True if summarization was performed.
        """
        recent_tokens = self._estimate_recent_tokens()

        if recent_tokens <= self.MAX_RECENT_TOKENS:
            return False

        if len(self._recent_messages) <= self.MIN_MESSAGES_TO_KEEP:
            return False

        # Find split point - keep last MIN_MESSAGES_TO_KEEP messages
        split_idx = len(self._recent_messages) - self.MIN_MESSAGES_TO_KEEP
        messages_to_summarize = self._recent_messages[:split_idx]

        if not messages_to_summarize:
            return False

        # Generate summary
        new_summary = await self._generate_summary(messages_to_summarize)

        # Update state
        self._recent_messages = self._recent_messages[split_idx:]

        # Merge with existing summary if present
        if self._summary:
            merged_content = f"{self._summary.content}\n\n{new_summary}"
            self._summary = ConversationSummary(
                content=merged_content,
                messages_summarized=self._summary.messages_summarized + len(messages_to_summarize),
                original_token_count=self._summary.original_token_count + sum(
                    m.estimate_tokens() for m in messages_to_summarize
                ),
            )

            # If merged summary is too long, re-summarize it
            if self._summary.estimate_tokens() > self.TARGET_SUMMARY_TOKENS * 2:
                self._summary = await self._compress_summary(self._summary)
        else:
            self._summary = ConversationSummary(
                content=new_summary,
                messages_summarized=len(messages_to_summarize),
                original_token_count=sum(m.estimate_tokens() for m in messages_to_summarize),
            )

        return True

    async def _generate_summary(self, messages: list[ConversationMessage]) -> str:
        """Generate a summary of messages."""
        if not self._claude:
            # Fallback: simple extraction of key info
            return self._simple_summary(messages)

        # Build prompt for Claude
        messages_text = []
        for msg in messages:
            prefix = "User" if msg.role == "user" else "Assistant"
            messages_text.append(f"{prefix}: {msg.content[:1000]}")
            if msg.files_modified:
                messages_text.append(f"  Files modified: {', '.join(msg.files_modified)}")
            if msg.outcome:
                messages_text.append(f"  Outcome: {msg.outcome}")

        prompt = f"""Summarize this conversation history concisely. Focus on:
1. What tasks were requested
2. What was accomplished
3. Key decisions or outcomes
4. Any important context for future tasks

Keep the summary under 500 words. Use bullet points.

Conversation:
{chr(10).join(messages_text)}

Summary:"""

        try:
            summary = await self._claude.complete(
                prompt=prompt,
                system="You are a conversation summarizer. Be concise and focus on actionable information.",
                max_tokens=800,
            )
            return summary
        except Exception:
            return self._simple_summary(messages)

    def _simple_summary(self, messages: list[ConversationMessage]) -> str:
        """Simple summary without Claude."""
        parts = ["Previous conversation summary:"]

        for msg in messages:
            if msg.role == "user" and msg.task:
                parts.append(f"- Task: {msg.task[:100]}")

            if msg.role == "assistant":
                if msg.outcome:
                    parts.append(f"  Outcome: {msg.outcome}")
                if msg.files_modified:
                    parts.append(f"  Files: {', '.join(msg.files_modified[:5])}")

        return "\n".join(parts)

    async def _compress_summary(self, summary: ConversationSummary) -> ConversationSummary:
        """Compress an overly long summary."""
        if not self._claude:
            # Simple truncation
            content = summary.content[:self.TARGET_SUMMARY_TOKENS * 4]
            return ConversationSummary(
                content=content,
                messages_summarized=summary.messages_summarized,
                original_token_count=summary.original_token_count,
            )

        prompt = f"""Compress this conversation summary to be more concise while keeping key information.
Keep it under 300 words.

Current summary:
{summary.content}

Compressed summary:"""

        try:
            compressed = await self._claude.complete(
                prompt=prompt,
                system="You compress conversation summaries. Be very concise.",
                max_tokens=500,
            )
            return ConversationSummary(
                content=compressed,
                messages_summarized=summary.messages_summarized,
                original_token_count=summary.original_token_count,
            )
        except Exception:
            # Truncate on failure
            return ConversationSummary(
                content=summary.content[:self.TARGET_SUMMARY_TOKENS * 4],
                messages_summarized=summary.messages_summarized,
                original_token_count=summary.original_token_count,
            )

    def get_context_for_prompt(self) -> str:
        """
        Get the full context to include in prompts.

        Returns a string combining the summary and recent messages.
        """
        parts = []

        if self._summary:
            parts.append("## Conversation History (Summarized)")
            parts.append(self._summary.content)
            parts.append("")

        if self._recent_messages:
            parts.append("## Recent Conversation")
            for msg in self._recent_messages:
                prefix = "User" if msg.role == "user" else "Assistant"
                # Include more detail for recent messages
                content = msg.content
                if len(content) > 2000:
                    content = content[:2000] + "..."
                parts.append(f"{prefix}: {content}")

                if msg.files_modified:
                    parts.append(f"  [Files modified: {', '.join(msg.files_modified)}]")

        return "\n".join(parts) if parts else ""

    def clear(self) -> None:
        """Clear all conversation history."""
        self._summary = None
        self._recent_messages = []
        self._total_messages = 0

    def to_dict(self) -> dict:
        """Serialize conversation state to a dictionary."""
        return {
            "summary": {
                "content": self._summary.content,
                "messages_summarized": self._summary.messages_summarized,
                "original_token_count": self._summary.original_token_count,
                "created_at": self._summary.created_at.isoformat(),
            } if self._summary else None,
            "recent_messages": [
                {
                    "role": m.role,
                    "content": m.content,
                    "timestamp": m.timestamp.isoformat(),
                    "task": m.task,
                    "files_modified": m.files_modified,
                    "tools_used": m.tools_used,
                    "outcome": m.outcome,
                }
                for m in self._recent_messages
            ],
            "total_messages": self._total_messages,
        }

    @classmethod
    def from_dict(cls, data: dict, claude_client=None) -> "ConversationManager":
        """Restore conversation state from a dictionary."""
        manager = cls(claude_client)

        if data.get("summary"):
            s = data["summary"]
            manager._summary = ConversationSummary(
                content=s["content"],
                messages_summarized=s["messages_summarized"],
                original_token_count=s["original_token_count"],
                created_at=datetime.fromisoformat(s["created_at"]),
            )

        for m in data.get("recent_messages", []):
            msg = ConversationMessage(
                role=m["role"],
                content=m["content"],
                timestamp=datetime.fromisoformat(m["timestamp"]),
                task=m.get("task"),
                files_modified=m.get("files_modified", []),
                tools_used=m.get("tools_used", []),
                outcome=m.get("outcome"),
            )
            manager._recent_messages.append(msg)

        manager._total_messages = data.get("total_messages", len(manager._recent_messages))
        return manager

    def save_to_file(self, path: str) -> None:
        """Save conversation state to a JSON file."""
        import json
        from pathlib import Path
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w') as f:
            json.dump(self.to_dict(), f, indent=2)

    @classmethod
    def load_from_file(cls, path: str, claude_client=None) -> "ConversationManager":
        """Load conversation state from a JSON file."""
        import json
        from pathlib import Path
        if not Path(path).exists():
            return cls(claude_client)
        with open(path) as f:
            return cls.from_dict(json.load(f), claude_client)

    @property
    def message_count(self) -> int:
        """Total number of messages (including summarized)."""
        return self._total_messages

    @property
    def has_history(self) -> bool:
        """Check if there's any conversation history."""
        return bool(self._summary or self._recent_messages)
