core/session_manager.py

import json
import re
import uuid
from datetime import datetime
from pathlib import Path


class SessionManager:

    def __init__(self, data_dir):
        self.sessions_dir = Path(data_dir) / "sessions"
        self.sessions_dir.mkdir(parents=True, exist_ok=True)

        self._current_id = None
        self._current_path = None



    def start_new_session(self, system_message):
        session_id = uuid.uuid4().hex[:12]
        now = datetime.now()

        session = {
            "id": session_id,
            "title": "Untitled Session",
            "created_at": now.isoformat(timespec="seconds"),
            "updated_at": now.isoformat(timespec="seconds"),
            "message_count": 0,
            "messages": [system_message],
        }

        file_name = self._build_filename(now, "untitled")
        file_path = self.sessions_dir / file_name

        self._current_id = session_id
        self._current_path = file_path
        self._write_session(file_path, session)
        return session

    def save_messages(self, messages):
        if self._current_path is None:
            return

        session = self._read_session(self._current_path)
        if session is None:
            session = {
                "id": self._current_id or uuid.uuid4().hex[:12],
                "title": "Restored Session",
                "created_at": datetime.now().isoformat(timespec="seconds"),
            }
            if not self._current_id:
                self._current_id = session["id"]

        # Update title from the first user message if still untitled
        if session["title"] == "Untitled Session":
            title = self._extract_title(messages)
            if title:
                session["title"] = title
                # Rename the file to include the real title
                new_path = self._rename_with_title(self._current_path, title)
                if new_path:
                    self._current_path = new_path

        session["messages"] = messages
        session["message_count"] = len([m for m in messages if m["role"] != "system"])
        session["updated_at"] = datetime.now().isoformat(timespec="seconds")
        self._write_session(self._current_path, session)

    def list_sessions(self):
        sessions = []
        for file_path in sorted(self.sessions_dir.glob("*.json"), reverse=True):
            session = self._read_session(file_path)
            if session is None:
                continue
            sessions.append({
                "path": file_path,
                "id": session.get("id", ""),
                "title": session.get("title", "Untitled"),
                "created_at": session.get("created_at", ""),
                "updated_at": session.get("updated_at", ""),
                "message_count": session.get("message_count", 0),
            })
        return sessions

    def load_session(self, index, system_message):
        sessions = self.list_sessions()
        if index < 1 or index > len(sessions):
            return None

        entry = sessions[index - 1]
        session = self._read_session(entry["path"])
        if session is None:
            return None

        messages = session.get("messages", [])

        # Ensure the system message is current (in case the prompt was updated)
        if messages and messages[0].get("role") == "system":
            messages[0] = system_message
        else:
            messages.insert(0, system_message)

        self._current_id = session.get("id", "")
        self._current_path = entry["path"]
        return messages

    def delete_session(self, index):
        sessions = self.list_sessions()
        if index < 1 or index > len(sessions):
            return None

        entry = sessions[index - 1]
        title = entry["title"]
        try:
            entry["path"].unlink(missing_ok=True)
        except OSError:
            return None

        # If we deleted the currently active session, clear state
        if self._current_path == entry["path"]:
            self._current_id = None
            self._current_path = None

        return title

    def format_session_list(self):
        sessions = self.list_sessions()
        if not sessions:
            return "No saved sessions found."

        lines = ["", "  Saved Sessions", "  " + "-" * 50]
        for index, entry in enumerate(sessions, start=1):
            count = entry["message_count"]
            updated = entry["updated_at"][:16].replace("T", " ")
            title = entry["title"][:50]
            lines.append(f"  {index}. {title}")
            lines.append(f"     {count} messages | Last updated: {updated}")
        lines.append("")
        return "\n".join(lines)



    def _build_filename(self, timestamp, slug):
        date_part = timestamp.strftime("%Y-%m-%d_%H-%M")
        safe_slug = self._slugify(slug)[:40]
        return f"{date_part}_{safe_slug}.json"

    def _rename_with_title(self, old_path, title):
        old_path = Path(old_path)
        if not old_path.exists():
            return None

        # Extract the date prefix from the old filename
        old_name = old_path.stem
        parts = old_name.split("_", 3)  # e.g. 2026-05-11_18-30_untitled
        if len(parts) >= 2:
            date_prefix = f"{parts[0]}_{parts[1]}"
        else:
            date_prefix = old_name[:16]

        safe_title = self._slugify(title)[:40]
        new_name = f"{date_prefix}_{safe_title}.json"
        new_path = old_path.parent / new_name

        if new_path == old_path:
            return old_path

        try:
            old_path.rename(new_path)
            return new_path
        except OSError:
            return old_path

    def _extract_title(self, messages):
        for message in messages:
            if message.get("role") == "user":
                content = message.get("content", "").strip()
                if not content:
                    continue
                # Remove source prefixes like "Source: PDF file (doc.pdf)\n\n..."
                if content.startswith("Source:"):
                    lines = content.split("\n", 2)
                    title = lines[0].strip()
                else:
                    title = content
                # Truncate to something readable
                if len(title) > 60:
                    title = title[:57] + "..."
                return title
        return ""

    def _slugify(self, text):
        slug = text.lower().strip()
        slug = re.sub(r"[^\w\s-]", "", slug)
        slug = re.sub(r"[\s_]+", "-", slug)
        slug = slug.strip("-")
        return slug or "untitled"

    def _read_session(self, file_path):
        try:
            return json.loads(Path(file_path).read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            return None

    def _write_session(self, file_path, session):
        Path(file_path).write_text(
            json.dumps(session, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )