rp/pr/core/session.py

import json
import os
from datetime import datetime
from typing import Dict, List, Optional

from pr.core.logging import get_logger

logger = get_logger("session")

SESSIONS_DIR = os.path.expanduser("~/.assistant_sessions")


class SessionManager:

    def __init__(self):
        os.makedirs(SESSIONS_DIR, exist_ok=True)

    def save_session(
        self, name: str, messages: List[Dict], metadata: Optional[Dict] = None
    ) -> bool:
        try:
            session_file = os.path.join(SESSIONS_DIR, f"{name}.json")

            session_data = {
                "name": name,
                "created_at": datetime.now().isoformat(),
                "messages": messages,
                "metadata": metadata or {},
            }

            with open(session_file, "w") as f:
                json.dump(session_data, f, indent=2)

            logger.info(f"Session saved: {name}")
            return True

        except Exception as e:
            logger.error(f"Error saving session {name}: {e}")
            return False

    def load_session(self, name: str) -> Optional[Dict]:
        try:
            session_file = os.path.join(SESSIONS_DIR, f"{name}.json")

            if not os.path.exists(session_file):
                logger.warning(f"Session not found: {name}")
                return None

            with open(session_file) as f:
                session_data = json.load(f)

            logger.info(f"Session loaded: {name}")
            return session_data

        except Exception as e:
            logger.error(f"Error loading session {name}: {e}")
            return None

    def list_sessions(self) -> List[Dict]:
        sessions = []

        try:
            for filename in os.listdir(SESSIONS_DIR):
                if filename.endswith(".json"):
                    filepath = os.path.join(SESSIONS_DIR, filename)
                    try:
                        with open(filepath) as f:
                            data = json.load(f)

                        sessions.append(
                            {
                                "name": data.get("name", filename[:-5]),
                                "created_at": data.get("created_at", "unknown"),
                                "message_count": len(data.get("messages", [])),
                                "metadata": data.get("metadata", {}),
                            }
                        )
                    except Exception as e:
                        logger.warning(f"Error reading session file {filename}: {e}")

            sessions.sort(key=lambda x: x["created_at"], reverse=True)

        except Exception as e:
            logger.error(f"Error listing sessions: {e}")

        return sessions

    def delete_session(self, name: str) -> bool:
        try:
            session_file = os.path.join(SESSIONS_DIR, f"{name}.json")

            if not os.path.exists(session_file):
                logger.warning(f"Session not found: {name}")
                return False

            os.remove(session_file)
            logger.info(f"Session deleted: {name}")
            return True

        except Exception as e:
            logger.error(f"Error deleting session {name}: {e}")
            return False

    def export_session(self, name: str, output_path: str, format: str = "json") -> bool:
        session_data = self.load_session(name)
        if not session_data:
            return False

        try:
            if format == "json":
                with open(output_path, "w") as f:
                    json.dump(session_data, f, indent=2)

            elif format == "markdown":
                with open(output_path, "w") as f:
                    f.write(f"# Session: {name}\n\n")
                    f.write(f"Created: {session_data['created_at']}\n\n")
                    f.write("---\n\n")

                    for msg in session_data["messages"]:
                        role = msg.get("role", "unknown")
                        content = msg.get("content", "")

                        f.write(f"## {role.capitalize()}\n\n")
                        f.write(f"{content}\n\n")
                        f.write("---\n\n")

            elif format == "txt":
                with open(output_path, "w") as f:
                    f.write(f"Session: {name}\n")
                    f.write(f"Created: {session_data['created_at']}\n")
                    f.write("=" * 80 + "\n\n")

                    for msg in session_data["messages"]:
                        role = msg.get("role", "unknown")
                        content = msg.get("content", "")

                        f.write(f"[{role.upper()}]\n")
                        f.write(f"{content}\n")
                        f.write("-" * 80 + "\n\n")

            else:
                logger.error(f"Unsupported export format: {format}")
                return False

            logger.info(f"Session exported to {output_path}")
            return True

        except Exception as e:
            logger.error(f"Error exporting session: {e}")
            return False