OpenSpace/openspace/communication/session_store.py

from __future__ import annotations

import json
import re
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

from openspace.utils.logging import Logger

from .types import ChannelAttachment, ChannelMessage, ChannelSession, ChannelSource

logger = Logger.get_logger(__name__)


class SessionStore:
    def __init__(
        self,
        sessions_dir: Path,
        *,
        workspace_root: Optional[Path] = None,
    ):
        self.sessions_dir = sessions_dir
        self.workspace_root = workspace_root
        self.sessions_dir.mkdir(parents=True, exist_ok=True)
        if self.workspace_root is not None:
            self.workspace_root.mkdir(parents=True, exist_ok=True)

    def get_or_create_session(self, source: ChannelSource) -> ChannelSession:
        session_key = build_session_key(source)
        session_dir = self.sessions_dir / session_key
        session_dir.mkdir(parents=True, exist_ok=True)

        metadata_path = session_dir / "session.json"
        transcript_path = session_dir / "transcript.jsonl"
        attachments_dir = session_dir / "attachments"
        workspace_dir = (
            self.workspace_root / session_key
            if self.workspace_root is not None
            else session_dir / "workspace"
        )
        attachments_dir.mkdir(parents=True, exist_ok=True)
        workspace_dir.mkdir(parents=True, exist_ok=True)

        now = _utcnow_iso()
        if metadata_path.exists():
            with open(metadata_path, "r", encoding="utf-8") as handle:
                data = json.load(handle)
            session = ChannelSession.from_dict(data)
            session.source = source
            session.updated_at = now
        else:
            session = ChannelSession(
                session_key=session_key,
                source=source,
                session_dir=str(session_dir),
                workspace_dir=str(workspace_dir),
                attachments_dir=str(attachments_dir),
                transcript_path=str(transcript_path),
                metadata_path=str(metadata_path),
                created_at=now,
                updated_at=now,
            )

        self._write_session_metadata(session)
        return session

    def append_user_message(self, session: ChannelSession, message: ChannelMessage) -> None:
        self._append_transcript_entry(
            session,
            {
                "entry_id": uuid.uuid4().hex,
                "role": "user",
                "content": message.text,
                "platform_message_id": message.message_id,
                "reply_to_message_id": message.reply_to_message_id,
                "reply_to_text": message.reply_to_text,
                "mentions_bot": message.mentions_bot,
                "attachments": [attachment.to_context_dict() for attachment in message.attachments],
                "source": message.source.to_dict(),
                "metadata": message.metadata,
                "timestamp": message.received_at.isoformat(),
            },
        )

    def append_assistant_message(
        self,
        session: ChannelSession,
        *,
        content: str,
        platform_message_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        self._append_transcript_entry(
            session,
            {
                "entry_id": uuid.uuid4().hex,
                "role": "assistant",
                "content": content,
                "platform_message_id": platform_message_id,
                "metadata": metadata or {},
                "timestamp": _utcnow_iso(),
            },
        )

    def load_history(self, session: ChannelSession, max_turns: int) -> List[Dict[str, str]]:
        entries = self._read_transcript_entries(session)
        if not entries:
            return []

        selected: List[Dict[str, str]] = []
        user_messages = 0
        for entry in reversed(entries):
            role = entry.get("role")
            if role not in {"user", "assistant"}:
                continue
            if role == "assistant" and not _assistant_entry_visible_in_history(entry):
                continue
            content = str(entry.get("content", "")).strip()
            if not content:
                continue
            selected.append({"role": role, "content": content})
            if role == "user":
                user_messages += 1
                if user_messages >= max_turns:
                    break
        selected.reverse()
        return selected

    def is_reply_to_assistant(self, session: ChannelSession, message_id: Optional[str]) -> bool:
        if not message_id:
            return False
        for entry in reversed(self._read_transcript_entries(session)):
            if entry.get("platform_message_id") == message_id:
                return entry.get("role") == "assistant" and _assistant_entry_visible_in_history(entry)
        return False

    def list_sessions(self) -> List[ChannelSession]:
        sessions: List[ChannelSession] = []
        for metadata_path in sorted(self.sessions_dir.glob("*/session.json")):
            try:
                with open(metadata_path, "r", encoding="utf-8") as handle:
                    sessions.append(ChannelSession.from_dict(json.load(handle)))
            except Exception as exc:
                logger.warning("Failed to load session metadata %s: %s", metadata_path, exc)
        return sessions

    def _append_transcript_entry(self, session: ChannelSession, entry: Dict[str, Any]) -> None:
        transcript_path = Path(session.transcript_path)
        transcript_path.parent.mkdir(parents=True, exist_ok=True)
        with open(transcript_path, "a", encoding="utf-8") as handle:
            handle.write(json.dumps(entry, ensure_ascii=False) + "\n")
        session.updated_at = _utcnow_iso()
        self._write_session_metadata(session)

    def _write_session_metadata(self, session: ChannelSession) -> None:
        with open(session.metadata_path, "w", encoding="utf-8") as handle:
            json.dump(session.to_dict(), handle, ensure_ascii=False, indent=2)

    def _read_transcript_entries(self, session: ChannelSession) -> List[Dict[str, Any]]:
        transcript_path = Path(session.transcript_path)
        if not transcript_path.exists():
            return []
        entries: List[Dict[str, Any]] = []
        with open(transcript_path, "r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if not line:
                    continue
                try:
                    entries.append(json.loads(line))
                except json.JSONDecodeError:
                    logger.warning("Skipping malformed transcript line in %s", transcript_path)
        return entries


def build_session_key(source: ChannelSource) -> str:
    parts = [source.platform.value, _sanitize(source.chat_id)]
    if source.thread_id:
        parts.append(_sanitize(source.thread_id))
    return "__".join(part for part in parts if part)


def _sanitize(value: str) -> str:
    value = re.sub(r"[^a-zA-Z0-9._-]+", "-", str(value).strip())
    value = value.strip("-._")
    return value or "unknown"


def _utcnow_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _assistant_entry_visible_in_history(entry: Dict[str, Any]) -> bool:
    metadata = entry.get("metadata")
    if not isinstance(metadata, dict):
        return True
    return metadata.get("send_success") is not False