"""Token estimation helpers for afterTurn budget calculation.

Mirrors the TypeScript estimateSessionTokenCountForAfterTurn logic
from engine.ts, ported to Python for use in Hippo's compaction pipeline.

Issue: #1160 (Hippo W2-9)
"""

from __future__ import annotations

from cortex.hippo.compaction import estimate_tokens

__all__ = ["estimate_tokens", "estimate_session_tokens"]


def _estimate_message_content_tokens(content: str | list) -> int:
    """Estimate tokens for a message's content field.

    Handles both plain string content and structured array content
    (list of text blocks, tool_use blocks, etc.).
    """
    if isinstance(content, str):
        return estimate_tokens(content)

    total = 0
    for block in content:
        if isinstance(block, str):
            total += estimate_tokens(block)
        elif isinstance(block, dict):
            text = block.get("text", "")
            if text:
                total += estimate_tokens(text)
            # tool_use blocks: estimate the input JSON as text
            inp = block.get("input")
            if inp is not None:
                total += estimate_tokens(str(inp))
    return total


def estimate_session_tokens(messages: list[dict]) -> int:
    """Estimate total token count for a list of session messages.

    Mirrors ``estimateSessionTokenCountForAfterTurn`` from engine.ts.

    Each message is inspected for:
    - ``content`` (str or list of blocks) → delegated to content estimator
    - ``command`` / ``output`` → estimated as ``command\\noutput``

    Returns the summed token estimate across all messages.
    """
    total = 0
    for message in messages:
        if "content" in message:
            total += _estimate_message_content_tokens(message["content"])
            continue
        if "command" in message or "output" in message:
            cmd = message.get("command", "") or ""
            out = message.get("output", "") or ""
            if not isinstance(cmd, str):
                cmd = ""
            if not isinstance(out, str):
                out = ""
            total += estimate_tokens(f"{cmd}\n{out}")
    return total
