core/session_memory.py

"""
core/session_memory.py — session snapshot persistence with optional HF Dataset sync.
"""

from __future__ import annotations

import io
import json
import os
import re
import shutil
import tempfile
from pathlib import Path
from typing import Any, Dict, Optional

try:
    from huggingface_hub import HfApi, hf_hub_download
except Exception:  # pragma: no cover
    HfApi = None  # type: ignore
    hf_hub_download = None  # type: ignore


class SessionMemoryManager:
    def __init__(
        self,
        workspace: str,
        dataset_repo: Optional[str] = None,
        token: Optional[str] = None,
    ) -> None:
        self.root = Path(workspace).resolve() / "sessions"
        self.root.mkdir(parents=True, exist_ok=True)
        self.dataset_repo = dataset_repo or os.environ.get("DEVAI_SESSION_DATASET") or os.environ.get("HF_SESSION_DATASET_REPO")
        self.token = token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
        self.api = HfApi(token=self.token) if (self.dataset_repo and self.token and HfApi is not None) else None

    @staticmethod
    def sanitize_session_id(session_id: str) -> str:
        cleaned = re.sub(r"[^A-Za-z0-9._-]+", "-", str(session_id or "anon")).strip("-._")
        return cleaned[:120] or "anon"

    def _session_dir(self, session_id: str) -> Path:
        d = self.root / self.sanitize_session_id(session_id)
        d.mkdir(parents=True, exist_ok=True)
        return d

    def _attachments_dir(self, session_id: str) -> Path:
        d = self._session_dir(session_id) / "attachments"
        d.mkdir(parents=True, exist_ok=True)
        return d

    def _snapshot_path(self, session_id: str) -> Path:
        return self._session_dir(session_id) / "latest.json"

    def _copy_attachment(self, src_path: Optional[str], dst_name: str, session_id: str) -> Optional[str]:
        if not src_path or not os.path.exists(src_path):
            return None
        dst = self._attachments_dir(session_id) / dst_name
        shutil.copyfile(src_path, dst)
        return str(dst)

    def _push_to_dataset(self, session_id: str, snapshot: Dict[str, Any]) -> None:
        if not self.api or not self.dataset_repo:
            return
        payload = json.dumps(snapshot, ensure_ascii=False, indent=2, default=str).encode("utf-8")
        self.api.upload_file(
            path_or_fileobj=io.BytesIO(payload),
            path_in_repo=f"sessions/{self.sanitize_session_id(session_id)}/latest.json",
            repo_id=self.dataset_repo,
            repo_type="dataset",
        )

    def _pull_from_dataset(self, session_id: str) -> Optional[Dict[str, Any]]:
        if not (self.dataset_repo and self.token and hf_hub_download is not None):
            return None
        try:
            local = hf_hub_download(
                repo_id=self.dataset_repo,
                repo_type="dataset",
                filename=f"sessions/{self.sanitize_session_id(session_id)}/latest.json",
                token=self.token,
                local_dir=tempfile.mkdtemp(prefix="devai_session_restore_"),
            )
            return json.loads(Path(local).read_text(encoding="utf-8"))
        except Exception:
            return None

    def save_ui_state(self, session_id: str, snapshot: Dict[str, Any]) -> Dict[str, Any]:
        sid = self.sanitize_session_id(session_id)
        payload = dict(snapshot or {})
        payload["session_id"] = sid
        payload["zip_input_path"] = self._copy_attachment(payload.get("zip_input_path"), "last_input.zip", sid)
        payload["zip_output_path"] = self._copy_attachment(payload.get("zip_output_path"), "last_output.zip", sid)
        path = self._snapshot_path(sid)
        path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, default=str), encoding="utf-8")
        try:
            self._push_to_dataset(sid, payload)
        except Exception:
            pass
        return payload

    def restore_ui_state(self, session_id: str) -> Optional[Dict[str, Any]]:
        sid = self.sanitize_session_id(session_id)
        path = self._snapshot_path(sid)
        if path.exists():
            try:
                return json.loads(path.read_text(encoding="utf-8"))
            except Exception:
                pass
        pulled = self._pull_from_dataset(sid)
        if pulled:
            try:
                path.write_text(json.dumps(pulled, ensure_ascii=False, indent=2, default=str), encoding="utf-8")
            except Exception:
                pass
        return pulled


__all__ = ["SessionMemoryManager"]