from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Any, Iterable


def rough_token_estimate(text: str) -> int:
    """Rough token estimator using word and character heuristics.

    For English text, ~4 chars per token. We take max between
    word-count and char/4 to cover short/long tokens.
    """
    if not text:
        return 0
    words = text.split()
    return max(len(words), int(len(text) / 4))


def sum_message_tokens(contents: Iterable[str]) -> int:
    total = 0
    for c in contents:
        total += rough_token_estimate(c)
    return total


def estimate_chat_prompt_tokens(body: Mapping[str, Any]) -> int:
    """Estimate prompt tokens from a chat completions request body."""
    messages = body.get("messages") or []
    contents: list[str] = []
    for msg in messages:
        if isinstance(msg, dict):
            content = msg.get("content", "")
            if isinstance(content, str):
                contents.append(content)
            elif isinstance(content, list):
                # Handle content arrays (e.g., multimodal)
                for part in content:
                    if isinstance(part, dict):
                        text = part.get("text", "")
                        if isinstance(text, str):
                            contents.append(text)
    return sum_message_tokens(contents)


def estimate_responses_prompt_tokens(body: Mapping[str, Any]) -> int:
    """Estimate prompt tokens from a responses request body.

    Extracts text from:
    - input: string or list of messages
    - messages: list of messages
    - instructions: system prompt
    """
    contents: list[str] = []

    # Extract instructions (system prompt)
    instructions = body.get("instructions")
    if isinstance(instructions, str):
        contents.append(instructions)

    # Extract from input field
    input_val = body.get("input")
    if isinstance(input_val, str):
        contents.append(input_val)
    elif isinstance(input_val, list):
        _extract_message_contents(input_val, contents)

    # Extract from messages field
    messages = body.get("messages")
    if isinstance(messages, list):
        _extract_message_contents(messages, contents)

    return sum_message_tokens(contents)


def _extract_message_contents(
    messages: Sequence[Any],
    out: list[str],
) -> None:
    """Extract text content from a list of messages into out.

    Handles standard text content, function_call arguments, and
    function_call_output results so that token estimates remain
    reasonable for multi-turn tool conversations.
    """
    for msg in messages:
        if not isinstance(msg, dict):
            continue

        # Top-level items in Responses API input (function_call, function_call_output)
        msg_type = msg.get("type")
        if msg_type == "function_call":
            args = msg.get("arguments")
            if isinstance(args, str):
                out.append(args)
            continue
        if msg_type == "function_call_output":
            output = msg.get("output")
            if isinstance(output, str):
                out.append(output)
            continue

        content = msg.get("content")
        if isinstance(content, str):
            out.append(content)
        elif isinstance(content, list):
            for part in content:
                if isinstance(part, dict):
                    text = part.get("text")
                    if isinstance(text, str):
                        out.append(text)
                    # Also count function call arguments
                    args = part.get("arguments")
                    if isinstance(args, str):
                        out.append(args)
                    # Also count function call output
                    output = part.get("output")
                    if isinstance(output, str):
                        out.append(output)
                elif isinstance(part, str):
                    out.append(part)
