"""Compaction — token estimation and context management."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from pi_agents.ai.types import Usage


# Rough token estimation: ~0.3 tokens per character (~3.3 chars per token)
CHARS_PER_TOKEN = 3.3


@dataclass
class CompactionSettings:
    """Settings for context compaction."""

    enabled: bool = True
    reserve_tokens: int = 16384
    keep_recent_tokens: int = 20000


DEFAULT_COMPACTION_SETTINGS = CompactionSettings()


@dataclass
class CompactionResult:
    """Result of compaction."""

    summary: str
    first_kept_entry_id: str
    tokens_before: int
    details: Any = None


@dataclass
class ContextUsageEstimate:
    """Estimated token usage."""

    tokens: int
    usage_tokens: int
    trailing_tokens: int
    last_usage_index: int


@dataclass
class CompactionDetails:
    """Details about file operations during compaction."""

    read_files: list[str] = field(default_factory=list)
    modified_files: list[str] = field(default_factory=list)


def estimate_tokens(text: str) -> int:
    """Rough token estimation (~0.3 tokens per character)."""
    if not text:
        return 0
    return max(1, int(len(text) / CHARS_PER_TOKEN))


def estimate_message_tokens(message: Any) -> int:
    """Estimate tokens for a single message."""
    if isinstance(message, dict):
        content = message.get("content", "")
    elif hasattr(message, "content"):
        content = message.content
    else:
        content = str(message)

    if isinstance(content, str):
        return estimate_tokens(content)
    elif isinstance(content, list):
        total = 0
        for part in content:
            if isinstance(part, dict):
                total += estimate_tokens(part.get("text", ""))
            elif hasattr(part, "text"):
                total += estimate_tokens(part.text)
        return total
    return estimate_tokens(str(content))


def calculate_context_tokens(usage: Usage) -> int:
    """Calculate total context tokens from usage."""
    return usage.input + usage.output + usage.cache_read + usage.cache_write


def should_compact(
    context_tokens: int,
    context_window: int,
    settings: CompactionSettings | None = None,
) -> bool:
    """Check if compaction is needed.

    Returns True if: context_tokens + reserve_tokens > context_window - keep_recent_tokens
    """
    s = settings or DEFAULT_COMPACTION_SETTINGS
    if not s.enabled:
        return False
    return context_tokens + s.reserve_tokens > context_window - s.keep_recent_tokens


# Summarization prompt
SUMMARIZATION_SYSTEM_PROMPT = (
    "You are a conversation summarizer. Summarize the conversation so far. "
    "Focus on key decisions, code changes, file operations, and important context. "
    "Do NOT continue the conversation or add new information. Only summarize."
)

COMPACTION_SUMMARY_PREFIX = "The conversation history before this point was compacted into the following summary:\n\n<summary>\n"
COMPACTION_SUMMARY_SUFFIX = "\n</summary>"

BRANCH_SUMMARY_PREFIX = "The following is a summary of a branch that this conversation came back from:\n\n<summary>\n"
BRANCH_SUMMARY_SUFFIX = "</summary>"


def serialize_conversation(messages: list[Any]) -> str:
    """Convert messages to text for summarization."""
    lines: list[str] = []
    for msg in messages:
        role = msg.get("role", "") if isinstance(msg, dict) else getattr(msg, "role", "")
        content = msg.get("content", "") if isinstance(msg, dict) else getattr(msg, "content", "")

        if isinstance(content, list):
            parts: list[str] = []
            for part in content:
                if isinstance(part, dict):
                    parts.append(part.get("text", ""))
                elif hasattr(part, "text"):
                    parts.append(part.text)
            content = "\n".join(parts)

        # Truncate long tool results
        if role == "toolResult" and len(content) > 2000:
            content = content[:2000] + "...[truncated]"

        label = {"user": "User", "assistant": "Assistant", "toolResult": "Tool result"}.get(role, role)
        lines.append(f"[{label}]: {content}")

    return "\n\n".join(lines)
