"""
Memory Compression Worker — Phase 23.5

Background worker that automatically clusters raw memory nodes and
compresses them into summary nodes in the DAG.

Pipeline:
  1. Scan recent uncompressed raw nodes
  2. Cluster by semantic similarity (keyword Jaccard)
  3. Summarize each cluster (LLM or rule-based)
  4. Write summary node to DAG
  5. Link summary → children (preserve raw nodes)

Scheduling:
  - Every N observations (triggered by MemoryEngine._auto_compress)
  - Or every 5 minutes via CompressionWorkerThread (background mode)
"""

import logging
import threading
import time
from typing import Dict, List, Optional

logger = logging.getLogger("openmemo.compression")

DEFAULT_CLUSTER_THRESHOLD = 0.25
DEFAULT_MIN_RAW = 3
DEFAULT_MAX_CLUSTER_SIZE = 10


class CompressionWorker:
    """
    Synchronous compression worker.

    Args:
        graph           — MemoryGraph instance
        store           — SQLiteStore (for reading raw observations)
        llm_fn          — Optional LLM callable: llm_fn(prompt: str) -> str
        cluster_threshold — Jaccard similarity threshold for clustering
    """

    def __init__(self, graph, store=None,
                 llm_fn=None,
                 cluster_threshold: float = DEFAULT_CLUSTER_THRESHOLD):
        self._graph = graph
        self._store = store
        self._llm_fn = llm_fn
        self._threshold = cluster_threshold

    def run_once(self, agent_id: str = "", scene: str = "",
                 min_raw: int = DEFAULT_MIN_RAW) -> Dict:
        """
        Run one compression pass.

        Returns:
            {
              "raw_processed": int,
              "clusters_found": int,
              "summaries_created": int,
              "skipped": int,     # clusters too small
            }
        """
        uncompressed = [
            n for n in self._graph.list_raw(
                agent_id=agent_id, scene=scene, limit=500
            )
            if not n.parent_ids
        ]

        if len(uncompressed) < min_raw:
            return {
                "raw_processed": 0,
                "clusters_found": 0,
                "summaries_created": 0,
                "skipped": 0,
            }

        clusters = self._cluster(uncompressed)

        summaries_created = 0
        skipped = 0

        for cluster in clusters:
            if len(cluster) < min_raw:
                skipped += 1
                continue

            summary_text = self._summarize(cluster)
            if not summary_text:
                skipped += 1
                continue

            child_ids = [n.id for n in cluster]
            node_agent = cluster[0].agent_id if cluster else agent_id
            node_scene = cluster[0].scene if cluster else scene
            node_task = cluster[0].task_id if cluster else ""

            avg_importance = sum(n.importance for n in cluster) / len(cluster)

            self._graph.add_summary(
                content=summary_text,
                child_ids=child_ids,
                agent_id=node_agent,
                scene=node_scene,
                task_id=node_task,
                importance=min(1.0, avg_importance * 1.1),
            )
            summaries_created += 1

        logger.info(
            "[compression] pass done: raw=%d clusters=%d summaries=%d skipped=%d",
            len(uncompressed), len(clusters), summaries_created, skipped
        )

        return {
            "raw_processed": len(uncompressed),
            "clusters_found": len(clusters),
            "summaries_created": summaries_created,
            "skipped": skipped,
        }

    def _cluster(self, nodes: list) -> List[list]:
        """
        Group nodes by semantic similarity using keyword Jaccard.
        Returns list of clusters (each cluster is a list of nodes).
        """
        if not nodes:
            return []

        def tokenize(text: str) -> set:
            stopwords = {"the", "a", "an", "is", "was", "to", "of",
                         "and", "in", "for", "with", "at"}
            return {w for w in text.lower().split() if w not in stopwords}

        def jaccard(a: set, b: set) -> float:
            if not a and not b:
                return 1.0
            union = a | b
            if not union:
                return 0.0
            return len(a & b) / len(union)

        token_sets = [tokenize(n.content) for n in nodes]
        assigned = [False] * len(nodes)
        clusters = []

        for i, node in enumerate(nodes):
            if assigned[i]:
                continue
            cluster = [node]
            assigned[i] = True
            for j in range(i + 1, len(nodes)):
                if assigned[j]:
                    continue
                if len(cluster) >= DEFAULT_MAX_CLUSTER_SIZE:
                    break
                sim = jaccard(token_sets[i], token_sets[j])
                if sim >= self._threshold:
                    cluster.append(nodes[j])
                    assigned[j] = True
            clusters.append(cluster)

        return clusters

    def _summarize(self, cluster: list) -> str:
        """
        Generate summary text for a cluster of nodes.
        Uses LLM if available, otherwise rule-based.
        """
        texts = [n.content for n in cluster]
        if self._llm_fn:
            return self._summarize_with_llm(texts)
        return self._summarize_rule_based(texts, cluster)

    def _summarize_with_llm(self, texts: List[str]) -> str:
        prompt = (
            "Summarize the following memory observations into a single concise "
            "sentence (max 2 sentences). Focus on key facts and outcomes.\n\n"
            "Observations:\n" + "\n".join(f"- {t}" for t in texts[:10])
        )
        try:
            return self._llm_fn(prompt).strip()
        except Exception as e:
            logger.warning("[compression] LLM summarization failed: %s", e)
            return self._summarize_rule_based(texts, [])

    def _summarize_rule_based(self, texts: List[str], cluster: list) -> str:
        if not texts:
            return ""

        types = set()
        if cluster:
            for n in cluster:
                meta = getattr(n, "metadata", {})
                if meta and meta.get("obs_type"):
                    types.add(meta["obs_type"])

        count = len(texts)
        type_str = "/".join(sorted(types)) if types else "memory"

        if count == 1:
            return texts[0]

        words_all = []
        for t in texts:
            words_all.extend(t.lower().split())
        from collections import Counter
        common = [w for w, _ in Counter(words_all).most_common(5)
                  if len(w) > 3]

        theme = ", ".join(common[:3]) if common else "related activity"
        return f"[Summary] {count} {type_str} observations about {theme}: {texts[0][:80]}..."


class CompressionWorkerThread:
    """
    Background thread that runs CompressionWorker periodically.

    Args:
        graph       — MemoryGraph
        store       — SQLiteStore
        interval    — seconds between compression passes (default 300 = 5 min)
        llm_fn      — optional LLM callable
    """

    def __init__(self, graph, store=None, interval: int = 300,
                 llm_fn=None):
        self._graph = graph
        self._store = store
        self._interval = interval
        self._llm_fn = llm_fn
        self._thread: Optional[threading.Thread] = None
        self._running = False

    def start(self):
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(
            target=self._loop,
            daemon=True,
            name="openmemo-compression",
        )
        self._thread.start()
        logger.info("[compression] background worker started (interval=%ds)",
                    self._interval)

    def stop(self):
        self._running = False
        if self._thread:
            self._thread.join(timeout=5)
        logger.info("[compression] background worker stopped")

    def _loop(self):
        while self._running:
            try:
                worker = CompressionWorker(
                    graph=self._graph,
                    store=self._store,
                    llm_fn=self._llm_fn,
                )
                worker.run_once(min_raw=DEFAULT_MIN_RAW)
            except Exception as e:
                logger.warning("[compression] worker error: %s", e)
            time.sleep(self._interval)
