"""Agent memory distiller -- condense persistent memory to query-relevant context.

Ported from BOSS tools/memory_distiller.  Agents accumulate facts, paths,
collaborators, and working memory over time.  Before injecting memory into
a prompt, this module distils the raw dump down to only what's relevant
for the current query.

Two-layer architecture:
    1. Raw assembly -- collect all memory sections from the agent store
    2. LLM distillation -- send raw + query to a cheap model for focused extraction

Falls back to keyword-based filtering if no LLM is available.

Usage::

    from cohort.pipeline.memory_distiller import MemoryDistiller

    distiller = MemoryDistiller(model="qwen2.5:1.5b")  # or None for keyword fallback

    # Raw memory dict (from agent_store or memory.json)
    memory = {
        "learned_facts": [...],
        "known_paths": [...],
        "active_tasks": [...],
        "collaborators": [...],
        "working_memory": [...],
        "preferences": {...},
    }

    result = distiller.distil(memory, query="How should we handle auth tokens?")
    # => "- Auth tokens stored in redis with 30m TTL\\n- Security agent flagged..."
"""

from __future__ import annotations

import hashlib
import json
import logging
import re
from collections import OrderedDict
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

MAX_BULLETS = 12
MAX_OUTPUT_CHARS = 2000        # ~500 tokens
CACHE_TTL_SECONDS = 120.0
CACHE_MAX_ENTRIES = 200
WORKING_MEMORY_LIMIT = 20
WORKING_MEMORY_CHAR_LIMIT = 200


# ---------------------------------------------------------------------------
# LRU cache with TTL
# ---------------------------------------------------------------------------

class _TTLCache:
    """LRU cache with per-entry TTL."""

    def __init__(self, max_size: int = CACHE_MAX_ENTRIES, ttl: float = CACHE_TTL_SECONDS):
        self._store: OrderedDict[str, tuple[float, str]] = OrderedDict()
        self._max = max_size
        self._ttl = ttl

    def get(self, key: str) -> Optional[str]:
        import time
        if key in self._store:
            ts, value = self._store[key]
            if (time.time() - ts) < self._ttl:
                self._store.move_to_end(key)
                return value
            else:
                del self._store[key]
        return None

    def put(self, key: str, value: str) -> None:
        import time
        self._store[key] = (time.time(), value)
        self._store.move_to_end(key)
        while len(self._store) > self._max:
            self._store.popitem(last=False)


# ---------------------------------------------------------------------------
# MemoryDistiller
# ---------------------------------------------------------------------------

class MemoryDistiller:
    """Distil agent memory to query-relevant context.

    Parameters
    ----------
    model:
        Ollama model tag for LLM distillation.  If ``None``, uses
        keyword-based filtering only.
    ollama_url:
        Base URL for Ollama API.
    timeout:
        Seconds before LLM call is abandoned.
    """

    def __init__(
        self,
        model: Optional[str] = None,
        ollama_url: str = "http://localhost:11434",
        timeout: float = 8.0,
    ) -> None:
        self.model = model
        self.ollama_url = ollama_url.rstrip("/")
        self.timeout = timeout
        self._cache = _TTLCache()

    def distil(
        self,
        memory: Dict[str, Any],
        query: str,
        agent_id: str = "",
    ) -> str:
        """Distil *memory* to content relevant to *query*.

        Returns
        -------
        str
            Bullet-point list of relevant memory, or ``"(no relevant memory)"``
            if nothing matches.
        """
        if not memory or not query:
            return "(no relevant memory)"

        # Cache lookup
        cache_key = self._cache_key(agent_id, query)
        cached = self._cache.get(cache_key)
        if cached is not None:
            return cached

        # Assemble raw memory text
        raw = self._assemble_raw(memory)
        if not raw.strip():
            return "(no relevant memory)"

        # Try LLM distillation, fall back to keyword filtering
        result: Optional[str] = None
        if self.model:
            result = self._distil_llm(raw, query)

        if result is None:
            result = self._distil_keyword(memory, query)

        self._cache.put(cache_key, result)
        return result

    # ------------------------------------------------------------------
    # Raw assembly
    # ------------------------------------------------------------------

    def _assemble_raw(self, memory: Dict[str, Any]) -> str:
        """Assemble all memory sections into a single text block."""
        sections: List[str] = []

        # Known paths
        paths = memory.get("known_paths", [])
        if paths:
            sections.append("## Known Paths\n" + "\n".join(f"- {p}" for p in paths[:20]))

        # Learned facts
        facts = memory.get("learned_facts", [])
        if facts:
            lines = []
            for f in facts[:30]:
                if isinstance(f, dict):
                    text = f.get("fact", f.get("text", str(f)))
                    conf = f.get("confidence", "")
                    src = f.get("source", "")
                    line = f"- {text}"
                    if conf:
                        line += f" (confidence: {conf})"
                    if src:
                        line += f" [source: {src}]"
                    lines.append(line)
                else:
                    lines.append(f"- {f}")
            sections.append("## Learned Facts\n" + "\n".join(lines))

        # Active tasks
        tasks = memory.get("active_tasks", [])
        if tasks:
            lines = []
            for t in tasks[:10]:
                if isinstance(t, dict):
                    desc = t.get("description", t.get("task", str(t)))
                    status = t.get("status", "")
                    lines.append(f"- {desc} ({status})" if status else f"- {desc}")
                else:
                    lines.append(f"- {t}")
            sections.append("## Active Tasks\n" + "\n".join(lines))

        # Collaborators
        collabs = memory.get("collaborators", [])
        if collabs:
            lines = []
            for c in collabs[:10]:
                if isinstance(c, dict):
                    name = c.get("name", c.get("agent_id", str(c)))
                    rel = c.get("relationship", "")
                    lines.append(f"- {name}: {rel}" if rel else f"- {name}")
                else:
                    lines.append(f"- {c}")
            sections.append("## Collaborators\n" + "\n".join(lines))

        # Preferences
        prefs = memory.get("preferences", {})
        if prefs and isinstance(prefs, dict):
            lines = [f"- {k}: {v}" for k, v in list(prefs.items())[:10]]
            sections.append("## Preferences\n" + "\n".join(lines))

        # Working memory (recent interactions)
        wm = memory.get("working_memory", [])
        if wm:
            lines = []
            for entry in wm[-WORKING_MEMORY_LIMIT:]:
                if isinstance(entry, dict):
                    text = entry.get("content", entry.get("text", str(entry)))
                else:
                    text = str(entry)
                lines.append(f"- {str(text)[:WORKING_MEMORY_CHAR_LIMIT]}")
            sections.append("## Recent Context\n" + "\n".join(lines))

        return "\n\n".join(sections)

    # ------------------------------------------------------------------
    # LLM distillation
    # ------------------------------------------------------------------

    _DISTIL_SYSTEM = """\
You are a memory distiller. Given an agent's full memory and a current query,
extract ONLY the facts, decisions, preferences, and context relevant to the query.

Rules:
- Output 5-12 bullet points (lines starting with "- ")
- Preserve exact names, file paths, config values
- Drop thinking noise, redundant entries, stale context
- Never repeat the same fact twice
- If nothing is relevant, output exactly: (no relevant memory)
"""

    def _distil_llm(self, raw_memory: str, query: str) -> Optional[str]:
        """Use LLM to extract relevant memory."""
        try:
            import urllib.request

            payload = {
                "model": self.model,
                "messages": [
                    {"role": "system", "content": self._DISTIL_SYSTEM},
                    {"role": "user", "content": (
                        f"## Query\n{query}\n\n"
                        f"## Agent Memory\n{raw_memory[:6000]}"
                    )},
                ],
                "stream": False,
                "options": {"num_predict": 512, "temperature": 0.1},
            }
            req = urllib.request.Request(
                f"{self.ollama_url}/api/chat",
                data=json.dumps(payload).encode(),
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                body = json.loads(resp.read())

            result = body.get("message", {}).get("content", "").strip()
            if result and len(result) > 5:
                return result[:MAX_OUTPUT_CHARS]

        except Exception as exc:
            logger.debug("[memory-distiller] LLM call failed: %s", exc)

        return None

    # ------------------------------------------------------------------
    # Keyword-based fallback
    # ------------------------------------------------------------------

    def _distil_keyword(self, memory: Dict[str, Any], query: str) -> str:
        """Extract memory entries matching query keywords."""
        # Extract keywords from query (words 3+ chars, lowered)
        keywords = set(
            w.lower() for w in re.findall(r"\w+", query)
            if len(w) >= 3
        )

        if not keywords:
            return "(no relevant memory)"

        matches: List[str] = []

        # Score each memory entry by keyword overlap
        for section_name, entries in memory.items():
            if section_name == "preferences":
                # Handle dict
                if isinstance(entries, dict):
                    for k, v in entries.items():
                        text = f"{k}: {v}"
                        if self._keyword_match(text, keywords):
                            matches.append(f"- [pref] {text}")
                continue

            if not isinstance(entries, list):
                continue

            for entry in entries:
                if isinstance(entry, dict):
                    text = " ".join(str(v) for v in entry.values())
                    display = entry.get("fact", entry.get("text",
                               entry.get("description", entry.get("content",
                               str(entry)))))
                else:
                    text = str(entry)
                    display = text

                if self._keyword_match(text, keywords):
                    bullet = str(display)[:WORKING_MEMORY_CHAR_LIMIT]
                    matches.append(f"- {bullet}")

                if len(matches) >= MAX_BULLETS:
                    break
            if len(matches) >= MAX_BULLETS:
                break

        if not matches:
            return "(no relevant memory)"

        return "\n".join(matches[:MAX_BULLETS])

    @staticmethod
    def _keyword_match(text: str, keywords: set) -> bool:
        """Check if *text* contains any of the keywords."""
        text_lower = text.lower()
        return any(kw in text_lower for kw in keywords)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _cache_key(self, agent_id: str, query: str) -> str:
        """Deterministic cache key from agent + query tail."""
        raw = f"{agent_id}:{query[-500:]}"
        return hashlib.md5(raw.encode()).hexdigest()
