"""Token tracking and conversation compaction.

Strategy:
- Estimate tokens from character count (rough: 1 token ~ 3 chars, conservative)
- When conversation approaches limit, compact by:
  1. Replacing old tool results with summaries
  2. Dropping oldest message pairs (preserving tool_call/result pairing)
  3. Preserving: system prompt, KRIM.md, recent messages
"""

from __future__ import annotations

import json


# rough estimate: ~3 chars per token (conservative; English ~4, CJK ~2, code ~3)
def estimate_tokens(text: str) -> int:
    if not text:
        return 0
    return len(text) // 3


def estimate_message_tokens(messages: list[dict]) -> int:
    total = 0
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            total += estimate_tokens(content)
        elif isinstance(content, list):
            for block in content:
                if isinstance(block, dict):
                    total += estimate_tokens(json.dumps(block))
                else:
                    total += estimate_tokens(str(block))
        # tool_calls in openai format
        if "tool_calls" in msg:
            total += estimate_tokens(json.dumps(msg["tool_calls"]))
    return total


def needs_compaction(messages: list[dict], max_tokens: int = 120_000, threshold: float = 0.75) -> bool:
    used = estimate_message_tokens(messages)
    return used > max_tokens * threshold


def compact(messages: list[dict], max_tokens: int = 120_000) -> list[dict]:
    """Compact conversation to fit within token budget.

    Preserves: system message (index 0), last N user/assistant exchanges.
    Replaces: old tool results with short summaries.
    """
    if len(messages) <= 4:
        return messages

    system = messages[0] if messages[0].get("role") == "system" else None
    rest = messages[1:] if system else messages[:]

    # phase 1: truncate old tool results
    compacted = []
    for i, msg in enumerate(rest):
        # keep last 6 messages intact
        if i >= len(rest) - 6:
            compacted.append(msg)
            continue

        content = msg.get("content", "")
        role = msg.get("role", "")

        # truncate tool results (claude format: list of tool_result blocks)
        if role == "user" and isinstance(content, list):
            new_content = []
            for block in content:
                if isinstance(block, dict) and block.get("type") == "tool_result":
                    result_text = block.get("content", "")
                    if isinstance(result_text, str) and len(result_text) > 200:
                        block = dict(block)
                        block["content"] = _smart_truncate(result_text)
                new_content.append(block)
            compacted.append(dict(msg, content=new_content))

        # truncate tool results (openai format)
        elif role == "tool":
            if isinstance(content, str) and len(content) > 200:
                compacted.append({**msg, "content": _smart_truncate(content)})
            else:
                compacted.append(msg)

        else:
            compacted.append(msg)

    # phase 2: if still too large, drop oldest message groups
    # We must keep tool_call assistant msgs paired with their tool_result msgs
    # to avoid orphaned references that cause API errors.
    result = ([system] if system else []) + compacted
    while len(result) > 4 and estimate_message_tokens(result) > max_tokens * 0.6:
        # find the oldest droppable group (starting after system msg)
        dropped = _drop_oldest_group(result, start=1)
        if not dropped:
            break  # nothing left to drop

    return result


_ERROR_KEYWORDS = {"error", "Error", "ERROR", "fail", "Fail", "FAIL",
                   "exception", "Exception", "traceback", "Traceback",
                   "denied", "not found", "No such", "exit code"}


def _smart_truncate(text: str, head: int = 80, max_error_lines: int = 5) -> str:
    """Truncate tool output while preserving error-related lines."""
    lines = text.splitlines()
    error_lines = [
        ln for ln in lines
        if any(kw in ln for kw in _ERROR_KEYWORDS)
    ][:max_error_lines]

    parts = [text[:head]]
    if error_lines:
        parts.append("...\n" + "\n".join(error_lines))
    parts.append("... [compacted]")
    return "\n".join(parts)


def _drop_oldest_group(messages: list[dict], start: int) -> bool:
    """Drop the oldest message group starting at `start`.

    Groups:
    - assistant(tool_calls) + following tool_result(s) = drop together
    - plain user + assistant pair = drop together
    - single message = drop alone

    Returns True if something was dropped.
    """
    if start >= len(messages):
        return False

    msg = messages[start]
    role = msg.get("role", "")

    # assistant with tool_calls: drop it + all following tool results
    has_tool_calls = (
        "tool_calls" in msg  # openai format
        or (role == "assistant" and isinstance(msg.get("content"), list)
            and any(b.get("type") == "tool_use" for b in msg["content"] if isinstance(b, dict)))
    )
    if has_tool_calls:
        # count how many following messages are tool results
        end = start + 1
        while end < len(messages):
            r = messages[end].get("role", "")
            c = messages[end].get("content", "")
            is_tool_result = (
                r == "tool"  # openai
                or (r == "user" and isinstance(c, list)
                    and all(isinstance(b, dict) and b.get("type") == "tool_result"
                            for b in c if isinstance(b, dict)))
            )
            if is_tool_result:
                end += 1
            else:
                break
        del messages[start:end]
        return True

    # user + assistant pair
    if role == "user" and start + 1 < len(messages) and messages[start + 1].get("role") == "assistant":
        del messages[start:start + 2]
        return True

    # fallback: drop single message
    del messages[start]
    return True
