"""MemoryConsolidator — phase 2 merge of similar memories during idle.

Phase 1 (ingest_knowledge_fast) adds all entries without conflict resolution,
recording similarity pairs in pending_merges.

Phase 2 (this module) runs when the digestor is idle:
- Read pending_merges → build connected components via UnionFind
- For each group: fetch records → LLM decides merge or keep separate
- Merge: delete old memories, insert single merged memory
"""

import json
import logging
import os
import time
from typing import TYPE_CHECKING

from neo_cortex import config

if TYPE_CHECKING:
    from neo_cortex.dedup import DedupStore
    from neo_cortex.embedder import JinaEmbedder
    from neo_cortex.graph import ConceptGraph
    from neo_cortex.llm import LLMClient
    from neo_cortex.memory_index import MemoryIndex
    from neo_cortex.models import StructuredFields
    from neo_cortex.store import MemoryStore

logger = logging.getLogger("neo-cortex-consolidator")

from functools import partial
from neo_cortex._log import boot_log
_bl = partial(boot_log, "CONSOLIDATOR")


MERGE_PROMPT = """You are a memory consolidation system. Given a group of similar memories, decide whether to MERGE them into one or KEEP them separate.

RULES:
- MERGE if they describe the same topic/fact and can be combined without losing information
- KEEP_SEPARATE if they describe genuinely different things that happen to share terminology
- When merging, produce a single text that preserves ALL important details from every memory

Return ONLY valid JSON:
{
  "action": "merge" or "keep_separate",
  "merged_text": "the combined text (only if action=merge)",
  "merged_title": "concise title (only if action=merge)",
  "reason": "brief explanation"
}

Memories to evaluate:
"""


class _UnionFind:
    """Disjoint set for building connected components from pairs."""

    def __init__(self) -> None:
        self._parent: dict[str, str] = {}

    def find(self, x: str) -> str:
        if x not in self._parent:
            self._parent[x] = x
        while self._parent[x] != x:
            self._parent[x] = self._parent[self._parent[x]]
            x = self._parent[x]
        return x

    def union(self, a: str, b: str) -> None:
        ra, rb = self.find(a), self.find(b)
        if ra != rb:
            self._parent[ra] = rb

    def groups(self) -> dict[str, list[str]]:
        """Return root → [members] mapping."""
        result: dict[str, list[str]] = {}
        for x in self._parent:
            root = self.find(x)
            result.setdefault(root, []).append(x)
        return result


class MemoryConsolidator:
    """Merges similar memories recorded in pending_merges."""

    def __init__(
        self,
        store: 'MemoryStore',
        index: 'MemoryIndex',
        embedder: 'JinaEmbedder',
        llm: 'LLMClient',
        dedup: 'DedupStore | None' = None,
        graph: 'ConceptGraph | None' = None,
    ):
        self._store = store
        self._index = index
        self._embedder = embedder
        self._llm = llm
        self._dedup = dedup
        self._graph = graph

    async def consolidate(self) -> int:
        """Process all pending merges. Returns number of groups merged."""
        pairs = self._index.get_pending_merges()
        if not pairs:
            return 0

        _bl(f"consolidate: {len(pairs)} pending pairs")

        groups = self._build_groups(pairs)
        row_ids = [p[3] for p in pairs]

        merged_count = 0
        for group_ids in groups:
            if len(group_ids) < 2:
                continue
            # Cap group size
            group_ids = group_ids[:config.CONSOLIDATION_MAX_GROUP_SIZE]
            try:
                did_merge = await self._consolidate_group(group_ids)
                if did_merge:
                    merged_count += 1
            except Exception as e:
                _bl(f"consolidate: group {group_ids[:2]}... FAILED: {e}")

        self._index.mark_merges_processed(row_ids)
        _bl(f"consolidate: done — {merged_count} groups merged, {len(row_ids)} pairs processed")
        return merged_count

    def _build_groups(self, pairs: list[tuple[str, str, float, int]]) -> list[list[str]]:
        """Build connected components from similarity pairs."""
        uf = _UnionFind()
        for id_a, id_b, sim, _ in pairs:
            uf.union(id_a, id_b)
        return list(uf.groups().values())

    async def _consolidate_group(self, memory_ids: list[str]) -> bool:
        """Fetch records, ask LLM to merge or keep. Returns True if merged."""
        # Fetch all records + structured fields BEFORE any deletion
        records = self._index.get_by_ids(memory_ids)
        if len(records) < 2:
            _bl(f"consolidate_group: <2 records found for {memory_ids}, skipping")
            return False

        structured_map: dict[str, 'StructuredFields | None'] = {}
        for mid in memory_ids:
            structured_map[mid] = self._index.get_structured(mid)

        # Build prompt
        memory_texts = []
        for i, rec in enumerate(records):
            sf = structured_map.get(rec.id)
            title = sf.title if sf else rec.topic
            memory_texts.append(f"[Memory {i+1}] id={rec.id}\nTitle: {title}\nContent: {rec.document}")

        combined = "\n\n".join(memory_texts)

        # Ask LLM
        try:
            response = await self._llm.call(MERGE_PROMPT, combined, max_tokens=1000, temperature=0.1)
            response = _strip_code_fences(response)
            decision = json.loads(response)
        except (json.JSONDecodeError, Exception) as e:
            _bl(f"consolidate_group: LLM parse error: {e}")
            return False

        action = decision.get("action", "keep_separate")

        if action != "merge":
            _bl(f"consolidate_group: keep_separate ({len(records)} memories)")
            return False

        merged_text = decision.get("merged_text", "")
        merged_title = decision.get("merged_title", "")
        if not merged_text:
            _bl("consolidate_group: merge decision but empty merged_text, skipping")
            return False

        _bl(f"consolidate_group: MERGING {len(records)} memories → '{merged_title}'")

        # Collect all concepts from the group before deleting
        all_concepts: list[str] = []
        for mid in memory_ids:
            sf = structured_map.get(mid)
            if sf and sf.concepts:
                for c in sf.concepts:
                    if c not in all_concepts:
                        all_concepts.append(c)

        # Delete old memories from all stores
        for rec in records:
            sf = structured_map.get(rec.id)
            if self._graph and sf and sf.concepts:
                self._graph.remove_memory(rec.id, sf.concepts)
            self._store.delete(rec.id)
            self._index.delete(rec.id)

        # Insert merged memory
        ts = time.time()
        merged_id = f"cn_{int(ts) % 100000}"

        from neo_cortex.models import Activity, MemoryRecord, StructuredFields as SF

        # Use first record's metadata as base
        base = records[0]
        merged_record = MemoryRecord(
            id=merged_id,
            session_id=base.session_id,
            timestamp=ts,
            question=merged_title,
            answer_preview=merged_text[:500],
            document=merged_text,
            project=base.project,
            topic=merged_title[:100],
            activity=base.activity,
            energy=1.0,
            source="consolidator",
        )

        embedding = await self._embedder.embed_passage(merged_text)
        self._store.insert(merged_record, embedding)

        merged_structured = SF(
            title=merged_title,
            summary=merged_text[:300],
            facts=[merged_text],
            concepts=all_concepts,
        )
        self._index.insert(merged_record, merged_structured)

        if self._dedup:
            self._dedup.add(merged_text, merged_id)

        if self._graph and all_concepts:
            self._graph.add_memory(merged_id, all_concepts)

        _bl(f"consolidate_group: DONE → {merged_id}")
        return True


def _strip_code_fences(text: str) -> str:
    """Strip markdown code fences from LLM output."""
    text = text.strip()
    if text.startswith("```"):
        text = "\n".join(text.split("\n")[1:])
    if text.endswith("```"):
        text = "\n".join(text.split("\n")[:-1])
    return text.strip()
