import json
import os

from .base import Session, SessionStore, sanitize_id


class JSONSessionStore(SessionStore):
    """A simple JSON file-based session store."""

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

    def _path(self, user_id: str, session_id: str) -> str:
        return os.path.join(
            self.base_dir,
            f"{sanitize_id(user_id, 'anon')}__{sanitize_id(session_id, 'default')}.json")

    def load(self, user_id: str, session_id: str) -> Session:
        path = self._path(user_id, session_id)
        messages = []
        if os.path.exists(path):
            try:
                with open(path, 'r', encoding='utf-8') as f:
                    messages = json.load(f).get('messages', [])
                    
            except (json.JSONDecodeError, OSError):
                messages = []
        
        return Session(session_id=session_id, user_id=user_id, messages=messages)


    def save(self, session: Session) -> None:
        path = self._path(session.user_id, session.session_id)
        tmp = path + ".tmp"
        
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(
                {
                    "user_id": session.user_id,
                    "session_id": session.session_id,
                    "messages": session.messages,
                },
                f,
                ensure_ascii=False,
                indent=2,
            )
        os.replace(tmp, path)

    def clear(self, user_id: str, session_id: str) -> None:
        path = self._path(user_id, session_id)
        if os.path.exists(path):
            os.remove(path)