import json
import os
import re
from typing import Callable, Dict, List


def update_folding_state(total_screenshots: int, folded_prefix_k: int, image_max: int, fold_size: int) -> int:
    while (total_screenshots - folded_prefix_k) > image_max:
        folded_prefix_k += fold_size
    if folded_prefix_k > total_screenshots:
        folded_prefix_k = total_screenshots
    return folded_prefix_k


def should_collapse_step(step_num_1based: int, folded_prefix_k: int) -> bool:
    return step_num_1based <= folded_prefix_k


def previous_actions_text(actions: List[str], start_step: int) -> str:
    previous_actions = [
        f"Step {i + 1}: {actions[i]}"
        for i in range(0, min(start_step - 1, len(actions)))
    ]
    return "\n".join(previous_actions) if previous_actions else "None"


def wrap_tool_response(parts: List[Dict]) -> List[Dict]:
    return (
        [{"type": "text", "text": "<tool_response>\n"}]
        + parts
        + [{"type": "text", "text": "\n</tool_response>"}]
    )


def build_messages(
    *,
    system_prompt: str,
    instruction_prompt: str,
    screenshots: List[str],
    responses: List[str],
    start_step: int,
    total_steps: int,
    folded_prefix_k: int,
    collapse_text: str,
    response_transform: Callable[[str], str] = lambda text: text,
) -> List[Dict]:
    messages: List[Dict] = [
        {"role": "system", "content": [{"type": "text", "text": system_prompt}]}
    ]

    for step_num in range(start_step, total_steps + 1):
        is_first_turn = step_num == start_step
        is_collapsed = should_collapse_step(step_num, folded_prefix_k)

        if is_collapsed:
            if is_first_turn:
                user_content = [{"type": "text", "text": instruction_prompt}]
            else:
                user_content = wrap_tool_response([{"type": "text", "text": collapse_text}])
            messages.append({"role": "user", "content": user_content})
        else:
            img_url = f"data:image/png;base64,{screenshots[step_num - 1]}"
            if is_first_turn:
                user_content = [
                    {"type": "image_url", "image_url": {"url": img_url}},
                    {"type": "text", "text": instruction_prompt},
                ]
            else:
                user_content = wrap_tool_response(
                    [{"type": "image_url", "image_url": {"url": img_url}}]
                )
            messages.append({"role": "user", "content": user_content})

        if step_num <= total_steps - 1 and (step_num - 1) < len(responses):
            messages.append(
                {
                    "role": "assistant",
                    "content": [
                        {
                            "type": "text",
                            "text": response_transform(responses[step_num - 1]),
                        }
                    ],
                }
            )

    return messages


def ensure_empty_think_prefix(response: str) -> str:
    text = response or ""
    if re.match(r"^\s*<think>.*?</think>\s*", text, re.DOTALL):
        return text
    return "<think>\n\n</think>\n\n" + text.lstrip("\n")


def sanitize_messages_for_dump(messages: List[Dict]) -> List[Dict]:
    sanitized: List[Dict] = []
    for message in messages:
        cloned = {"role": message.get("role"), "content": []}
        for part in message.get("content", []) or []:
            if isinstance(part, dict) and part.get("type") == "image_url":
                url = ((part.get("image_url") or {}).get("url")) or ""
                if url.startswith("data:image/"):
                    cloned["content"].append(
                        {
                            "type": "image_url",
                            "image_url": {"url": url[:40] + "...<omitted>"},
                        }
                    )
                else:
                    cloned["content"].append(part)
            else:
                cloned["content"].append(part)
        sanitized.append(cloned)
    return sanitized


def dump_debug_messages(messages: List[Dict], filename: str, logger=None) -> None:
    try:
        draft_dir = "./draft/message_cache"
        os.makedirs(draft_dir, exist_ok=True)
        message_file_path = os.path.join(draft_dir, filename)
        with open(message_file_path, "w", encoding="utf-8") as file_obj:
            json.dump(sanitize_messages_for_dump(messages), file_obj, ensure_ascii=False, indent=2)
    except Exception as exc:
        if logger:
            logger.warning("failed to dump debug messages: %s", exc)
