"""Thread / channel context condenser.

Compresses historical channel messages into structured JSON summaries,
keeping token budgets under control for downstream agent prompts.

Ported from BOSS smack_context_condenser -- adapted for Cohort's
zero-infrastructure flat-file design.  Works with any OpenAI-compatible
LLM endpoint (Ollama, vLLM, LiteLLM, etc.) or falls back to a cheap
extractive summariser that needs no LLM at all.

Usage::

    from cohort.pipeline.condenser import ChannelCondenser

    condenser = ChannelCondenser()                       # extractive fallback
    condenser = ChannelCondenser(model="qwen2.5:1.5b")   # LLM-powered

    summary = condenser.condense(messages, agent_id="backend_dev")
    # => CondensedContext(message_count=42, participants=[...], ...)

Design principles:
    - Graceful degradation: LLM unavailable -> extractive fallback -> raw passthrough
    - Non-destructive: original messages are never mutated or deleted
    - Cache-friendly: LRU keyed on (agent_id, channel_id, last_message_id)
    - Token-budgeted: output capped at ~800 tokens regardless of input size
"""

from __future__ import annotations

import hashlib
import json
import logging
import re
from collections import OrderedDict
from dataclasses import asdict, dataclass, field
from typing import Any, Dict, List, Optional, Sequence

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Output schema
# ---------------------------------------------------------------------------

@dataclass
class CondensedContext:
    """Structured summary of a channel conversation."""

    message_count: int = 0
    participants: List[str] = field(default_factory=list)
    thread_summary: str = ""
    key_decisions: List[str] = field(default_factory=list)
    open_questions: List[str] = field(default_factory=list)
    last_speaker: str = ""
    last_topic: str = ""
    mentions_of_me: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    def to_prompt_block(self) -> str:
        """Render as a compact prompt section (~200-400 tokens)."""
        lines = [f"## Channel Context ({self.message_count} messages)"]
        if self.thread_summary:
            lines.append(self.thread_summary)
        if self.key_decisions:
            lines.append("**Decisions:** " + "; ".join(self.key_decisions))
        if self.open_questions:
            lines.append("**Open:** " + "; ".join(self.open_questions))
        if self.mentions_of_me:
            lines.append("**Mentions of you:** " + "; ".join(self.mentions_of_me))
        return "\n".join(lines)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> CondensedContext:
        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})


# ---------------------------------------------------------------------------
# LRU cache (simple, thread-safe-ish via OrderedDict)
# ---------------------------------------------------------------------------

class _LRUCache:
    """Bounded LRU dict.  Not perfectly thread-safe but good enough for
    a single-process MCP tool server."""

    def __init__(self, max_size: int = 500) -> None:
        self._store: OrderedDict[str, CondensedContext] = OrderedDict()
        self._max = max_size

    def get(self, key: str) -> Optional[CondensedContext]:
        if key in self._store:
            self._store.move_to_end(key)
            return self._store[key]
        return None

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

    def invalidate(self, prefix: str) -> int:
        """Drop all entries whose key starts with *prefix*."""
        to_drop = [k for k in self._store if k.startswith(prefix)]
        for k in to_drop:
            del self._store[k]
        return len(to_drop)


# ---------------------------------------------------------------------------
# ChannelCondenser
# ---------------------------------------------------------------------------

_SYSTEM_PROMPT = """\
You are a context condenser.  Given a sequence of channel messages,
produce a JSON object with EXACTLY these fields:

{
  "message_count": <int>,
  "participants": [<unique sender IDs>],
  "thread_summary": "<2-3 sentence recap of the conversation>",
  "key_decisions": ["<decision 1>", ...],
  "open_questions": ["<unresolved question>", ...],
  "last_speaker": "<sender of last message>",
  "last_topic": "<1-sentence topic of last message>",
  "mentions_of_me": ["<@mention with surrounding context>", ...]
}

Rules:
- Output ONLY valid JSON, no markdown fences, no commentary.
- thread_summary must be factual and concise.
- key_decisions: only things explicitly agreed/decided.
- open_questions: things asked but not yet answered.
- mentions_of_me: messages containing @{agent_id} (omit if none).
- If the conversation is trivial, keep fields short or empty lists.
"""


class ChannelCondenser:
    """Compress channel history into structured context.

    Parameters
    ----------
    model:
        Ollama model tag (e.g. ``"qwen2.5:1.5b"``).  If ``None``, the
        extractive fallback is used (no LLM required).
    ollama_url:
        Base URL for the Ollama API.  Defaults to ``http://localhost:11434``.
    timeout:
        Seconds before LLM call is abandoned and extractive fallback kicks in.
    cache_size:
        Max entries in the in-memory LRU cache.
    """

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

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def condense(
        self,
        messages: Sequence[Dict[str, Any]],
        agent_id: str = "",
        channel_id: str = "",
    ) -> CondensedContext:
        """Return a condensed summary of *messages*.

        Parameters
        ----------
        messages:
            List of message dicts.  Each must have at least ``sender``
            and ``content`` keys.  ``id`` and ``timestamp`` are optional.
        agent_id:
            The agent this context is being built for (drives @mention
            extraction).
        channel_id:
            Used for cache keying.

        Returns
        -------
        CondensedContext
            Always returns a value -- never raises.  On failure the result
            is an extractive summary built without an LLM.
        """
        if not messages:
            return CondensedContext()

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

        # Try LLM, fall back to extractive
        result: Optional[CondensedContext] = None
        if self.model:
            result = self._condense_llm(messages, agent_id)

        if result is None:
            result = self._condense_extractive(messages, agent_id)

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

    def invalidate_cache(self, channel_id: str = "", agent_id: str = "") -> int:
        """Drop cached entries for a channel/agent combo."""
        prefix = f"{agent_id}:{channel_id}" if agent_id else channel_id
        return self._cache.invalidate(prefix)

    # ------------------------------------------------------------------
    # LLM-powered condenser
    # ------------------------------------------------------------------

    def _condense_llm(
        self,
        messages: Sequence[Dict[str, Any]],
        agent_id: str,
    ) -> Optional[CondensedContext]:
        """Call Ollama chat endpoint.  Returns None on any failure."""
        try:
            import urllib.request

            transcript = self._render_transcript(messages)
            system = _SYSTEM_PROMPT.replace("{agent_id}", agent_id)
            payload = {
                "model": self.model,
                "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": transcript},
                ],
                "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())

            raw = body.get("message", {}).get("content", "")
            return self._parse_llm_output(raw, messages, agent_id)

        except Exception as exc:
            logger.debug("[condenser] LLM call failed, falling back: %s", exc)
            return None

    def _parse_llm_output(
        self,
        raw: str,
        messages: Sequence[Dict[str, Any]],
        agent_id: str,
    ) -> Optional[CondensedContext]:
        """Parse JSON from LLM output, tolerating markdown fences."""
        # Strip markdown code fences if present
        cleaned = re.sub(r"^```(?:json)?\s*", "", raw.strip())
        cleaned = re.sub(r"\s*```$", "", cleaned)
        try:
            data = json.loads(cleaned)
            if not isinstance(data, dict):
                return None
            return CondensedContext.from_dict(data)
        except (json.JSONDecodeError, TypeError):
            logger.debug("[condenser] Failed to parse LLM JSON output")
            return None

    # ------------------------------------------------------------------
    # Extractive fallback (no LLM needed)
    # ------------------------------------------------------------------

    def _condense_extractive(
        self,
        messages: Sequence[Dict[str, Any]],
        agent_id: str,
    ) -> CondensedContext:
        """Build a summary using simple heuristics -- no LLM needed."""
        participants: List[str] = []
        seen_senders: set[str] = set()
        mentions: List[str] = []
        questions: List[str] = []

        for msg in messages:
            sender = msg.get("sender", "unknown")
            if sender not in seen_senders:
                seen_senders.add(sender)
                participants.append(sender)

            content = msg.get("content", "")

            # Collect @mentions of our agent
            if agent_id and f"@{agent_id}" in content:
                snippet = content[:200]
                mentions.append(f"{sender}: {snippet}")

            # Heuristic: lines ending with ? are likely questions
            for line in content.split("\n"):
                stripped = line.strip()
                if stripped.endswith("?") and len(stripped) > 10:
                    questions.append(stripped)

        last_msg = messages[-1]
        last_sender = last_msg.get("sender", "")
        last_content = last_msg.get("content", "")

        # Build a rough summary from the first and last few messages
        summary_parts: List[str] = []
        if len(messages) <= 3:
            for m in messages:
                summary_parts.append(f"{m.get('sender', '?')}: {m.get('content', '')[:120]}")
        else:
            first = messages[0]
            summary_parts.append(
                f"Started with {first.get('sender', '?')}: "
                f"{first.get('content', '')[:120]}"
            )
            summary_parts.append(
                f"Most recently {last_sender}: {last_content[:120]}"
            )
            summary_parts.append(f"{len(messages)} messages from {len(participants)} participants.")

        return CondensedContext(
            message_count=len(messages),
            participants=participants,
            thread_summary=" ".join(summary_parts),
            key_decisions=[],  # extractive can't reliably detect decisions
            open_questions=questions[:5],
            last_speaker=last_sender,
            last_topic=last_content[:150],
            mentions_of_me=mentions[:5],
        )

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

    def _render_transcript(self, messages: Sequence[Dict[str, Any]]) -> str:
        """Render messages into a compact transcript string."""
        lines: List[str] = []
        for msg in messages:
            sender = msg.get("sender", "?")
            content = msg.get("content", "")
            ts = msg.get("timestamp", "")
            if ts:
                lines.append(f"[{ts}] {sender}: {content}")
            else:
                lines.append(f"{sender}: {content}")
        return "\n".join(lines)

    def _cache_key(
        self,
        messages: Sequence[Dict[str, Any]],
        agent_id: str,
        channel_id: str,
    ) -> str:
        """Deterministic cache key from agent + channel + last message."""
        last_id = ""
        if messages:
            last = messages[-1]
            last_id = last.get("id", "") or last.get("timestamp", "") or str(len(messages))
        raw = f"{agent_id}:{channel_id}:{last_id}:{len(messages)}"
        return hashlib.md5(raw.encode()).hexdigest()
