"""Context-management policies for the Anthropic computer-use agent.

Four flag-selectable policies for the cache-aware context experiment:

  keep_all     (P0): never edit history. Append-only; maximal prompt-cache reuse.
  image_trim   (P1): keep the last `keep_n_images` screenshots, removing older
                     ones in chunks of `image_chunk_size` (cache-aligned removal:
                     the prefix is only invalidated once per chunk, not per step).
  summarize    (P2): Pointer-style threshold compaction. When the estimated
                     context exceeds 60% of `context_budget_tokens`, flatten the
                     oldest ~80% of messages to text, summarize them with an LLM
                     call, and replace them with a <CONTEXT_SUMMARY> block.
  focus_chain  (P3, focuschain-ab variant): append-only notepad. The model
                     returns a running `task_notes` field in each tool call;
                     the harness appends the latest note at the END of the next
                     user turn. History is NEVER edited — no screenshot
                     dropping, no rewrites; context is keep_all + notes.

Adapted from mm_agents/pointer/llm_context_manager.py (not modified in place).
All summarizer API usage is returned to the caller so it can be logged into
traj.jsonl for honest cost accounting.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Optional

logger = logging.getLogger("desktopenv.agent")

POLICY_KEEP_ALL = "keep_all"
POLICY_IMAGE_TRIM = "image_trim"
POLICY_SUMMARIZE = "summarize"
POLICY_FOCUS_CHAIN = "focus_chain"
ALL_POLICIES = (
    POLICY_KEEP_ALL,
    POLICY_IMAGE_TRIM,
    POLICY_SUMMARIZE,
    POLICY_FOCUS_CHAIN,
)

# 960x540 screenshot ~= (960*540)/750 ~= 691 tokens; round up a little for
# base64/overhead. Only used for threshold estimation, not billing.
IMAGE_TOKEN_ESTIMATE = 700
MAX_FLATTENED_HISTORY_CHARS = 100_000

SUMMARIZE_PROMPT_ADDITION = (
    "CONTEXT COMPACTION: On long tasks, your older conversation history may be "
    "replaced by a <CONTEXT_SUMMARY> block that summarizes everything done so far. "
    "Trust it as an accurate record of your own earlier actions and continue the "
    "task from the current screenshot."
)

FOCUS_CHAIN_PROMPT_ADDITION = (
    "TASK NOTES (IMPORTANT): In EVERY computer tool call, you MUST also fill the "
    "top-level `task_notes` field (alongside `actions`) with a concise running "
    "note (max ~120 words) of your task state: the goal, key facts discovered so "
    "far (exact values, file paths, cell references, settings), what is already "
    "done, and what remains. Your latest task_notes are echoed back to you at "
    "the end of the next user turn — use them to stay focused on the goal and "
    "to avoid re-deriving facts you already discovered. Keep them "
    "self-contained and up to date."
)

TASK_NOTES_SCHEMA_PROPERTY = {
    "type": "string",
    "description": (
        "REQUIRED running note of task state (max ~120 words): goal, key facts "
        "(exact values, paths, cell refs), what is done, what remains. It is "
        "echoed back to you next turn as a persistent focus aid."
    ),
}


def estimate_tokens(messages: list[dict[str, Any]], system_text: str = "") -> int:
    """Heuristic token estimate (~4 chars/token text, IMAGE_TOKEN_ESTIMATE/image)."""
    total = len(system_text) // 4

    def _block_tokens(block: dict[str, Any]) -> int:
        btype = block.get("type")
        if btype == "text":
            return len(block.get("text", "")) // 4
        if btype == "image":
            return IMAGE_TOKEN_ESTIMATE
        if btype == "tool_use":
            return len(json.dumps(block.get("input", {}))) // 4 + 10
        if btype == "tool_result":
            sub = block.get("content", [])
            if isinstance(sub, str):
                return len(sub) // 4
            return sum(_block_tokens(b) for b in sub if isinstance(b, dict))
        if btype in ("thinking", "redacted_thinking"):
            return len(block.get("thinking", "")) // 4
        return 10

    for msg in messages:
        content = msg.get("content", [])
        if isinstance(content, str):
            total += len(content) // 4
        elif isinstance(content, list):
            total += sum(_block_tokens(b) for b in content if isinstance(b, dict))
    return total


def count_images(messages: list[dict[str, Any]]) -> int:
    n = 0
    for msg in messages:
        content = msg.get("content", [])
        if not isinstance(content, list):
            continue
        for block in content:
            if not isinstance(block, dict):
                continue
            if block.get("type") == "image":
                n += 1
            elif block.get("type") == "tool_result":
                sub = block.get("content", [])
                if isinstance(sub, list):
                    n += sum(
                        1
                        for b in sub
                        if isinstance(b, dict) and b.get("type") == "image"
                    )
    return n


# MARK: P2 helpers (adapted from Pointer's LLMContextManager)

def flatten_messages_to_text(messages: list[dict[str, Any]]) -> str:
    """Text-only rendering of messages for the summarizer."""
    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", {}))[:400]
                parts.append(f"TOOL_CALL [{block.get('name')}]: {inp}")
            elif btype == "tool_result":
                sub = block.get("content", [])
                texts = [
                    b.get("text", "")
                    for b in sub
                    if isinstance(b, dict) and b.get("type") == "text"
                ] if isinstance(sub, list) else [str(sub)]
                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


def find_compaction_boundary(messages: list[dict[str, Any]]) -> Optional[int]:
    """Split index: messages[:i] summarized, messages[i:] kept (~20% tail).

    The boundary 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

    tail_size = max(4, int(0.2 * n))
    candidate = n - tail_size

    def _has_tool_use(msg: dict[str, Any]) -> 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
        )

    for i in range(candidate, 0, -1):
        if _has_tool_use(messages[i]):
            return i
    for i in range(candidate + 1, n - 2):
        if _has_tool_use(messages[i]):
            return i
    return None


def summarize_history(
    history_text: str,
    client: Any,
    model: str,
    task_instruction: str,
) -> tuple[str, Optional[dict[str, Any]]]:
    """LLM call producing a structured summary. Returns (summary, usage_dict)."""
    prompt = (
        "You are given a log of an AI agent's task execution history on an Ubuntu VM. "
        "Produce a structured summary with these sections:\n\n"
        "## Actions\nList each tool call and its result:\n"
        "- tool_name: <max 60 char summary of what was done and the result>\n\n"
        "## Key facts\nExact values, file paths, cell references, settings, or UI "
        "locations discovered so far that are needed to finish the task.\n\n"
        "## State\nOne or two sentences on current progress and what remains.\n\n"
        "Rules: include ALL tool calls; prefix errored calls with ERROR:.\n\n"
        "<HISTORY>\n"
        f"{history_text}\n"
        "</HISTORY>"
    )
    try:
        response = client.beta.messages.create(
            max_tokens=1024,
            messages=[{"role": "user", "content": prompt}],
            model=model,
            system=[{
                "type": "text",
                "text": "You are a concise summarizer of AI agent task logs.",
            }],
        )
        usage = getattr(response, "usage", None)
        usage_dict = None
        if usage is not None:
            usage_dict = {
                "input_tokens": getattr(usage, "input_tokens", None),
                "output_tokens": getattr(usage, "output_tokens", None),
                "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", None),
                "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", None),
            }
            raw = usage.model_dump() if hasattr(usage, "model_dump") else {}
            if raw.get("cost_details"):
                usage_dict["cost_details"] = raw.get("cost_details")
        for block in response.content:
            if hasattr(block, "text") and block.text:
                summary = block.text
                break
        else:
            summary = "[Summary unavailable]"
        # Keep the original task verbatim: paraphrased goals lose exact
        # target values, which would conflate summarization cost with a
        # lossy-instruction implementation bug.
        summary = (
            f"## Original task (verbatim)\n{task_instruction}\n\n{summary}"
        )
        return summary, usage_dict
    except Exception as e:  # noqa: BLE001
        logger.warning(f"Context summarization failed: {e}")
        return "", None


def compact_messages(
    messages: list[dict[str, Any]],
    client: Any,
    model: str,
    task_instruction: str,
) -> tuple[list[dict[str, Any]], Optional[dict[str, Any]], bool]:
    """Pointer-style 80/20 compaction. Returns (messages, summarizer_usage, did_compact)."""
    boundary = find_compaction_boundary(messages)
    if boundary is None:
        return messages, None, False

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

    history_text = flatten_messages_to_text(cold)
    summary, usage = summarize_history(history_text, client, model, task_instruction)
    if not summary:
        # Summarizer failed: keep history intact rather than dropping it.
        return messages, usage, False

    summary_block = {
        "type": "text",
        "text": f"<CONTEXT_SUMMARY>\n{summary}\n</CONTEXT_SUMMARY>",
    }
    logger.info(
        f"Context compaction: summarized {boundary} messages "
        f"({len(history_text)} flat chars) into {len(summary)} chars, kept {len(tail)} tail messages"
    )

    if tail[0]["role"] == "assistant":
        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 messages have orphaned tool_use ids once the cold
        # section is gone — replace rather than prepend.
        if is_pure_tool_result:
            tail[0]["content"] = [summary_block]
        else:
            tail[0]["content"] = [summary_block] + first_content

    return tail, usage, True