"""会话历史：数据库保存完整对话，每次请求按 token 预算带最近的轮次，最近一轮总是带上

裁剪在服务层组装请求时做一次，意图判断、问答和 screen agent 看到的是同一段历史
每轮是 {user, assistant, outcome}，有轮次记录时另带 changes、queries、clarifications（字段同 Turn）：
用户原话和回复转成对话消息，记录由 history_records 放进本次请求的 JSON，两者共用一个预算
"""

import json
from functools import cache

from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.messages.utils import count_tokens_approximately
from loguru import logger

from app.schemas.base import JsonObject
from app.schemas.screen import Screen

RECORD_KEYS = ("changes", "queries", "clarifications")
# 放进请求的单条 SQL 上限（字符）；更长的不带，原文仍在会话记录里
RECORD_SQL_MAX_CHARS = 2000
# 每条消息在请求里的固定开销（角色、分隔），计入历史预算
MESSAGE_OVERHEAD = 4
# 最近一轮超过预算被截短时，回复和用户原话各自至少保留的 token
MIN_KEPT_TOKENS = 200


@cache
def _encoding():
    """o200k 近似各家模型的 token；编码表第一次使用时由 tiktoken 下载并缓存，拿不到时退回按字符估算"""
    try:
        import tiktoken

        return tiktoken.get_encoding("o200k_base")
    except Exception:  # noqa: BLE001 计数只用来裁剪历史，不能让请求失败
        logger.exception("tiktoken 编码表加载失败，历史预算改按字符估算，中文会低估")
        return None


def count_tokens(text: str) -> int:
    encoding = _encoding()
    if encoding is None:
        return count_tokens_approximately([HumanMessage(content=text)])
    return len(encoding.encode(text, disallowed_special=()))


def _record(turn: JsonObject, ago: int, present: set[str] | None) -> JsonObject | None:
    """一轮放进请求的记录，没有的字段不出现；present 是当前大屏的组件 id，None 时按都不在算（只用于计数）

    过长的 SQL 截断了就不能执行，这条查询整条不带，原文仍在会话记录里
    """
    record: JsonObject = {key: turn[key] for key in RECORD_KEYS if turn.get(key)}
    if "queries" in record:
        record["queries"] = [
            query
            for query in record["queries"]
            if len(query["sql"]) <= RECORD_SQL_MAX_CHARS
        ] or None
    if "changes" in record:
        record["changes"] = [
            {**change, "on_screen": _on_screen(change, present or set())}
            for change in record["changes"]
        ]
    record = {key: value for key, value in record.items() if value}
    return {"ago": ago, **record} if record else None


def _on_screen(change: JsonObject, present: set[str]) -> bool:
    if change["type"] == "canvas":
        return True
    return change["node_id"] in present


def turn_tokens(turn: JsonObject) -> int:
    """一轮在请求里的大小：两条消息（含每条的固定开销）和放进请求的记录"""
    tokens = (
        count_tokens(turn["user"])
        + count_tokens(turn["assistant"])
        + MESSAGE_OVERHEAD * 2
    )
    if record := _record(turn, ago=1, present=None):
        tokens += count_tokens(json.dumps(record, ensure_ascii=False))
    return tokens


def history_messages(history: list[JsonObject]) -> list[BaseMessage]:
    """每轮转成一条用户消息和一条助手消息，按时间从早到晚排列；记录不进消息"""
    messages: list[BaseMessage] = []
    for turn in history:
        messages.append(HumanMessage(content=turn["user"]))
        messages.append(AIMessage(content=turn["assistant"]))
    return messages


def trim_history(history: list[JsonObject], max_tokens: int) -> list[JsonObject]:
    """从最近一轮往前保留，消息和记录合计不超过 max_tokens；更早的轮次放不下就整轮不带

    最近一轮总是保留：它单独就超过预算时截短再带上（见 _fit），不能因为一条长回答把整段历史丢掉
    """
    if not history:
        return []
    latest = _fit(history[-1], max_tokens)
    used = turn_tokens(latest)
    kept = [latest]
    for turn in reversed(history[:-1]):
        used += turn_tokens(turn)
        if used > max_tokens:
            break
        kept.append(turn)
    return kept[::-1]


def _fit(turn: JsonObject, max_tokens: int) -> JsonObject:
    """一轮超过预算时先截回复、再截用户原话，都是去掉中间、保留开头和结尾（提议和结论通常在回复末尾）

    记录不截：查询最多 10 条、每条 SQL 不超过 RECORD_SQL_MAX_CHARS，默认预算下放得下
    """
    for key in ("assistant", "user"):
        over = turn_tokens(turn) - max_tokens
        if over <= 0:
            break
        keep = max(count_tokens(turn[key]) - over, MIN_KEPT_TOKENS)
        turn = {**turn, key: truncate_middle(turn[key], keep)}
    return turn


def truncate_middle(text: str, max_tokens: int) -> str:
    """去掉中间、保留开头和结尾各一半，中间写明省略了多少字；结果不超过 max_tokens"""
    total = count_tokens(text)
    if total <= max_tokens:
        return text
    keep = len(text) * max_tokens // total
    while keep > 0:
        head, tail = text[: keep // 2], text[len(text) - (keep - keep // 2) :]
        shortened = f"{head}\n…（中间省略 {len(text) - keep} 字）…\n{tail}"
        if count_tokens(shortened) <= max_tokens:
            return shortened
        keep = keep * 9 // 10
    return f"…（省略 {len(text)} 字）…"


def intent_history(history: list[JsonObject]) -> list[JsonObject]:
    """意图判断看的历史：原话、回复、结果和澄清问答；SQL 和修改记录不给"""
    projected = []
    for turn in history:
        item = {key: turn[key] for key in ("user", "assistant", "outcome")}
        if turn.get("clarifications"):
            item["clarifications"] = turn["clarifications"]
        projected.append(item)
    return projected


def history_records(history: list[JsonObject], screen: Screen) -> list[JsonObject]:
    """保留下来的各轮的记录，最近的在前；ago 是往前数第几轮，1 是上一轮

    on_screen 按当前大屏现算：只表示这个 id 现在还在大屏上，不代表那次修改已经生效
    澄清恢复时可能换了大屏，所以在节点里调用，不在组装请求时算好
    """
    present = {node.id for node in screen.nodes}
    records = []
    for ago, turn in enumerate(reversed(history), start=1):
        if record := _record(turn, ago, present):
            records.append(record)
    return records