"""LLMContextManager — token counting, summarization, and context compaction for LLM conversations."""

from __future__ import annotations

import json
import os
from typing import Any

from anthropic import Anthropic
from anthropic.types.beta import BetaMessageParam

from .task_logger import get_agent_logger
from .utils import (
    APIProvider,
    _maybe_filter_to_n_most_recent_images,
    get_model_name,
)

MAX_FLATTENED_HISTORY_CHARS = 100_000


class LLMContextManager:
    """Manages token counting, summarization, and context compaction.

    Static methods can be used standalone (e.g. to summarize tool results
    before injecting into agent context). Instance methods handle the full
    compaction flow with client state.
    """

    def __init__(
        self,
        name: str,
        model: str,
        summarization_model: str,
        provider: APIProvider,
        client: Any,
        betas: list[str],
        only_n_most_recent_images: int | None = 10,
    ) -> None:
        self.name = name
        self.model = model
        self.counting_model = get_model_name(APIProvider.ANTHROPIC, model)
        self.summarization_model = summarization_model
        self.provider = provider
        self.client = client
        self.betas = betas
        self.only_n_most_recent_images = only_n_most_recent_images
        self.logger = get_agent_logger(self.name)

    # MARK: Token counting

    def _get_counting_client(self) -> Anthropic:
        if not hasattr(self, "_counting_client"):
            self._counting_client = Anthropic(
                api_key=os.environ["ANTHROPIC_API_KEY"],
            )
        return self._counting_client

    def count_tokens(
        self,
        messages: list[BetaMessageParam],
        system_prompt: str,
        tools: list[dict[str, Any]],
    ) -> int:
        """Return accurate input token count via the Anthropic count_tokens API."""
        try:
            client = self._get_counting_client()
            system = [{"type": "text", "text": system_prompt}]

            result = client.beta.messages.count_tokens(
                messages=messages,
                model=self.counting_model,
                system=system,
                tools=tools,
                betas=self.betas,
            )
            return result.input_tokens
        except Exception:
            return self._estimate_tokens(messages)

    @staticmethod
    def _estimate_tokens(messages: list[BetaMessageParam]) -> int:
        """Heuristic fallback: ~4 chars per token for text, ~1600 tokens per image."""
        total = 0
        for msg in messages:
            content = msg.get("content", [])
            if isinstance(content, str):
                total += len(content) // 4
                continue
            if not isinstance(content, list):
                continue
            for block in content:
                if not isinstance(block, dict):
                    continue
                btype = block.get("type")
                if btype == "text":
                    total += len(block.get("text", "")) // 4
                elif btype == "image":
                    total += 1600
                elif btype == "tool_use":
                    total += len(json.dumps(block.get("input", {}))) // 4
                elif btype == "tool_result":
                    result_content = block.get("content", [])
                    if isinstance(result_content, list):
                        for sub_block in result_content:
                            if isinstance(sub_block, dict):
                                if sub_block.get("type") == "text":
                                    total += len(sub_block.get("text", "")) // 4
                                elif sub_block.get("type") == "image":
                                    total += 1600
                    elif isinstance(result_content, str):
                        total += len(result_content) // 4
                elif btype in ("thinking", "redacted_thinking"):
                    total += len(block.get("thinking", "")) // 4
        return total

    # MARK: Summarization

    @staticmethod
    def flatten_messages_to_text(messages: list[BetaMessageParam]) -> str:
        """Produce a text-only representation of messages for summarization."""
        parts: list[str] = []
        for msg in messages:
            role = msg["role"].upper()
            content = msg.get("content", [])
            if isinstance(content, str):
                parts.append(f"{role}: {content[:1000]}")
                continue
            for block in content:
                if not isinstance(block, dict):
                    continue
                btype = block.get("type")
                if btype == "text":
                    parts.append(f"{role}: {block.get('text', '')[:1000]}")
                elif btype == "tool_use":
                    inp = json.dumps(block.get("input", {}))[:300]
                    parts.append(f"TOOL_CALL [{block.get('name')}]: {inp}")
                elif btype == "tool_result":
                    result_content = block.get("content", [])
                    texts = [
                        b.get("text", "")
                        for b in result_content
                        if isinstance(b, dict) and b.get("type") == "text"
                    ]
                    result_text = " ".join(texts)[:500]
                    prefix = "ERROR_RESULT" if block.get("is_error") else "TOOL_RESULT"
                    parts.append(f"{prefix}: {result_text}")

        full_text = "\n".join(parts)
        if len(full_text) > MAX_FLATTENED_HISTORY_CHARS:
            full_text = (
                "...[earlier history truncated]...\n"
                + full_text[-MAX_FLATTENED_HISTORY_CHARS:]
            )
        return full_text

    @staticmethod
    def summarize_history(
        history_text: str,
        client: Any,
        model: str,
        betas: list[str] | None = None,
    ) -> str:
        """Call the API to produce a structured summary of flattened history."""
        prompt = (
            "You are given a log of an AI agent's task execution history on an Ubuntu VM. "
            "Produce a structured summary with two sections:\n\n"
            "## Task\nOne sentence describing the original task.\n\n"
            "## Actions\nList each tool call and its result in this exact format:\n"
            "- tool_name: <max 50 char summary of result>\n\n"
            "Rules:\n"
            "- Each result summary MUST be 50 characters or fewer\n"
            "- Include ALL tool calls from the history\n"
            "- If a tool errored, prefix with ERROR:\n"
            "- End with a ## State section (one sentence on current progress)\n\n"
            "<HISTORY>\n"
            f"{history_text}\n"
            "</HISTORY>"
        )
        try:
            kwargs: dict[str, Any] = {
                "max_tokens": 1024,
                "messages": [{"role": "user", "content": prompt}],
                "model": model,
                "system": [
                    {
                        "type": "text",
                        "text": "You are a concise summarizer of AI agent task logs.",
                    }
                ],
            }
            if betas:
                response = client.beta.messages.create(**kwargs, betas=betas)
            else:
                response = client.messages.create(**kwargs)

            for block in response.content:
                if hasattr(block, "text"):
                    return block.text
                if isinstance(block, dict) and block.get("type") == "text":
                    return block.get("text", "[Summary unavailable]")
        except Exception as e:
            return f"[Summary failed: {e}]"
        return "[Summary unavailable]"

    @staticmethod
    def summarize_tool_results(
        results: list[dict[str, Any]],
        client: Any,
        model: str,
        betas: list[str] | None = None,
    ) -> str:
        """Summarize a list of tool call results for injection into agent context.

        Args:
            results: List of dicts with keys like {"name", "output", "error"}.
        """
        parts: list[str] = []
        for r in results:
            name = r.get("name", "unknown_tool")
            output = r.get("output", "")[:500]
            error = r.get("error", "")
            if error:
                parts.append(f"ERROR [{name}]: {error[:200]}")
            else:
                parts.append(f"TOOL_RESULT [{name}]: {output}")

        history_text = "\n".join(parts)
        return LLMContextManager.summarize_history(history_text, client, model, betas)

    # MARK: Compaction

    @staticmethod
    def strip_thinking_blocks(messages: list[BetaMessageParam]) -> None:
        """Remove thinking/redacted_thinking from all assistant messages in-place."""
        for msg in messages:
            if msg["role"] != "assistant":
                continue
            content = msg.get("content", [])
            if not isinstance(content, list):
                continue
            filtered = [
                b
                for b in content
                if not (
                    isinstance(b, dict)
                    and b.get("type") in ("thinking", "redacted_thinking")
                )
            ]
            msg["content"] = filtered if filtered else [{"type": "text", "text": ""}]

    @staticmethod
    def find_compaction_boundary(messages: list[BetaMessageParam]) -> int | None:
        """Return a split index: messages[:boundary] get summarized, messages[boundary:] are kept.

        Keeps ~20% of messages as the tail. The boundary always lands on an assistant message
        with tool_use blocks so the tail opens with a clean [tool_use → tool_result] cycle.
        """
        n = len(messages)
        if n <= 4:
            return None

        # Target tail covers ~20% of messages, minimum 4
        tail_size = max(4, int(0.2 * n))
        candidate = n - tail_size  # ~80% mark

        def _has_tool_use(msg: BetaMessageParam) -> bool:
            if msg["role"] != "assistant":
                return False
            content = msg.get("content", [])
            if not isinstance(content, list):
                return False
            return any(
                isinstance(b, dict) and b.get("type") == "tool_use" for b in content
            )

        # Walk backward from candidate — maximizes the cold section being compacted
        for i in range(candidate, 0, -1):
            if _has_tool_use(messages[i]):
                return i

        # Walk forward if nothing found behind candidate
        for i in range(candidate + 1, n - 2):
            if _has_tool_use(messages[i]):
                return i

        return None

    @staticmethod
    def aggressive_compaction(
        messages: list[BetaMessageParam],
        only_n_most_recent_images: int | None = 10,
    ) -> None:
        """Apply aggressive size reduction by truncating tool results and text blocks."""
        _maybe_filter_to_n_most_recent_images(
            messages,
            images_to_keep=only_n_most_recent_images,
            min_removal_threshold=1,
        )

        for msg in messages:
            if msg["role"] != "user":
                continue
            content = msg.get("content", [])
            if not isinstance(content, list):
                continue
            for block in content:
                if not isinstance(block, dict):
                    continue
                if block.get("type") == "tool_result":
                    result_content = block.get("content", [])
                    if isinstance(result_content, list):
                        for sub_block in result_content:
                            if (
                                isinstance(sub_block, dict)
                                and sub_block.get("type") == "text"
                            ):
                                text = sub_block.get("text", "")
                                if len(text) > 200:
                                    sub_block["text"] = text[:200] + "...[truncated]"
                elif block.get("type") == "text":
                    text = block.get("text", "")
                    if len(text) > 2000 and "<CONTEXT_SUMMARY>" not in text:
                        block["text"] = text[:2000] + "...[truncated]"

    # MARK: Orchestration

    def compact(
        self,
        messages: list[BetaMessageParam],
        system_prompt: str,
        tools: list[dict[str, Any]],
        token_threshold: int,
    ) -> list[BetaMessageParam]:
        """Run the full compaction flow and return the compacted message list."""
        self.strip_thinking_blocks(messages)

        boundary = self.find_compaction_boundary(messages)
        if boundary is None:
            self.aggressive_compaction(messages, self.only_n_most_recent_images)
            return messages

        cold = messages[:boundary]
        tail = list(messages[boundary:])

        history_text = self.flatten_messages_to_text(cold)
        summary = self.summarize_history(
            history_text, self.client, self.summarization_model, None
        )
        summary_block: dict[str, Any] = {
            "type": "text",
            "text": f"<CONTEXT_SUMMARY>\n{summary}\n</CONTEXT_SUMMARY>",
        }
        self.logger.info(f"<CONTEXT_SUMMARY>\n{summary}\n</CONTEXT_SUMMARY>")

        if tail[0]["role"] == "assistant":
            # Boundary lands on an assistant message — insert a synthetic user message
            # before the tail so the API sees a valid alternating-role sequence.
            tail.insert(0, {"role": "user", "content": [summary_block]})
        else:
            first_content = tail[0]["content"]
            if isinstance(first_content, str):
                first_content = [{"type": "text", "text": first_content}]
            else:
                first_content = list(first_content)
            is_pure_tool_result = bool(first_content) and all(
                isinstance(b, dict) and b.get("type") == "tool_result"
                for b in first_content
            )
            # Pure tool_result user messages have orphaned IDs once their matching tool_use
            # blocks are in the summarized cold section — replace rather than prepend.
            if is_pure_tool_result:
                tail[0]["content"] = [summary_block]
            else:
                tail[0]["content"] = [summary_block] + first_content

        if self.only_n_most_recent_images is not None:
            _maybe_filter_to_n_most_recent_images(
                tail,
                images_to_keep=self.only_n_most_recent_images,
                min_removal_threshold=1,
            )

        # If still too large after standard compaction, escalate
        estimated = self.count_tokens(tail, system_prompt, tools)
        if estimated > 0.8 * token_threshold:
            self.aggressive_compaction(tail, self.only_n_most_recent_images)

        return tail
