from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any

from niannian.chat.chat_dedup import chat_dedup_key, consolidate_remember_text, dedup_exists_recently
from niannian.chat.remember_gate import _effective_text
from niannian.chat.session import Turn
from niannian.llm.router import LLMRouter
from niannian.llm.types import LLMUnavailable, Message
from niannian.token.budget import resolve_chat_llm_level


@dataclass
class ConsolidationCandidate:
    user_turn: Turn
    assistant_turn: Turn
    text: str
    tags: list[str]
    confidence: float
    skip_reason: str | None = None


def _pair_index(pairs: list[tuple[Turn, Turn]], user_id: str, assistant_id: str) -> int | None:
    for i, (user, assistant) in enumerate(pairs):
        if user.id == user_id and assistant.id == assistant_id:
            return i
    return None


def consolidation_gate(
    pairs: list[tuple[Turn, Turn]],
    *,
    config: dict[str, Any],
    budget: dict[str, Any],
    llm: LLMRouter | None = None,
    left_store: Any | None = None,
    project: str = "global",
) -> list[ConsolidationCandidate]:
    """LLM gate for batch consolidation; rules-only path returns []."""
    if resolve_chat_llm_level(budget) != "heavy":
        return []
    if llm is None or not llm.is_module_configured("left"):
        return []

    chat_cfg = config.get("chat") or {}
    consolidation = chat_cfg.get("consolidation") or {}
    min_confidence = float(consolidation.get("min_batch_confidence", 0.85))
    min_novelty = int(consolidation.get("min_novelty_chars", 20))
    max_candidates = int(consolidation.get("max_candidates_per_batch", 3))
    max_chars = int(chat_cfg.get("max_ingest_chars", 2000))
    dedup_hours = int(config.get("memory", {}).get("dedup_window_hours", 24))

    lines: list[str] = []
    transcript_limit = 8000
    used = 0
    for idx, (user, assistant) in enumerate(pairs):
        block = (
            f"[{idx}] 用户: {_effective_text(user)}\n"
            f"[{idx}] 助手: {_effective_text(assistant)}"
        )
        if used + len(block) > transcript_limit:
            break
        lines.append(block)
        used += len(block) + 1
    if not lines:
        return []

    system = (
        "你是睡眠整合阶段的小忆。只输出值得长期记住的 0~3 条原子事实。"
        "必须拒绝：寒暄、语气词、纯情绪回应、无新信息的确认、临时状态、"
        "对助手外观/能力的闲聊、可从当下对话外轻易推断的废话。"
        "必须保留：用户偏好、承诺、决定、人名关系、需下次还记得的纠错。"
        "输出 JSON 数组：[{\"pair_index\":0,\"should_ingest\":true,\"text\":\"...\","
        "\"tags\":[\"chat\"],\"confidence\":0.9,\"skip_reason\":null}]"
    )
    user_msg = "会话 transcript：\n" + "\n\n".join(lines)

    try:
        result = llm.complete(
            "left",
            [
                Message(role="system", content=system),
                Message(role="user", content=user_msg),
            ],
            json_mode=True,
            temperature=0.2,
        )
        raw_items = json.loads(result.text)
    except (LLMUnavailable, json.JSONDecodeError, TypeError, KeyError):
        return []

    if isinstance(raw_items, dict):
        raw_items = raw_items.get("candidates") or raw_items.get("items") or []
    if not isinstance(raw_items, list):
        return []

    out: list[ConsolidationCandidate] = []
    seen_keys: set[str] = set()
    for item in raw_items:
        if not isinstance(item, dict):
            continue
        if len(out) >= max_candidates:
            break
        if not item.get("should_ingest"):
            continue
        confidence = float(item.get("confidence") or 0)
        if confidence < min_confidence:
            continue
        text = str(item.get("text") or "").strip()
        if len(text) < min_novelty:
            continue
        pair_idx = item.get("pair_index")
        if pair_idx is None:
            continue
        try:
            pair_idx = int(pair_idx)
        except (TypeError, ValueError):
            continue
        if pair_idx < 0 or pair_idx >= len(pairs):
            continue
        user_turn, assistant_turn = pairs[pair_idx]
        dedup_key = chat_dedup_key(text[:max_chars])
        if dedup_key in seen_keys:
            continue
        if left_store is not None and dedup_exists_recently(
            left_store, dedup_key, project, hours=dedup_hours
        ):
            continue
        seen_keys.add(dedup_key)
        tags = [str(t) for t in (item.get("tags") or ["chat"])]
        out.append(
            ConsolidationCandidate(
                user_turn=user_turn,
                assistant_turn=assistant_turn,
                text=text[:max_chars],
                tags=tags,
                confidence=confidence,
                skip_reason=item.get("skip_reason"),
            )
        )
    return out
