from __future__ import annotations

from collections import defaultdict
from datetime import datetime, timezone
from typing import Any

from niannian.memory.types import ReflectResult
from niannian.personas import BRAINSTEM, HIPPOCAMPUS
from niannian.schemas.models import Atom


def _atom_kind(atom: Atom) -> str:
    return str(atom.meta.get("kind") or "note")


def _atom_thread(atom: Atom) -> str | None:
    rel = atom.meta.get("relations") or {}
    if isinstance(rel, dict) and rel.get("thread_id"):
        return str(rel["thread_id"])
    scope = atom.meta.get("scope") or {}
    if scope.get("loop_id") and scope.get("node_id"):
        return f"{scope['loop_id']}:{scope['node_id']}"
    return None


class ReflectService:
    """② 反思 Reflect — thread 摘要 + dream。"""

    def __init__(self, orchestrator, left, hippocampus, config: dict[str, Any]) -> None:
        self._orchestrator = orchestrator
        self._left = left
        self._hippocampus = hippocampus
        self._config = config
        self._min_frags = int(config.get("reflect", {}).get("thread_summary_min_fragments", 3))

    def reflect(self, *, scope: dict[str, Any] | None = None, project: str | None = None) -> ReflectResult:
        trace: list[str] = [f"[{BRAINSTEM}] reflect 开始"]
        req_scope = dict(scope or {})
        proj = project or req_scope.get("tool") or "global"
        loop_id = req_scope.get("loop_id")

        atoms = self._left.list_atoms(project=proj, limit=2000)
        if loop_id:
            atoms = [a for a in atoms if (a.meta.get("scope") or {}).get("loop_id") == loop_id]

        by_thread: dict[str, list[Atom]] = defaultdict(list)
        for atom in atoms:
            tid = _atom_thread(atom)
            if tid:
                by_thread[tid].append(atom)

        summaries = self._hippocampus.store.read_summaries()
        lifecycle = self._hippocampus.store.read_lifecycle()
        now = datetime.now(timezone.utc).isoformat()
        merged = 0

        for tid, group in by_thread.items():
            if len(group) < self._min_frags:
                continue
            group.sort(key=lambda a: a.created_at)
            parts: list[str] = []
            for a in group:
                parts.append(f"{_atom_kind(a)}: {a.raw_text[:80]}")
            narrative = " → ".join(parts[-5:])
            summaries[f"thread:{tid}"] = {
                "text": narrative,
                "updated_at": now,
                "status": "distilled",
                "atom_ids": [a.id for a in group],
            }
            merged += 1
            for a in group:
                if _atom_kind(a) == "write" and _atom_kind(group[-1]) == "adopt":
                    lifecycle[a.id] = {"status": "distilled", "updated_at": now}

        if merged:
            self._hippocampus.store.write_summaries(summaries)
            self._hippocampus.store.write_lifecycle(lifecycle)
            trace.append(f"[{HIPPOCAMPUS}] thread 摘要 {merged} 条")

        dream_stats = self._orchestrator.dream(project=proj)
        trace.extend(dream_stats.get("trace", []))

        return ReflectResult(
            threads_merged=merged,
            summaries_written=merged,
            atoms_touched=len(atoms),
            trace=trace,
        )
