"""
cortex/sleep/s1_consolidator.py — S1: Session Consolidator

Issue #360, Stage 1 of the Sleep Pipeline.

Responsibilities:
  1. Gather all session claims for a given session_id
  2. Generate a concise session summary via LLM (capped at 50 claims to
     stay within token budgets)
  3. Persist the summary to the handoff_packet table (status='session_summary')
     and update session_record.summary

Design principles:
  - All methods are async
  - LLM failures are non-fatal: fallback to a deterministic summary
  - No dependency on the full CircadianEngine / SleepConsolidator stack
  - Storage interactions use the existing StorageAdapter interface

Persistence strategy:
  - handoff_packet row with status='session_summary' stores the LLM text
    and the raw claims (truncated) as payload_json
  - session_record.summary is also updated so retrieval pipelines can access
    the summary without querying handoff_packet directly
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING, List, Optional

if TYPE_CHECKING:
    from cortex.storage.adapter import StorageAdapter
    from cortex.types import LLMProvider, SemanticClaim

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# LLM prompts
# ---------------------------------------------------------------------------

_S1_SYSTEM = """\
You are a memory consolidation assistant for an AI agent.
Produce a concise, factual summary of the key takeaways from this session.
Focus on: durable facts, explicit decisions, user preferences, and stated goals.
Exclude: greetings, filler, transient context, and raw tool outputs.
Maximum length: 5 sentences. Write in plain prose."""

_S1_USER = """\
Session claims (subject-predicate-object triples extracted from this session):

{claims_json}

Write a brief summary capturing the session's most important insights."""


# ---------------------------------------------------------------------------
# Result dataclass
# ---------------------------------------------------------------------------

@dataclass
class S1Result:
    """Result of the S1 session consolidation stage."""

    session_id: str
    owner_id: str
    claim_count: int = 0
    summary: str = ""
    handoff_packet_id: Optional[str] = None
    completed_at: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )
    success: bool = True
    error: Optional[str] = None

    def to_dict(self) -> dict:
        return {
            "session_id": self.session_id,
            "owner_id": self.owner_id,
            "claim_count": self.claim_count,
            "summary": self.summary,
            "handoff_packet_id": self.handoff_packet_id,
            "completed_at": self.completed_at,
            "success": self.success,
            "error": self.error,
        }


# ---------------------------------------------------------------------------
# S1Consolidator
# ---------------------------------------------------------------------------


class S1Consolidator:
    """
    S1: Session Consolidator.

    Summarises the session's semantic claims using an LLM and persists
    the result to the handoff_packet table.

    Args:
        storage:  StorageAdapter instance (must support get_session_claims,
                  update_session, and create_handoff_packet).
        llm:      LLMProvider for summary generation.
        max_claims: Cap on claims sent to the LLM (default: 50).
    """

    def __init__(
        self,
        storage: "StorageAdapter",
        llm: "LLMProvider",
        max_claims: int = 50,
    ) -> None:
        self._storage = storage
        self._llm = llm
        self._max_claims = max_claims

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

    async def run(self, session_id: str, owner_id: str = "default") -> S1Result:
        """
        Run S1 for the given session.

        Steps:
          1. Gather session claims from storage
          2. Generate summary via LLM (with fallback)
          3. Persist to handoff_packet + update session_record.summary
          4. Return S1Result

        Args:
            session_id:  The session to consolidate.
            owner_id:    Owner/tenant identifier.

        Returns:
            S1Result with summary, packet id, and metadata.
        """
        logger.info("[S1] consolidating session=%s owner=%s", session_id, owner_id)

        result = S1Result(session_id=session_id, owner_id=owner_id)

        # Step 1: Gather claims
        claims = await self._gather_claims(session_id)
        result.claim_count = len(claims)

        if not claims:
            logger.info("[S1] no claims found for session=%s — empty summary", session_id)
            result.summary = "(no session claims to summarise)"
        else:
            # Step 2: Generate summary
            result.summary = await self._generate_summary(claims)

        # Step 3: Persist
        try:
            packet_id = await self._persist(
                session_id=session_id,
                owner_id=owner_id,
                summary=result.summary,
                claims=claims,
            )
            result.handoff_packet_id = packet_id
        except Exception as exc:
            logger.error("[S1] persist failed for session=%s: %s", session_id, exc)
            result.success = False
            result.error = str(exc)

        result.completed_at = datetime.now(timezone.utc).isoformat()
        logger.info(
            "[S1] done: session=%s claims=%d summary_len=%d packet=%s",
            session_id, result.claim_count, len(result.summary),
            result.handoff_packet_id,
        )
        return result

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    async def _gather_claims(self, session_id: str) -> List["SemanticClaim"]:
        """Retrieve all claims for this session from storage."""
        try:
            return await self._storage.get_session_claims(session_id)
        except Exception as exc:
            logger.error("[S1] get_session_claims failed: %s", exc)
            return []

    async def _generate_summary(self, claims: List["SemanticClaim"]) -> str:
        """
        Generate an LLM-based session summary.

        Falls back to a deterministic template if the LLM call fails.
        """
        capped = claims[: self._max_claims]
        claims_json = json.dumps(
            [
                {
                    "subject": c.subject,
                    "predicate": c.predicate,
                    "object": c.object_value,
                    "confidence": round(getattr(c, "confidence", 1.0), 3),
                }
                for c in capped
            ],
            ensure_ascii=False,
            indent=2,
        )

        try:
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(self._llm, "consolidation")
            resp = await self._llm.complete(
                system=_S1_SYSTEM,
                user=_S1_USER.format(claims_json=claims_json),
                response_format="text",
                **({"model": _routed_model} if _routed_model else {}),
            )
            text = resp.text if hasattr(resp, "text") else str(resp)
            summary = text.strip()
            logger.debug("[S1] LLM summary generated (%d chars)", len(summary))
            return summary
        except Exception as exc:
            logger.warning("[S1] LLM summary failed, using fallback: %s", exc)
            return self._fallback_summary(claims)

    def _fallback_summary(self, claims: List["SemanticClaim"]) -> str:
        """Deterministic summary when LLM is unavailable."""
        subjects = sorted({c.subject for c in claims})
        subject_str = ", ".join(subjects[:5])
        return (
            f"Session contained {len(claims)} memory claim(s) "
            f"covering: {subject_str or 'various topics'}. "
            f"(LLM summary unavailable — deterministic fallback)"
        )

    async def _persist(
        self,
        session_id: str,
        owner_id: str,
        summary: str,
        claims: List["SemanticClaim"],
    ) -> str:
        """
        Persist the session summary.

        Writes to:
          - handoff_packet (status='session_summary', payload_json contains claim stubs)
          - session_record.summary via update_session()

        Returns the handoff_packet id.
        """
        now = datetime.now(timezone.utc).isoformat()

        # Build compact payload (not full claims — just stubs for auditability)
        payload = {
            "stage": "S1",
            "claim_count": len(claims),
            "claim_stubs": [
                {
                    "id": c.id,
                    "subject": c.subject,
                    "predicate": c.predicate,
                    "object": c.object_value,
                }
                for c in claims[:20]  # top 20 stubs only
            ],
        }

        # Persist handoff_packet row
        packet_id = await self._storage.create_handoff_packet(
            owner_id=owner_id,
            session_id=session_id,
            status="session_summary",
            summary_text=summary,
            payload_json=json.dumps(payload, ensure_ascii=False),
            generated_at=now,
        )

        # Update session_record.summary for easy retrieval
        try:
            await self._storage.update_session(
                session_id,
                summary=summary,
            )
        except Exception as exc:
            logger.warning("[S1] update_session(summary=) failed: %s", exc)

        return packet_id
