notebook_store.py

from __future__ import annotations

import json
import shutil
import re
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable, Optional


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


def storage_root() -> Path:
    return Path(__file__).with_name("storage")


def notebooks_root() -> Path:
    return storage_root() / "notebooks"


@dataclass(frozen=True)
class NotebookMeta:
    id: str
    name: str
    created_at: str


def notebook_dir(notebook_id: str) -> Path:
    return notebooks_root() / notebook_id


def notebook_meta_path(notebook_id: str) -> Path:
    return notebook_dir(notebook_id) / "notebook.json"


def _safe_folder_name(name: str) -> str:
    """
    Make a Windows-safe folder name that still looks like the notebook name.
    """
    name = (name or "").strip()
    if not name:
        return "Notebook"

    # Replace characters that are illegal in Windows file/folder names.
    name = re.sub(r'[<>:"/\\\\|?*]', "_", name)
    name = re.sub(r"\s+", " ", name).strip()
    # Avoid trailing dots/spaces (Windows doesn't like them)
    name = name.rstrip(" .")
    return name or "Notebook"


def _unique_folder_name(base: str) -> str:
    base = _safe_folder_name(base)
    root = notebooks_root()
    root.mkdir(parents=True, exist_ok=True)

    if not (root / base).exists():
        return base

    i = 2
    while True:
        candidate = f"{base} ({i})"
        if not (root / candidate).exists():
            return candidate
        i += 1


def _rewrite_jsonl_notebook_id(path: Path, *, old_id: str, new_id: str) -> None:
    if not path.exists():
        return
    lines = path.read_text(encoding="utf-8").splitlines()
    out_lines: list[str] = []
    for line in lines:
        raw = line.strip()
        if not raw:
            continue
        try:
            obj = json.loads(raw)
        except json.JSONDecodeError:
            continue
        if obj.get("notebook_id") == old_id:
            obj["notebook_id"] = new_id
        out_lines.append(json.dumps(obj, ensure_ascii=False))
    path.write_text("\n".join(out_lines) + ("\n" if out_lines else ""), encoding="utf-8")


def migrate_notebook_folders_to_names() -> list[NotebookMeta]:
    """
    If older notebooks were stored under random IDs (uuid hex), rename
    their folders to match the notebook name (filesystem-safe).
    """
    root = notebooks_root()
    if not root.exists():
        return []

    for d in list(root.iterdir()):
        if not d.is_dir():
            continue
        meta_file = d / "notebook.json"
        if not meta_file.exists():
            continue
        try:
            meta = _read_json(meta_file)
            name = str(meta.get("name") or d.name)
        except Exception:
            continue

        desired = _safe_folder_name(name)
        if desired != d.name:
            # If target already exists, choose a unique variant.
            desired = _unique_folder_name(desired)
            src = d
            dst = root / desired
            try:
                src.rename(dst)
            except Exception:
                continue

            # Update notebook.json id to match folder name.
            meta_path = dst / "notebook.json"
            try:
                raw = _read_json(meta_path)
                old_id = str(raw.get("id") or src.name)
                raw["id"] = desired
                raw["name"] = name
                raw.setdefault("created_at", _now_iso())
                _write_json(meta_path, raw)

                # Rewrite jsonl notebook_id fields (best-effort).
                _rewrite_jsonl_notebook_id(dst / "sources" / "sources.jsonl", old_id=old_id, new_id=desired)
                _rewrite_jsonl_notebook_id(dst / "chat" / "history.jsonl", old_id=old_id, new_id=desired)
            except Exception:
                pass

    return list_notebooks()


def ensure_notebook_layout(notebook_id: str) -> Path:
    nb = notebook_dir(notebook_id)
    (nb / "sources" / "raw").mkdir(parents=True, exist_ok=True)
    (nb / "sources" / "text").mkdir(parents=True, exist_ok=True)
    (nb / "chat").mkdir(parents=True, exist_ok=True)
    (nb / "quizzes").mkdir(parents=True, exist_ok=True)
    (nb / "reports").mkdir(parents=True, exist_ok=True)
    (nb / "chroma_db").mkdir(parents=True, exist_ok=True)
    return nb


def _read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _write_json(path: Path, obj: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(obj, indent=2, ensure_ascii=False), encoding="utf-8")


def _iter_jsonl(path: Path) -> Iterable[dict[str, Any]]:
    if not path.exists():
        return []
    out: list[dict[str, Any]] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            out.append(json.loads(line))
        except json.JSONDecodeError:
            continue
    return out


def _append_jsonl(path: Path, obj: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(obj, ensure_ascii=False) + "\n")


def list_notebooks() -> list[NotebookMeta]:
    root = notebooks_root()
    if not root.exists():
        return []

    notebooks: list[NotebookMeta] = []
    for d in root.iterdir():
        if not d.is_dir():
            continue
        meta_file = d / "notebook.json"
        if not meta_file.exists():
            continue
        try:
            raw = _read_json(meta_file)
            notebooks.append(
                NotebookMeta(
                    id=str(raw.get("id") or d.name),
                    name=str(raw.get("name") or d.name),
                    created_at=str(raw.get("created_at") or ""),
                )
            )
        except Exception:
            continue

    notebooks.sort(key=lambda n: n.name.lower())
    return notebooks


def create_notebook(name: str) -> NotebookMeta:
    name = (name or "").strip()
    if not name:
        raise ValueError("Notebook name cannot be empty.")

    nb_id = _unique_folder_name(name)
    ensure_notebook_layout(nb_id)
    meta = NotebookMeta(id=nb_id, name=name, created_at=_now_iso())
    _write_json(notebook_meta_path(nb_id), asdict(meta))
    return meta


def rename_notebook(notebook_id: str, new_name: str) -> NotebookMeta:
    new_name = (new_name or "").strip()
    if not new_name:
        raise ValueError("New name cannot be empty.")

    old_dir = notebook_dir(notebook_id)
    if not old_dir.exists():
        raise FileNotFoundError("Notebook does not exist.")

    new_id = _unique_folder_name(new_name) if _safe_folder_name(new_name) != notebook_id else notebook_id

    created_at = ""
    try:
        raw_existing = _read_json(notebook_meta_path(notebook_id))
        created_at = str(raw_existing.get("created_at") or "")
    except Exception:
        created_at = ""

    if new_id != notebook_id:
        new_dir = notebook_dir(new_id)
        old_dir.rename(new_dir)

        # Rewrite jsonl notebook_id fields (best-effort).
        _rewrite_jsonl_notebook_id(new_dir / "sources" / "sources.jsonl", old_id=notebook_id, new_id=new_id)
        _rewrite_jsonl_notebook_id(new_dir / "chat" / "history.jsonl", old_id=notebook_id, new_id=new_id)

    meta = NotebookMeta(id=new_id, name=new_name, created_at=created_at or _now_iso())
    _write_json(notebook_meta_path(new_id), asdict(meta))
    return meta


def delete_notebook(notebook_id: str) -> None:
    nb = notebook_dir(notebook_id)
    if nb.exists():
        shutil.rmtree(nb)


def duplicate_notebook(notebook_id: str, *, new_name: Optional[str] = None) -> NotebookMeta:
    src = notebook_dir(notebook_id)
    if not src.exists():
        raise FileNotFoundError("Notebook does not exist.")

    src_meta = _read_json(notebook_meta_path(notebook_id))
    name = (new_name or "").strip() or f"{src_meta.get('name', 'Notebook')} (copy)"

    meta = create_notebook(name)
    dst = notebook_dir(meta.id)

    # Copy user data/artifacts; keep the new notebook.json intact.
    for rel in ["sources", "chat", "quizzes", "reports", "chroma_db"]:
        s = src / rel
        if s.exists():
            shutil.copytree(s, dst / rel, dirs_exist_ok=True)

    return meta


def ensure_default_notebooks(names: list[str]) -> list[NotebookMeta]:
    migrate_notebook_folders_to_names()
    existing = list_notebooks()
    if existing:
        return existing
    created: list[NotebookMeta] = []
    for n in names:
        created.append(create_notebook(n))
    return created


def sources_manifest_path(notebook_id: str) -> Path:
    return notebook_dir(notebook_id) / "sources" / "sources.jsonl"


def append_source_event(notebook_id: str, event: dict[str, Any]) -> None:
    event = dict(event)
    event.setdefault("ts", _now_iso())
    event.setdefault("notebook_id", notebook_id)
    _append_jsonl(sources_manifest_path(notebook_id), event)


def list_source_events(notebook_id: str) -> list[dict[str, Any]]:
    return list(_iter_jsonl(sources_manifest_path(notebook_id)))


def chat_history_path(notebook_id: str) -> Path:
    return notebook_dir(notebook_id) / "chat" / "history.jsonl"


def append_chat_event(notebook_id: str, *, role: str, content: str) -> None:
    _append_jsonl(
        chat_history_path(notebook_id),
        {"ts": _now_iso(), "notebook_id": notebook_id, "role": role, "content": content},
    )