from __future__ import annotations

import json
import logging
from typing import Any
from zoneinfo import ZoneInfo

from temporal.core.models import ContentBlock, Message

logger = logging.getLogger(__name__)


class DefaultMessageFormatter:
    """Default message formatter that renders messages as timestamped plain text.

    Handles text, tool_call, and tool_result content blocks. Applies
    progressive compression at higher filter levels.
    """

    def __init__(self, timezone: str = "UTC") -> None:
        self.timezone = timezone

    def format_message(self, message: Message, filter_level: int = 0) -> str | None:
        ts = message.timestamp
        if ts is None:
            time_str = "??:??:??"
        else:
            local = ts.astimezone(ZoneInfo(self.timezone))
            time_str = local.strftime("%H:%M:%S") + f".{local.microsecond // 1000:03d}"

        if isinstance(message.content, str):
            if not message.content.strip():
                return None
            if message.role == "system":
                return None
            label = message.role.upper()
            return f"\n{time_str}: [{label}] {message.content}"

        for block in message.content:
            if block.type == "text" and block.text:
                if "This is your memory context" in block.text:
                    return None

        parts: list[str] = []
        for block in message.content:
            formatted = self._format_block(block, filter_level)
            if formatted:
                parts.append(formatted)

        if not parts:
            return None

        return f"\n{time_str}: " + " | ".join(parts)

    def serialize_for_summary(self, message: Message) -> dict[str, Any]:
        return message.to_dict()

    def _format_block(self, block: ContentBlock, filter_level: int) -> str | None:
        if block.type == "text":
            if not block.text or not block.text.strip():
                return None
            return f"[TEXT] {block.text}"

        if block.type == "tool_call":
            name = block.tool_name or "unknown"
            args = block.tool_args or {}
            args_str = json.dumps(args) if args else ""
            if filter_level >= 1 and name == "wait":
                return None
            return f"[TOOL_CALL] {name}({args_str})"

        if block.type == "tool_result":
            output = block.tool_output or ""
            if filter_level >= 1:
                confirmations = [
                    "spoke message to user",
                    "message sent",
                    "action completed",
                    "tool executed successfully",
                ]
                if any(c in output.lower() for c in confirmations):
                    return None

            if filter_level >= 2 and len(output) > 500:
                output = output[:500] + "..."

            return f"[TOOL_RESULT] {output}"

        return None

    def apply_compression_strategies(
        self,
        messages: list[Message],
        available_tokens: int,
        enc,
    ) -> list[str]:
        """Format messages with progressive compression to fit within a token budget."""
        formatted = []
        for msg in messages:
            text = self.format_message(msg, filter_level=0)
            if text:
                formatted.append(text)

        total_tokens = sum(len(enc.encode(msg)) for msg in formatted)
        if total_tokens <= available_tokens:
            return formatted

        final = []
        remaining = available_tokens
        for msg in formatted:
            msg_tokens = len(enc.encode(msg))
            if msg_tokens <= remaining:
                final.append(msg)
                remaining -= msg_tokens
            else:
                if remaining > 100:
                    truncated = self._truncate_to_tokens(msg, remaining - 50, enc)
                    final.append(truncated)
                break
        return final

    @staticmethod
    def _truncate_to_tokens(message: str, max_tokens: int, enc) -> str:
        current_tokens = len(enc.encode(message))
        if current_tokens <= max_tokens:
            return message
        left, right = 0, len(message)
        while left < right:
            mid = (left + right + 1) // 2
            truncated = message[:mid] + "... (truncated due to token limit)"
            if len(enc.encode(truncated)) <= max_tokens:
                left = mid
            else:
                right = mid - 1
        return message[:left] + "... (truncated due to token limit)"
