"""按消息条数与粗算 token 预算裁剪历史，用于启动时恢复上下文。"""

from __future__ import annotations

import json
from typing import Any

from aicd.agent.chat_serialize import estimate_message_tokens
from aicd.db.store import ChatMessageRow


def _tool_call_ids_from_assistant(msg: dict[str, Any]) -> list[str]:
    tcs = msg.get("tool_calls") or []
    if not isinstance(tcs, list):
        return []
    ids: list[str] = []
    for tc in tcs:
        if not isinstance(tc, dict):
            continue
        tid = tc.get("id")
        if tid:
            ids.append(str(tid))
    return ids


def openai_message_list_valid(msgs: list[dict[str, Any]]) -> bool:
    """
    检查消息序列是否满足 OpenAI 兼容 API 约定：
    无首部孤立 ``tool``；每条带 ``tool_calls`` 的 ``assistant`` 后须紧跟对应数量的 ``tool``，
    且 ``tool_call_id`` 与 ``tool_calls[].id`` 一致（多重集）。
    """
    if not msgs:
        return True
    if msgs[0].get("role") == "tool":
        return False
    i = 0
    while i < len(msgs):
        role = msgs[i].get("role")
        if role == "tool":
            return False
        if role == "system":
            i += 1
            continue
        if role == "user":
            i += 1
            continue
        if role == "assistant":
            ids = _tool_call_ids_from_assistant(msgs[i])
            if not ids:
                i += 1
                continue
            need = len(ids)
            want = sorted(ids)
            got: list[str] = []
            for k in range(need):
                j = i + 1 + k
                if j >= len(msgs):
                    return False
                tm = msgs[j]
                if tm.get("role") != "tool":
                    return False
                got.append(str(tm.get("tool_call_id") or ""))
            if sorted(got) != want:
                return False
            i += 1 + need
            continue
        return False
    return True


def repair_openai_chat_messages(messages: list[dict[str, Any]]) -> None:
    """
    原地修复：去掉首部孤立 ``tool``；去掉不完整 tool 轮次（含仅跟了部分 ``tool`` 的 ``assistant``）；
    去掉尾部仍含 ``tool_calls`` 且无对应 ``tool`` 的 ``assistant``。
    """
    changed = True
    while changed and messages:
        changed = False
        while messages and messages[0].get("role") == "tool":
            messages.pop(0)
            changed = True
        i = 0
        while i < len(messages):
            m = messages[i]
            if m.get("role") != "assistant":
                i += 1
                continue
            ids = _tool_call_ids_from_assistant(m)
            if not ids:
                i += 1
                continue
            need = len(ids)
            j = i + 1
            tool_roles = 0
            while j < len(messages) and tool_roles < need:
                if messages[j].get("role") != "tool":
                    break
                tool_roles += 1
                j += 1
            if tool_roles < need:
                del messages[i:j]
                changed = True
                continue
            got = [str(messages[i + 1 + k].get("tool_call_id") or "") for k in range(need)]
            if sorted(got) != sorted(ids):
                del messages[i:j]
                changed = True
                continue
            i = j
        while messages:
            last = messages[-1]
            if last.get("role") == "assistant" and _tool_call_ids_from_assistant(last):
                messages.pop()
                changed = True
                continue
            break


def longest_valid_openai_suffix(
    body: list[dict[str, Any]], max_len: int
) -> list[dict[str, Any]]:
    """
    取 ``body`` 的尾部子列表，长度不超过 ``max_len``，且对 API 合法；优先更长、更近的尾部。
    """
    n = len(body)
    if n == 0:
        return []
    cap = min(max_len, n)
    for length in range(cap, 0, -1):
        candidate = body[-length:]
        if openai_message_list_valid(candidate):
            return list(candidate)
    for j in range(n - 1, -1, -1):
        one = [body[j]]
        if openai_message_list_valid(one):
            return one
    return []


def chat_row_to_openai_message(row: ChatMessageRow) -> dict[str, Any] | None:
    if row.role == "tool":
        return {
            "role": "tool",
            "tool_call_id": row.tool_call_id or "",
            "content": row.content or "",
        }
    if row.role == "assistant":
        msg: dict[str, Any] = {"role": "assistant", "content": row.content or ""}
        if row.tool_calls_json:
            try:
                tc = json.loads(row.tool_calls_json)
                if isinstance(tc, list):
                    msg["tool_calls"] = tc
            except json.JSONDecodeError:
                pass
        return msg
    if row.role in ("user", "system"):
        return {"role": row.role, "content": row.content or ""}
    return {"role": row.role, "content": row.content or ""}


def rows_to_messages(rows: list[ChatMessageRow]) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    for r in rows:
        m = chat_row_to_openai_message(r)
        if m:
            out.append(m)
    return out


def trim_messages_to_token_budget(
    messages: list[dict[str, Any]], max_tokens: int
) -> list[dict[str, Any]]:
    """从头部丢弃消息直至估算 token 不超过 ``max_tokens``（至少保留最后一条）。"""
    if not messages:
        return messages
    m = list(messages)
    while len(m) > 1:
        total = sum(estimate_message_tokens(x) for x in m)
        if total <= max_tokens:
            break
        m.pop(0)
    repair_openai_chat_messages(m)
    return m


def build_hydrated_messages(
    *,
    rows: list[ChatMessageRow],
    truncated_by_count: bool,
    latest_summary: str | None,
    context_budget_tokens: int,
) -> list[dict[str, Any]]:
    msgs = rows_to_messages(rows)
    if truncated_by_count and latest_summary:
        msgs.insert(
            0,
            {
                "role": "user",
                "content": "[历史摘记 — 更早轮次未载入上下文]\n" + latest_summary.strip(),
            },
        )
    return trim_messages_to_token_budget(msgs, context_budget_tokens)
