"""Agent Team 的轻量持久化层。

这里没有引入数据库，
而是参考轻量 file store 设计：
1. 数据结构尽量稳定；
2. 使用原子写覆盖，避免半写状态；
3. 单文件规模保持小而可读，便于排查与测试。
"""

from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Any

from nanobot.agent.run_result import contains_placeholder_summary, has_meaningful_summary
from nanobot.agent_team.types import (
    BridgeResult,
    ExecutionMode,
    ProcedureRecord,
    RunRecord,
    now_iso,
)

# ASCII token 用于英文/agent id/命令片段匹配。
_ASCII_TOKEN_RE = re.compile(r"[a-z0-9_:-]+")
# 中文任务没有自然空格，这里退而求其次按单字切分，保证最小可匹配能力。
_CJK_CHAR_RE = re.compile(r"[\u4e00-\u9fff]")


def _memory_root(workspace: Path) -> Path:
    """返回 agent team memory 根目录。

    Demo 输出：
        `/workspace/agent_team`
    """
    # 独立目录便于用户直接查看 procedure/runs 文件，不和其他 runtime 状态混在一起。
    root = workspace / "agent_team"
    root.mkdir(parents=True, exist_ok=True)
    return root


def _load_json(path: Path, default: Any) -> Any:
    """从磁盘加载 JSON；损坏或不存在时回退到默认值。

    Demo 输出：
        `[]`
    """
    # agent team memory 不应因为单个文件损坏就拖垮主链路，所以统一做软失败。
    if not path.exists():
        return default
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError, json.JSONDecodeError):
        return default


def _atomic_write_json(path: Path, payload: Any) -> None:
    """把 JSON 原子写入目标路径。

    Demo 输出：
        `None`
    """
    # 先写临时文件再 `os.replace`，这样即使进程中断也不会留下半截 JSON。
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_suffix(path.suffix + ".tmp")
    tmp_path.write_text(
        json.dumps(payload, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    os.replace(str(tmp_path), str(path))


def task_tokens(text: str) -> list[str]:
    """把任务文本压成可匹配的轻量 token 列表。

    Demo 输出：
        `["生成", "周报", "writer-agent", "publish"]`
    """
    # 统一小写，保证 agent id、英文命令和 task keywords 比较时大小写无关。
    lowered = (text or "").strip().lower()
    if not lowered:
        return []

    # 英文 token 适合匹配 agent id、命令词和常见英文任务描述。
    ascii_tokens = [token for token in _ASCII_TOKEN_RE.findall(lowered) if len(token) > 1]
    # 中文这里按单字匹配，虽然粗糙，但比整句更利于无分词依赖的第一版实现。
    cjk_tokens = _CJK_CHAR_RE.findall(lowered)

    # 用 `dict.fromkeys` 去重并保持原始顺序，便于后续测试断言更稳定。
    return list(dict.fromkeys([*ascii_tokens, *cjk_tokens]))


def similarity_score(query_tokens: list[str], candidate_tokens: list[str]) -> float:
    """按 token 重叠度计算相似度。

    Demo 输出：
        `0.67`
    """
    # 任一侧为空都说明没有稳定的匹配依据，直接给 0。
    if not query_tokens or not candidate_tokens:
        return 0.0

    # 这里故意不做复杂权重，保持算法透明、可预测、可测试。
    query_set = set(query_tokens)
    candidate_set = set(candidate_tokens)
    overlap = len(query_set & candidate_set)
    if overlap <= 0:
        return 0.0

    # 使用 `max(len(query), len(candidate))` 作为分母，让长任务模板不会被短查询轻易误命中。
    return overlap / max(len(query_set), len(candidate_set))


def clip_confidence(value: float) -> float:
    """把置信度裁剪到 `[0.0, 1.0]`。

    Demo 输出：
        `0.8`
    """
    # 所有 confidence 更新都收口到这里，避免散落的边界处理不一致。
    return max(0.0, min(1.0, round(value, 4)))


class ProcedureMemory:
    """管理 learned procedure 的持久化和匹配。

    公开方法都带了 Demo 输出说明，便于用户直接对照磁盘结果和测试脚本理解行为。
    """

    def __init__(
        self,
        workspace: Path,
        *,
        min_confidence: float = 0.55,
        match_threshold: float = 0.2,
    ) -> None:
        """初始化 procedure memory。

        Demo 输出：
            `ProcedureMemory(workspace=/tmp/demo-workspace, procedures.json ready)`
        """
        # `procedures.json` 用数组存储，人工排查时最直观。
        self.workspace = workspace
        self.path = _memory_root(workspace) / "procedures.json"
        # 低于该值的 procedure 即使匹配到关键词，也不建议作为复用提示。
        self.min_confidence = min_confidence
        # 匹配阈值保持较低，只作为 AutoSwarmBuilder / planner 的参考提示。
        self.match_threshold = match_threshold

    def list_procedures(self) -> list[ProcedureRecord]:
        """读取全部 procedure 记录并按置信度排序。

        Demo 输出：
            `[ProcedureRecord(...), ProcedureRecord(...)]`
        """
        # 文件损坏或不存在时直接回空列表，主流程会自动退回探索模式。
        raw = _load_json(self.path, [])
        records = [
            ProcedureRecord.from_dict(item)
            for item in raw
            if isinstance(item, dict)
        ]
        # 高置信度、最近更新的记录更靠前，方便测试和人工查看。
        records.sort(key=lambda item: (item.confidence, item.updated_at), reverse=True)
        return records

    def match_procedure(self, task: str) -> ProcedureRecord | None:
        """为当前任务匹配最合适的 procedure。

        Demo 输出：
            `ProcedureRecord(id='procedure-a1b2c3d4', task_template='生成周报', ...)`
        """
        # 没有 token 说明任务文本几乎为空，此时不应命中任何 procedure。
        query_tokens = task_tokens(task)
        if not query_tokens:
            return None

        best_record: ProcedureRecord | None = None
        best_score = 0.0
        for record in self.list_procedures():
            # 明显是占位/空结果的历史 procedure 直接忽略，避免污染后续路由。
            if contains_placeholder_summary(record.summary):
                continue
            # 优先用关键词匹配；任务模板是人工兜底线索。
            candidate_tokens = record.task_keywords or task_tokens(record.task_template)
            score = similarity_score(query_tokens, candidate_tokens)
            # task_template 全量包含时，给一个小额加分，提高近似重跑命中率。
            if record.task_template and record.task_template.lower() in task.lower():
                score += 0.1
            # 最终排序同时考虑相似度、置信度和失败率，避免高失败 procedure 反复被选中。
            weighted = score + record.confidence * 0.2 - record.failure_rate() * 0.2
            if weighted > best_score:
                best_record = record
                best_score = weighted

        # 分数不足则视为没有可靠命中，让上层走探索式执行。
        if best_record is None or best_score < self.match_threshold:
            return None
        return best_record

    async def record_candidate(self, task: str, result: BridgeResult) -> ProcedureRecord | None:
        """把探索阶段产出的候选 procedure 写入 memory。

        Demo 输出：
            `ProcedureRecord(id='procedure-a1b2c3d4', confidence=0.6, success_count=1, ...)`
        """
        # 只有 bridge 显式产出候选 procedure 时才会落盘。
        candidate = result.candidate_procedure
        if candidate is None:
            return None
        if not has_meaningful_summary(candidate.summary):
            return None

        # 记录写入时间统一在这里刷新，保证磁盘上的排序行为可预测。
        timestamp = now_iso()
        # 任务 token 统一在持久化层补齐，保证不依赖具体 bridge 的实现细节。
        merged_keywords = list(dict.fromkeys([*candidate.task_keywords, *task_tokens(task)]))
        candidate.task_keywords = merged_keywords
        candidate.task_template = candidate.task_template or task
        candidate.summary = candidate.summary or result.summary
        candidate.confidence = clip_confidence(candidate.confidence or 0.55)
        candidate.created_at = candidate.created_at or timestamp
        candidate.updated_at = timestamp

        records = self.list_procedures()
        best_index: int | None = None
        best_score = 0.0
        for index, record in enumerate(records):
            # 完全相同 agent 组合视为强相关；否则退回关键词重叠比对。
            same_agents = (
                record.strategy == candidate.strategy
                and record.agent_ids == candidate.agent_ids
            )
            score = 1.0 if same_agents else similarity_score(candidate.task_keywords, record.task_keywords)
            if score > best_score:
                best_index = index
                best_score = score

        if best_index is not None and best_score >= 0.5:
            # 合并已有记录，避免每次探索都生成一条几乎重复的 procedure。
            current = records[best_index]
            current.task_template = candidate.task_template or current.task_template
            current.summary = candidate.summary or current.summary
            current.agent_ids = list(candidate.agent_ids) or current.agent_ids
            current.strategy = candidate.strategy or current.strategy
            current.task_keywords = list(dict.fromkeys([*current.task_keywords, *candidate.task_keywords]))
            current.confidence = clip_confidence(max(current.confidence, candidate.confidence))
            current.success_count += 1
            current.updated_at = timestamp
            current.metadata.update(candidate.metadata)
            current.source_run_id = candidate.source_run_id or current.source_run_id
            stored = current
        else:
            # 新候选第一次入库时直接记为一次成功学习。
            candidate.success_count = max(candidate.success_count, 1)
            candidate.failure_count = max(candidate.failure_count, 0)
            candidate.created_at = candidate.created_at or timestamp
            candidate.updated_at = timestamp
            records.append(candidate)
            stored = candidate

        _atomic_write_json(self.path, [item.to_dict() for item in records])
        return stored

    async def update_confidence(self, procedure_id: str, delta: float) -> ProcedureRecord | None:
        """更新某条 procedure 的置信度与成败计数。

        Demo 输出：
            `ProcedureRecord(id='procedure-a1b2c3d4', confidence=0.75, success_count=2, failure_count=0, ...)`
        """
        # 没有主键时直接回空，避免误更新所有记录。
        if not procedure_id:
            return None

        records = self.list_procedures()
        updated: ProcedureRecord | None = None
        for record in records:
            if record.id != procedure_id:
                continue
            # 所有状态变更都集中在这里，保证计数和 confidence 始终同步。
            record.confidence = clip_confidence(record.confidence + delta)
            # 统一刷新“最近一次使用”和“最近一次更新时间”，这两个字段都服务于路由与排障。
            timestamp = now_iso()
            record.updated_at = timestamp
            record.last_used_at = timestamp
            if delta >= 0:
                record.success_count += 1
            else:
                record.failure_count += 1
            updated = record
            break

        if updated is None:
            return None

        _atomic_write_json(self.path, [item.to_dict() for item in records])
        return updated


class RunMemory:
    """管理 run 级别的历史记录。"""

    def __init__(self, workspace: Path, *, max_records: int = 200) -> None:
        """初始化 run memory。

        Demo 输出：
            `RunMemory(workspace=/tmp/demo-workspace, runs.json ready)`
        """
        # `runs.json` 保持轻量滚动窗口，避免长期运行后无限膨胀。
        self.workspace = workspace
        self.path = _memory_root(workspace) / "runs.json"
        self.max_records = max(1, max_records)

    def list_runs(self) -> list[RunRecord]:
        """读取全部 run 记录。

        Demo 输出：
            `[RunRecord(...), RunRecord(...)]`
        """
        raw = _load_json(self.path, [])
        return [
            RunRecord.from_dict(item)
            for item in raw
            if isinstance(item, dict)
        ]

    async def record_run(
        self,
        task: str,
        mode: ExecutionMode,
        result: BridgeResult,
        procedure_id: str | None = None,
    ) -> RunRecord:
        """把一次 agent team 运行结果落盘。

        Demo 输出：
            `RunRecord(id='run-1a2b3c4d', mode=<ExecutionMode.SWARMS: 'swarms'>, success=True, ...)`
        """
        # 把 attempt/原始 bridge 结果也带进 metadata，后面排查 swarms 执行很有用。
        record = RunRecord(
            task=task,
            mode=mode,
            success=result.success,
            summary=result.summary,
            error=result.error,
            procedure_id=procedure_id or (result.matched_procedure.id if result.matched_procedure else None),
            metadata={
                "attempts": [attempt.to_dict() for attempt in result.attempts],
                "bridge_result": result.to_dict(),
            },
        )
        runs = self.list_runs()
        runs.append(record)
        # 只保留最近 N 条，保证 JSON 文件体积可控。
        if len(runs) > self.max_records:
            runs = runs[-self.max_records:]
        _atomic_write_json(self.path, [item.to_dict() for item in runs])
        return record