"""llm_context.py — Per-person conversation context management.

No ROS2 dependencies. Manages sliding-window conversation history per person_id,
builds llama.cpp prompt, handles summary compression when context is full.

Used by conversation_node.py.
Tested by test/test_llm_context.py.
"""

from __future__ import annotations

import json
import os
import time
from typing import List, Optional, Tuple

# ── Constants ──────────────────────────────────────────────────────────────────
MAX_TURNS = 20          # Max turns to keep before summary compression
SUMMARY_KEEP = 4        # Keep last N turns after summarizing older ones
SYSTEM_PROMPT_PATH = "/soul/SOUL.md"
DEFAULT_SYSTEM_PROMPT = (
    "You are Salty, a friendly social robot. "
    "You are helpful, warm, and concise. "
    "Keep responses under 2 sentences unless asked to elaborate."
)


# ── Conversation history ───────────────────────────────────────────────────────

class PersonContext:
    """Conversation history + relationship memory for one person."""

    def __init__(self, person_id: str, person_name: str = "") -> None:
        self.person_id = person_id
        self.person_name = person_name or person_id
        self.turns: List[dict] = []  # {"role": "user"|"assistant", "content": str, "ts": float}
        self.summary: str = ""
        self.last_seen: float = time.time()
        self.interaction_count: int = 0

    def add_user(self, text: str) -> None:
        self.turns.append({"role": "user", "content": text, "ts": time.time()})
        self.last_seen = time.time()
        self.interaction_count += 1

    def add_assistant(self, text: str) -> None:
        self.turns.append({"role": "assistant", "content": text, "ts": time.time()})

    def needs_compression(self) -> bool:
        return len(self.turns) > MAX_TURNS

    def compress(self, summary_text: str) -> None:
        """Replace old turns with a summary, keeping last SUMMARY_KEEP turns."""
        old_summary = self.summary
        if old_summary:
            self.summary = old_summary + " " + summary_text
        else:
            self.summary = summary_text
        # Keep most recent turns
        self.turns = self.turns[-SUMMARY_KEEP:]

    def to_dict(self) -> dict:
        return {
            "person_id": self.person_id,
            "person_name": self.person_name,
            "turns": self.turns,
            "summary": self.summary,
            "last_seen": self.last_seen,
            "interaction_count": self.interaction_count,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "PersonContext":
        ctx = cls(d["person_id"], d.get("person_name", ""))
        ctx.turns = d.get("turns", [])
        ctx.summary = d.get("summary", "")
        ctx.last_seen = d.get("last_seen", 0.0)
        ctx.interaction_count = d.get("interaction_count", 0)
        return ctx


# ── Context store ──────────────────────────────────────────────────────────────

class ContextStore:
    """Persistent per-person conversation context with optional JSON backing."""

    def __init__(self, db_path: str = "/social_db/conversation_context.json") -> None:
        self._db_path = db_path
        self._contexts: dict = {}   # person_id → PersonContext
        self._load()

    def get(self, person_id: str, person_name: str = "") -> PersonContext:
        if person_id not in self._contexts:
            self._contexts[person_id] = PersonContext(person_id, person_name)
        elif person_name and not self._contexts[person_id].person_name:
            self._contexts[person_id].person_name = person_name
        return self._contexts[person_id]

    def save(self) -> None:
        os.makedirs(os.path.dirname(self._db_path), exist_ok=True)
        with open(self._db_path, "w") as f:
            json.dump(
                {pid: ctx.to_dict() for pid, ctx in self._contexts.items()},
                f, indent=2
            )

    def _load(self) -> None:
        if os.path.exists(self._db_path):
            try:
                with open(self._db_path) as f:
                    raw = json.load(f)
                self._contexts = {
                    pid: PersonContext.from_dict(d) for pid, d in raw.items()
                }
            except Exception:
                self._contexts = {}

    def all_persons(self) -> List[str]:
        return list(self._contexts.keys())


# ── Prompt builder ─────────────────────────────────────────────────────────────

def load_system_prompt(path: str = SYSTEM_PROMPT_PATH) -> str:
    """Load SOUL.md as system prompt. Falls back to default."""
    if os.path.exists(path):
        with open(path) as f:
            return f.read().strip()
    return DEFAULT_SYSTEM_PROMPT


def build_llama_prompt(
    ctx: PersonContext,
    user_text: str,
    system_prompt: str,
    group_persons: Optional[List[str]] = None,
) -> str:
    """Build Phi-3 / Llama-3 style chat prompt string.

    Uses the ChatML format supported by llama-cpp-python.
    """
    lines = [f"<|system|>\n{system_prompt}"]

    if ctx.summary:
        lines.append(
            f"<|system|>\n[Memory of {ctx.person_name}]: {ctx.summary}"
        )

    if group_persons:
        others = [p for p in group_persons if p != ctx.person_id]
        if others:
            lines.append(
                f"<|system|>\n[Also present]: {', '.join(others)}"
            )

    for turn in ctx.turns:
        role = "user" if turn["role"] == "user" else "assistant"
        lines.append(f"<|{role}|>\n{turn['content']}")

    lines.append(f"<|user|>\n{user_text}")
    lines.append("<|assistant|>")

    return "\n".join(lines)


def needs_summary_prompt(ctx: PersonContext) -> str:
    """Build a prompt asking the LLM to summarize old conversation turns."""
    turns_text = "\n".join(
        f"{t['role'].upper()}: {t['content']}"
        for t in ctx.turns[:-SUMMARY_KEEP]
    )
    return (
        f"<|system|>\nSummarize this conversation between {ctx.person_name} and "
        "Salty in one sentence, focusing on what was discussed and their relationship.\n"
        f"<|user|>\n{turns_text}\n<|assistant|>"
    )
