"""
conversation_summarizer.py

Context-aware, LLM-powered rolling conversation summarizer.
Strategy : Progressive — auto-compresses as the conversation grows.
Preserved : User intent & goals, key decisions, action items.

Provider-agnostic: wire any LLM backend via LLMClient.complete().
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from enum import Enum


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------

class Role(str, Enum):
    SYSTEM    = "system"
    USER      = "user"
    ASSISTANT = "assistant"


@dataclass
class Message:
    role: Role
    content: str

    def to_dict(self) -> dict[str, str]:
        return {"role": self.role.value, "content": self.content}


@dataclass
class ConversationSummary:
    """
    Structured snapshot of a conversation produced by the summarizer.

    Fields
    ------
    overview      : 2-3 sentence narrative of what was discussed.
    user_intent   : The user's underlying goal (even if never stated explicitly).
    key_decisions : Conclusions or choices that were reached.
    action_items  : Pending tasks / next steps that still need to happen.
    """
    overview:      str
    user_intent:   str
    key_decisions: list[str] = field(default_factory=list)
    action_items:  list[str] = field(default_factory=list)

    def to_readable_text(self) -> str:
        """Plain-text summary suitable for displaying to a human."""
        lines = [
            f"Summary\n{'─' * 40}",
            f"{self.overview}",
            f"\nUser intent: {self.user_intent}",
        ]
        if self.key_decisions:
            lines.append("\nKey decisions:")
            for d in self.key_decisions:
                lines.append(f"  • {d}")
        if self.action_items:
            lines.append("\nAction items:")
            for a in self.action_items:
                lines.append(f"  • {a}")
        return "\n".join(lines)

    def to_context_block(self) -> str:
        """
        Compact text block to prepend as a system message in future LLM calls,
        so the model has full context without seeing every raw turn.
        """
        lines = [
            "=== CONVERSATION SUMMARY ===",
            f"Overview     : {self.overview}",
            f"User intent  : {self.user_intent}",
        ]
        if self.key_decisions:
            lines.append("Key decisions: " + " | ".join(self.key_decisions))
        if self.action_items:
            lines.append("Action items : " + " | ".join(self.action_items))
        lines.append("=== END SUMMARY ===")
        return "\n".join(lines)


# ---------------------------------------------------------------------------
# LLM client -- swap complete() body for any provider
# ---------------------------------------------------------------------------

class LLMClient:
    """
    Thin wrapper around any chat-completion SDK.

    Azure OpenAI example
    --------------------
    from openai import AzureOpenAI

    raw = AzureOpenAI(
        azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
        api_key=os.environ["AZURE_OPENAI_API_KEY"],
        api_version=os.environ.get("AZURE_OPENAI_API_VERSION", "2024-02-01"),
    )
    client = LLMClient(raw, model=os.environ["AZURE_OPENAI_DEPLOYMENT"])
    """

    def __init__(self, client, model: str):
        self._client = client
        self._model  = model

    def complete(self, messages: list[dict[str, str]]) -> str:
        response = self._client.chat.completions.create(
            model=self._model,
            messages=messages,
        )
        return response.choices[0].message.content


# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------

_FULL_SUMMARY_PROMPT = """\
You are an expert conversation analyst.

Produce a structured, context-aware summary of the conversation below.
Focus on what matters for continuity -- not a chronological replay.

Rules:
- Identify the user's underlying goal even if never stated explicitly.
- Capture every decision that was reached so far.
- List only action items that are still pending.
- Be concise but complete; this summary will replace the full history.

Respond ONLY with a valid JSON object -- no markdown fences, no extra text:
{
  "overview":      "<2-3 sentence narrative>",
  "user_intent":   "<one sentence>",
  "key_decisions": ["<decision>", ...],
  "action_items":  ["<item>", ...]
}"""

_INCREMENTAL_SUMMARY_PROMPT = """\
You are an expert conversation analyst performing an incremental summary update.

You will receive:
  1. An existing summary of earlier conversation turns.
  2. New turns that occurred after that summary.

Produce an updated summary that merges both coherently.
- Carry forward still-relevant decisions and action items.
- Mark items as resolved if the new turns show they were completed.
- Do NOT simply concatenate -- integrate and deduplicate.

Respond ONLY with a valid JSON object -- no markdown fences, no extra text:
{
  "overview":      "<2-3 sentence narrative>",
  "user_intent":   "<one sentence>",
  "key_decisions": ["<decision>", ...],
  "action_items":  ["<item>", ...]
}"""


# ---------------------------------------------------------------------------
# Core summarizer
# ---------------------------------------------------------------------------

class ConversationSummarizer:
    """
    Progressive, context-aware conversation summarizer.

    How it works
    ------------
    Feed turns in via add_message() / add_messages(). When the buffer reaches
    `auto_summarize_after` messages, the summarizer automatically compresses:

      * First compression  -> full summary of all buffered turns.
      * Later compressions -> incremental merge of the existing summary
                             with the newly buffered turns (context never lost).

    After compression, the last `keep_recent_turns` messages are kept verbatim
    so the model always sees the freshest raw context alongside the summary.

    Usage
    -----
        summarizer = ConversationSummarizer(llm_client)

        for turn in incoming_turns:
            summarizer.add_message(turn)

        # Inject into your next LLM call:
        messages = summarizer.get_context_messages()
        messages.append({"role": "user", "content": next_user_input})
        response = llm_client.complete(messages)
    """

    def __init__(
        self,
        llm_client: LLMClient,
        auto_summarize_after: int = 30,
        keep_recent_turns:    int = 6,
    ):
        """
        Parameters
        ----------
        llm_client            : Configured LLMClient instance.
        auto_summarize_after  : Compress when the buffer hits this many messages.
        keep_recent_turns     : Verbatim messages to retain after each compression.
        """
        self._llm          = llm_client
        self._threshold    = auto_summarize_after
        self._keep_recent  = keep_recent_turns
        self._buffer:  list[Message]               = []
        self._summary: ConversationSummary | None  = None

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

    def add_message(self, message: Message) -> None:
        """Append one turn; triggers compression if the buffer is full."""
        self._buffer.append(message)
        if len(self._buffer) >= self._threshold:
            self._compress()

    def add_messages(self, messages: list[Message]) -> None:
        """Append multiple turns; may trigger compression mid-batch."""
        for msg in messages:
            self.add_message(msg)

    def get_context_messages(self) -> list[dict[str, str]]:
        """
        Returns an optimised message list ready for your next LLM call:

            [system: summary block]   <- present only after first compression
            [recent verbatim turns]   <- last keep_recent_turns messages

        Pass the return value directly as the `messages` parameter, then
        append your new user turn before calling the LLM.
        """
        result: list[dict[str, str]] = []

        if self._summary:
            result.append({
                "role":    "system",
                "content": self._summary.to_context_block(),
            })

        recent = self._buffer[-self._keep_recent:] if self._buffer else []
        result.extend(m.to_dict() for m in recent)
        return result

    def summarize_to_text(self, conversation_text: str) -> str:
        """
        Simplest interface: pass the full conversation as plain text, get a
        readable summary back as plain text.

        The conversation_text can be in any format, e.g.:
            User: Hi I need help with X
            Assistant: Sure, tell me more about X
            User: ...

        Returns a human-readable summary string.
        """
        if not conversation_text.strip():
            raise ValueError("Conversation text is empty.")

        payload = [
            {"role": "system", "content": _FULL_SUMMARY_PROMPT},
            {"role": "user", "content": f"Conversation:\n\n{conversation_text}"},
        ]
        raw = self._llm.complete(payload)
        summary = _parse_summary(raw)
        return summary.to_readable_text()

    def force_compress(self) -> ConversationSummary:
        """
        Manually trigger compression of the entire buffer, regardless of size.
        Unlike auto-compression, this does not retain any recent turns verbatim —
        everything in the buffer is folded into the summary.
        """
        if not self._buffer:
            if self._summary is None:
                raise ValueError("Nothing to compress: buffer is empty and no summary exists.")
            return self._summary

        to_summarize = self._buffer
        self._buffer = []

        if self._summary is None:
            self._summary = self._full_summarize(to_summarize)
        else:
            self._summary = self._incremental_update(self._summary, to_summarize)

        return self._summary

    @property
    def current_summary(self) -> ConversationSummary | None:
        """The latest rolling summary (None until first compression fires)."""
        return self._summary

    @property
    def buffered_count(self) -> int:
        """Number of messages in the buffer not yet compressed."""
        return len(self._buffer)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _compress(self) -> None:
        cutoff       = max(0, len(self._buffer) - self._keep_recent)
        to_summarize = self._buffer[:cutoff]
        self._buffer = self._buffer[cutoff:]

        if not to_summarize:
            return

        if self._summary is None:
            self._summary = self._full_summarize(to_summarize)
        else:
            self._summary = self._incremental_update(self._summary, to_summarize)

    def _full_summarize(self, messages: list[Message]) -> ConversationSummary:
        payload = [
            {"role": "system", "content": _FULL_SUMMARY_PROMPT},
            {"role": "user",   "content": "Conversation:\n\n" + _format_transcript(messages)},
        ]
        return _parse_summary(self._llm.complete(payload))

    def _incremental_update(
        self,
        existing: ConversationSummary,
        new_messages: list[Message],
    ) -> ConversationSummary:
        payload = [
            {"role": "system", "content": _INCREMENTAL_SUMMARY_PROMPT},
            {
                "role":    "user",
                "content": (
                    f"Existing summary:\n{existing.to_context_block()}\n\n"
                    f"New turns:\n{_format_transcript(new_messages)}"
                ),
            },
        ]
        return _parse_summary(self._llm.complete(payload))


# ---------------------------------------------------------------------------
# Module-level helpers
# ---------------------------------------------------------------------------

def _format_transcript(messages: list[Message]) -> str:
    return "\n\n".join(
        f"[{m.role.value.upper()}]: {m.content}" for m in messages
    )


def _parse_summary(raw: str) -> ConversationSummary:
    # Strip common LLM markdown fence variants
    stripped = raw.strip()
    for prefix in ("```json", "```", "json"):
        if stripped.startswith(prefix):
            stripped = stripped[len(prefix):]
            break
    if stripped.endswith("```"):
        stripped = stripped[:-3]
    cleaned = stripped.strip()
    try:
        data: dict = json.loads(cleaned)
    except json.JSONDecodeError as exc:
        raise ValueError(f"LLM returned non-JSON output:\n{raw}") from exc

    return ConversationSummary(
        overview=data.get("overview", ""),
        user_intent=data.get("user_intent", ""),
        key_decisions=data.get("key_decisions", []),
        action_items=data.get("action_items", []),
    )
