"""autoplay_sdk.summarizer — Client-side session summarizer for context-window management.

Accumulates real-time user actions per session.  When the number of actions
reaches a configurable threshold, the summarizer calls your LLM with a
prompt to produce a concise prose summary, then fires an ``on_summary``
callback with the result and resets the session's history.

This keeps your RAG context window compact regardless of how long a session
runs — instead of growing unboundedly, each session is represented by a
rolling summary of the most recent ``threshold`` actions.

Sync usage::

    import openai
    from autoplay_sdk import ConnectorClient
    from autoplay_sdk.summarizer import SessionSummarizer

    openai_client = openai.OpenAI()

    def my_llm(prompt: str) -> str:
        return openai_client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
        ).choices[0].message.content

    def store_summary(session_id: str, summary: str) -> None:
        print(f"[{session_id}] {summary}")

    summarizer = SessionSummarizer(
        llm=my_llm,
        threshold=10,          # summarize every 10 actions
        on_summary=store_summary,
        # prompt=custom_prompt,  # optional — uses built-in default
    )

    ConnectorClient(url=URL, token=TOKEN) \\
        .on_actions(summarizer.add) \\
        .run()

Async usage::

    from autoplay_sdk.summarizer import AsyncSessionSummarizer

    async def my_async_llm(prompt: str) -> str:
        r = await async_openai.chat.completions.create(...)
        return r.choices[0].message.content

    async def store_summary(session_id: str, summary: str) -> None:
        await db.save(session_id, summary)

    summarizer = AsyncSessionSummarizer(
        llm=my_async_llm,
        threshold=10,
        on_summary=store_summary,
    )

    async with AsyncConnectorClient(url=URL, token=TOKEN) as client:
        client.on_actions(summarizer.add)
        await client.run()

Composing with RagPipeline::

    from autoplay_sdk.rag import RagPipeline
    from autoplay_sdk.summarizer import SessionSummarizer

    summarizer = SessionSummarizer(llm=my_llm, threshold=10)
    pipeline   = RagPipeline(embed=embed_fn, upsert=upsert_fn, summarizer=summarizer)

    client.on_actions(pipeline.on_actions).on_summary(pipeline.on_summary).run()
"""

from __future__ import annotations

import logging
import threading
import time
from collections import defaultdict
from collections.abc import Callable
from typing import Any

from autoplay_sdk.models import ActionsPayload

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Default prompt
# ---------------------------------------------------------------------------

DEFAULT_PROMPT = """\
You are summarising a user's in-app session for a RAG pipeline.
Your summary will be embedded and stored in a vector database to provide
context for a chatbot or AI assistant.

Write a concise 2-3 sentence summary of what the user did, focusing on:
- Which features or pages they visited
- What actions they performed
- Any clear intent or goal you can infer

Actions:
{actions}

Summary:"""


def _format_actions(payloads: list[ActionsPayload]) -> str:
    """Format a list of ActionsPayload into a numbered action list for the prompt."""
    lines: list[str] = []
    n = 1
    for payload in payloads:
        for action in payload.actions:
            lines.append(f"{n}. {action.to_text()}")
            n += 1
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Sync summarizer
# ---------------------------------------------------------------------------


class SessionSummarizer:
    """Accumulates actions per session and summarises them when a threshold is reached.

    Args:
        llm:        ``(prompt: str) -> str`` — any synchronous LLM callable.
                    Receives the formatted prompt and must return the summary
                    as a plain string.
        threshold:  Number of **individual actions** (not batches) to
                    accumulate before triggering a summarisation.  Must be
                    >= 1.  Default 10.
        prompt:     Custom prompt template string.  Must contain ``{actions}``
                    as a placeholder for the formatted action list.  If
                    ``None``, the built-in default prompt is used.
        on_summary: ``(session_id: str, summary: str) -> None`` — called
                    after each summarisation.  Wire this to your storage,
                    vector store, or ``RagPipeline``.  Can also be set later
                    by assigning to ``summarizer.on_summary``.
    """

    def __init__(
        self,
        llm: Callable[[str], str],
        threshold: int = 10,
        prompt: str | None = None,
        on_summary: Callable[[str, str], None] | None = None,
    ) -> None:
        if threshold < 1:
            raise ValueError(f"threshold must be >= 1, got {threshold!r}")
        self._llm = llm
        self._threshold = threshold
        self._prompt_template = prompt or DEFAULT_PROMPT
        self.on_summary = on_summary

        # session_id → list of ActionsPayload batches
        self._history: dict[str, list[ActionsPayload]] = defaultdict(list)
        # session_id → total action count
        self._counts: dict[str, int] = defaultdict(int)
        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def add(self, payload: ActionsPayload) -> None:
        """Receive an actions batch, accumulate, and summarise if threshold is hit.

        Wire directly to ``on_actions``::

            client.on_actions(summarizer.add)

        Args:
            payload: Typed ``ActionsPayload`` from the event stream.
        """
        session_id = payload.session_id or "unknown"
        action_count = len(payload.actions)

        with self._lock:
            self._history[session_id].append(payload)
            self._counts[session_id] += action_count
            should_summarise = self._counts[session_id] >= self._threshold
            if should_summarise:
                history_snapshot = list(self._history[session_id])
                self._history[session_id] = []
                self._counts[session_id] = 0

        if should_summarise:
            self._summarise(session_id, history_snapshot)

    def get_context(self, session_id: str) -> str:
        """Return the accumulated (not-yet-summarised) actions as text.

        Useful for including partial session context in a RAG query before
        the threshold is reached.

        Args:
            session_id: The session identifier.

        Returns:
            Formatted action list string, or empty string if no history.
        """
        with self._lock:
            history = list(self._history.get(session_id, []))
        return _format_actions(history)

    def reset(self, session_id: str) -> None:
        """Clear accumulated history for a session without summarising.

        Args:
            session_id: The session to reset.
        """
        with self._lock:
            self._history.pop(session_id, None)
            self._counts.pop(session_id, None)

    @property
    def active_sessions(self) -> list[str]:
        """List of session IDs with pending (not-yet-summarised) actions."""
        with self._lock:
            return [sid for sid, count in self._counts.items() if count > 0]

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _summarise(self, session_id: str, history: list[ActionsPayload]) -> None:
        actions_text = _format_actions(history)
        prompt = self._prompt_template.format(actions=actions_text)
        action_count = sum(len(p.actions) for p in history)

        summary: str | None = None
        t0 = time.perf_counter()
        try:
            summary = self._llm(prompt)
            elapsed_ms = (time.perf_counter() - t0) * 1000
            logger.debug(
                "summarizer: summarised session=%s actions=%d elapsed_ms=%.1f",
                session_id,
                action_count,
                elapsed_ms,
            )
        except Exception as exc:
            logger.error(
                "summarizer: LLM call failed for session=%s — restoring history: %s",
                session_id,
                exc,
                exc_info=True,
            )
            with self._lock:
                self._history[session_id] = history + self._history[session_id]
                self._counts[session_id] += action_count
            return

        if self.on_summary is not None:
            try:
                self.on_summary(session_id, summary)
            except Exception as exc:
                logger.error(
                    "summarizer: on_summary callback raised for session=%s: %s",
                    session_id,
                    exc,
                    exc_info=True,
                )


# ---------------------------------------------------------------------------
# Async summarizer
# ---------------------------------------------------------------------------


class AsyncSessionSummarizer:
    """Async version of ``SessionSummarizer`` for use with ``AsyncConnectorClient``.

    ``llm`` and ``on_summary`` are async callables (coroutines).

    Args:
        llm:        ``async (prompt: str) -> str``
        threshold:  Number of individual actions before summarisation.  Must
                    be >= 1.  Default 10.
        prompt:     Custom prompt template (must contain ``{actions}``).
        on_summary: ``async (session_id: str, summary: str) -> None``
    """

    def __init__(
        self,
        llm: Callable[[str], Any],
        threshold: int = 10,
        prompt: str | None = None,
        on_summary: Callable[[str, str], Any] | None = None,
    ) -> None:
        if threshold < 1:
            raise ValueError(f"threshold must be >= 1, got {threshold!r}")
        self._llm = llm
        self._threshold = threshold
        self._prompt_template = prompt or DEFAULT_PROMPT
        self.on_summary = on_summary

        self._history: dict[str, list[ActionsPayload]] = defaultdict(list)
        self._counts: dict[str, int] = defaultdict(int)
        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    async def add(self, payload: ActionsPayload) -> None:
        """Receive an actions batch, accumulate, and summarise if threshold is hit.

        Wire directly to ``on_actions``::

            client.on_actions(summarizer.add)
        """
        session_id = payload.session_id or "unknown"
        action_count = len(payload.actions)

        with self._lock:
            self._history[session_id].append(payload)
            self._counts[session_id] += action_count
            should_summarise = self._counts[session_id] >= self._threshold
            if should_summarise:
                history_snapshot = list(self._history[session_id])
                self._history[session_id] = []
                self._counts[session_id] = 0

        if should_summarise:
            await self._summarise(session_id, history_snapshot)

    def get_context(self, session_id: str) -> str:
        """Return the accumulated (not-yet-summarised) actions as text."""
        with self._lock:
            history = list(self._history.get(session_id, []))
        return _format_actions(history)

    def reset(self, session_id: str) -> None:
        """Clear accumulated history for a session without summarising."""
        with self._lock:
            self._history.pop(session_id, None)
            self._counts.pop(session_id, None)

    @property
    def active_sessions(self) -> list[str]:
        """List of session IDs with pending (not-yet-summarised) actions."""
        with self._lock:
            return [sid for sid, count in self._counts.items() if count > 0]

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    async def _summarise(self, session_id: str, history: list[ActionsPayload]) -> None:
        actions_text = _format_actions(history)
        prompt = self._prompt_template.format(actions=actions_text)
        action_count = sum(len(p.actions) for p in history)

        summary: str | None = None
        t0 = time.perf_counter()
        try:
            summary = await self._llm(prompt)
            elapsed_ms = (time.perf_counter() - t0) * 1000
            logger.debug(
                "summarizer: summarised session=%s actions=%d elapsed_ms=%.1f",
                session_id,
                action_count,
                elapsed_ms,
            )
        except Exception as exc:
            logger.error(
                "summarizer: LLM call failed for session=%s — restoring history: %s",
                session_id,
                exc,
                exc_info=True,
            )
            with self._lock:
                self._history[session_id] = history + self._history[session_id]
                self._counts[session_id] += action_count
            return

        if self.on_summary is not None:
            try:
                result = self.on_summary(session_id, summary)
                if hasattr(result, "__await__"):
                    await result
            except Exception as exc:
                logger.error(
                    "summarizer: on_summary callback raised for session=%s: %s",
                    session_id,
                    exc,
                    exc_info=True,
                )
