phase1/scripts/project_store.py

"""Persistent, Clerk-owned project and knowledge-document storage.

The store deliberately requires an explicit root directory. On Hugging Face this
must point at a read/write Storage Bucket mounted into the Space; silently using
the Space's ephemeral filesystem would make legal files disappear on restart.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import shutil
import threading
import time
import uuid
import zipfile
from dataclasses import dataclass
from pathlib import Path


MIB = 1024 * 1024
ALLOWED_EXTENSIONS = {".pdf", ".docx", ".txt", ".md"}
MEDIA_TYPES = {
    ".pdf": "application/pdf",
    ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    ".txt": "text/plain",
    ".md": "text/markdown",
}


class ProjectStoreError(Exception):
    status_code = 400
    code = "project_store_error"

    def __init__(self, message: str):
        super().__init__(message)
        self.message = message


class StorageUnavailable(ProjectStoreError):
    status_code = 503
    code = "project_storage_unavailable"


class NotFound(ProjectStoreError):
    status_code = 404
    code = "not_found"


class QuotaExceeded(ProjectStoreError):
    status_code = 413
    code = "knowledge_quota_exceeded"


class UnsupportedDocument(ProjectStoreError):
    status_code = 415
    code = "unsupported_document"


@dataclass(frozen=True)
class ProjectLimits:
    max_projects: int = 20
    max_documents: int = 25
    max_file_bytes: int = 10 * MIB
    max_project_bytes: int = 50 * MIB
    max_user_bytes: int = 250 * MIB

    @classmethod
    def from_env(cls) -> "ProjectLimits":
        def number(suffix: str, default: int) -> int:
            raw = os.environ.get(
                f"MOONLEY_PROJECT_{suffix}",
                os.environ.get(f"THEMIS_PROJECT_{suffix}", ""),
            ).strip()
            try:
                value = int(raw) if raw else default
            except ValueError:
                value = default
            return max(1, value)

        return cls(
            max_projects=number("MAX_PROJECTS", 20),
            max_documents=number("MAX_DOCUMENTS", 25),
            max_file_bytes=number("MAX_FILE_BYTES", 10 * MIB),
            max_project_bytes=number("MAX_BYTES", 50 * MIB),
            max_user_bytes=number("MAX_USER_BYTES", 250 * MIB),
        )

    def public_dict(self) -> dict:
        return {
            "max_projects": self.max_projects,
            "max_documents_per_project": self.max_documents,
            "max_file_bytes": self.max_file_bytes,
            "max_project_bytes": self.max_project_bytes,
            "max_user_bytes": self.max_user_bytes,
            "allowed_extensions": sorted(ALLOWED_EXTENSIONS),
        }


def _now() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def _clean_name(name: str, *, limit: int = 80) -> str:
    value = re.sub(r"\s+", " ", str(name or "")).strip()
    if not value:
        raise ProjectStoreError("A project name is required.")
    if len(value) > limit:
        raise ProjectStoreError(f"Project names must be {limit} characters or fewer.")
    return value


def _clean_filename(filename: str) -> str:
    value = Path(str(filename or "").replace("\\", "/")).name.strip()
    value = re.sub(r"[\x00-\x1f\x7f]", "", value)
    if not value or value in {".", ".."}:
        raise UnsupportedDocument("A valid document filename is required.")
    if len(value) > 180:
        stem, suffix = Path(value).stem[:150], Path(value).suffix[:20]
        value = stem + suffix
    return value


def _valid_id(value: str) -> str:
    try:
        return str(uuid.UUID(str(value)))
    except (ValueError, TypeError, AttributeError) as exc:
        raise NotFound("Project not found.") from exc


class ProjectStore:
    def __init__(self, root: str | Path | None, limits: ProjectLimits | None = None):
        self.root = Path(root).expanduser().resolve() if root else None
        self.limits = limits or ProjectLimits.from_env()
        self._lock = threading.RLock()
        self._configuration_error = ""
        if self.root:
            probe = self.root / f".moonley-write-probe-{uuid.uuid4().hex}"
            moved_probe = probe.with_suffix(".moved")
            try:
                self.root.mkdir(parents=True, exist_ok=True)
                probe.write_text("ok", encoding="utf-8")
                os.replace(probe, moved_probe)
            except OSError as exc:
                self._configuration_error = type(exc).__name__
            finally:
                for candidate in (probe, moved_probe):
                    try:
                        candidate.unlink(missing_ok=True)
                    except OSError:
                        pass

    @classmethod
    def from_env(cls) -> "ProjectStore":
        root = os.environ.get(
            "MOONLEY_PROJECT_STORAGE_ROOT",
            os.environ.get("THEMIS_PROJECT_STORAGE_ROOT", ""),
        )
        return cls(root.strip() or None)

    @property
    def configured(self) -> bool:
        return self.root is not None and not self._configuration_error

    def status(self) -> dict:
        return {
            "configured": self.configured,
            "persistent": self.configured,
            "provider": "mounted_volume" if self.configured else "unconfigured",
            "knowledge_ready": self.configured,
            "limits": self.limits.public_dict(),
            "message": (
                "Project files use the configured persistent mounted volume."
                if self.configured
                else "Attach a read/write persistent volume and set MOONLEY_PROJECT_STORAGE_ROOT."
            ),
        }

    def _require_configured(self) -> None:
        if not self.configured:
            raise StorageUnavailable(self.status()["message"])

    @staticmethod
    def _owner_key(owner_id: str) -> str:
        if not owner_id:
            raise ProjectStoreError("Authenticated user identity is required.")
        return hashlib.sha256(owner_id.encode("utf-8")).hexdigest()

    def _projects_dir(self, owner_id: str) -> Path:
        self._require_configured()
        assert self.root is not None
        return self.root / "users" / self._owner_key(owner_id) / "projects"

    def _project_dir(self, owner_id: str, project_id: str) -> Path:
        return self._projects_dir(owner_id) / _valid_id(project_id)

    @staticmethod
    def _manifest_path(project_dir: Path) -> Path:
        return project_dir / "project.json"

    def _read_manifest(self, owner_id: str, project_id: str) -> dict:
        path = self._manifest_path(self._project_dir(owner_id, project_id))
        try:
            payload = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as exc:
            raise NotFound("Project not found.") from exc
        return payload

    @staticmethod
    def _write_manifest(project_dir: Path, payload: dict) -> None:
        project_dir.mkdir(parents=True, exist_ok=True)
        temp = project_dir / f".project-{uuid.uuid4().hex}.tmp"
        temp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        os.replace(temp, project_dir / "project.json")

    @staticmethod
    def _public_document(document: dict) -> dict:
        return {
            key: value
            for key, value in document.items()
            if key not in {"stored_name", "sha256"}
        }

    @staticmethod
    def _public_project(payload: dict) -> dict:
        documents = payload.get("documents") or []
        total = sum(int(doc.get("size_bytes") or 0) for doc in documents)
        return {
            "id": payload["id"],
            "name": payload["name"],
            "created_at": payload["created_at"],
            "updated_at": payload["updated_at"],
            "document_count": len(documents),
            "knowledge_bytes": total,
            "knowledge_ready": any(
                document.get("status") in {"ready", "stored"} for document in documents
            ),
            "documents": [ProjectStore._public_document(document) for document in documents],
        }

    def list_projects(self, owner_id: str) -> list[dict]:
        with self._lock:
            directory = self._projects_dir(owner_id)
            if not directory.exists():
                return []
            projects = []
            for manifest in directory.glob("*/project.json"):
                try:
                    projects.append(self._public_project(json.loads(manifest.read_text(encoding="utf-8"))))
                except (OSError, KeyError, json.JSONDecodeError):
                    continue
            return sorted(projects, key=lambda item: item["updated_at"], reverse=True)

    def create_project(self, owner_id: str, name: str) -> dict:
        with self._lock:
            projects = self.list_projects(owner_id)
            if len(projects) >= self.limits.max_projects:
                raise QuotaExceeded(f"A user can have at most {self.limits.max_projects} projects.")
            project_id = str(uuid.uuid4())
            timestamp = _now()
            payload = {
                "version": 1,
                "id": project_id,
                "name": _clean_name(name),
                "created_at": timestamp,
                "updated_at": timestamp,
                "documents": [],
            }
            self._write_manifest(self._project_dir(owner_id, project_id), payload)
            return self._public_project(payload)

    def get_project(self, owner_id: str, project_id: str) -> dict:
        with self._lock:
            return self._public_project(self._read_manifest(owner_id, project_id))

    def rename_project(self, owner_id: str, project_id: str, name: str) -> dict:
        with self._lock:
            payload = self._read_manifest(owner_id, project_id)
            payload["name"] = _clean_name(name)
            payload["updated_at"] = _now()
            self._write_manifest(self._project_dir(owner_id, project_id), payload)
            return self._public_project(payload)

    def delete_project(self, owner_id: str, project_id: str) -> None:
        with self._lock:
            project_dir = self._project_dir(owner_id, project_id)
            if not self._manifest_path(project_dir).exists():
                raise NotFound("Project not found.")
            shutil.rmtree(project_dir)

    @staticmethod
    def _validate_content(extension: str, content: bytes) -> None:
        if not content:
            raise UnsupportedDocument("Empty documents cannot be uploaded.")
        if extension == ".pdf" and not content.startswith(b"%PDF-"):
            raise UnsupportedDocument("The file does not contain a valid PDF header.")
        if extension == ".docx":
            try:
                from io import BytesIO

                with zipfile.ZipFile(BytesIO(content)) as archive:
                    names = set(archive.namelist())
                    if "[Content_Types].xml" not in names or "word/document.xml" not in names:
                        raise UnsupportedDocument("The file is not a valid DOCX document.")
            except zipfile.BadZipFile as exc:
                raise UnsupportedDocument("The file is not a valid DOCX document.") from exc
        if extension in {".txt", ".md"}:
            try:
                content.decode("utf-8-sig")
            except UnicodeDecodeError as exc:
                raise UnsupportedDocument("Text and Markdown documents must use UTF-8 encoding.") from exc

    def _user_bytes(self, owner_id: str) -> int:
        return sum(project["knowledge_bytes"] for project in self.list_projects(owner_id))

    def add_document(self, owner_id: str, project_id: str, filename: str, content: bytes) -> dict:
        with self._lock:
            safe_name = _clean_filename(filename)
            extension = Path(safe_name).suffix.lower()
            if extension not in ALLOWED_EXTENSIONS:
                raise UnsupportedDocument("Allowed file types are PDF, DOCX, TXT, and Markdown.")
            size = len(content)
            if size > self.limits.max_file_bytes:
                raise QuotaExceeded(f"Each document must be {self.limits.max_file_bytes // MIB} MiB or smaller.")
            self._validate_content(extension, content)

            payload = self._read_manifest(owner_id, project_id)
            documents = payload.get("documents") or []
            digest = hashlib.sha256(content).hexdigest()
            duplicate = next((doc for doc in documents if doc.get("sha256") == digest), None)
            if duplicate:
                return self._public_document(duplicate)
            if len(documents) >= self.limits.max_documents:
                raise QuotaExceeded(f"A project can contain at most {self.limits.max_documents} documents.")
            project_bytes = sum(int(doc.get("size_bytes") or 0) for doc in documents)
            if project_bytes + size > self.limits.max_project_bytes:
                raise QuotaExceeded(f"Project knowledge is limited to {self.limits.max_project_bytes // MIB} MiB.")
            if self._user_bytes(owner_id) + size > self.limits.max_user_bytes:
                raise QuotaExceeded(f"User knowledge storage is limited to {self.limits.max_user_bytes // MIB} MiB.")

            document_id = str(uuid.uuid4())
            project_dir = self._project_dir(owner_id, project_id)
            document_dir = project_dir / "documents"
            document_dir.mkdir(parents=True, exist_ok=True)
            stored_name = document_id + extension
            temp = document_dir / f".{document_id}.tmp"
            temp.write_bytes(content)
            os.replace(temp, document_dir / stored_name)
            document = {
                "id": document_id,
                "name": safe_name,
                "size_bytes": size,
                "media_type": MEDIA_TYPES[extension],
                "sha256": digest,
                "status": "stored",
                "created_at": _now(),
                "stored_name": stored_name,
            }
            documents.append(document)
            payload["documents"] = documents
            payload["updated_at"] = document["created_at"]
            self._write_manifest(project_dir, payload)
            return self._public_document(document)

    def delete_document(self, owner_id: str, project_id: str, document_id: str) -> None:
        with self._lock:
            try:
                normalized_id = str(uuid.UUID(str(document_id)))
            except (ValueError, TypeError, AttributeError) as exc:
                raise NotFound("Document not found.") from exc
            payload = self._read_manifest(owner_id, project_id)
            documents = payload.get("documents") or []
            document = next((item for item in documents if item.get("id") == normalized_id), None)
            if not document:
                raise NotFound("Document not found.")
            project_dir = self._project_dir(owner_id, project_id)
            (project_dir / "documents" / document["stored_name"]).unlink(missing_ok=True)
            knowledge_dir = project_dir / "knowledge" / normalized_id
            if knowledge_dir.exists():
                shutil.rmtree(knowledge_dir)
            payload["documents"] = [item for item in documents if item.get("id") != normalized_id]
            payload["updated_at"] = _now()
            self._write_manifest(project_dir, payload)

    def document_record(self, owner_id: str, project_id: str, document_id: str) -> dict:
        """Return one private manifest record after owner/project validation."""
        with self._lock:
            normalized_id = _valid_id(document_id)
            payload = self._read_manifest(owner_id, project_id)
            document = next(
                (item for item in (payload.get("documents") or []) if item.get("id") == normalized_id),
                None,
            )
            if not document:
                raise NotFound("Document not found.")
            return dict(document)

    def document_path(self, owner_id: str, project_id: str, document_id: str) -> Path:
        document = self.document_record(owner_id, project_id, document_id)
        documents_dir = (self._project_dir(owner_id, project_id) / "documents").resolve()
        path = (documents_dir / str(document["stored_name"])).resolve()
        if path.parent != documents_dir or not path.is_file():
            raise NotFound("Document file not found.")
        return path

    def knowledge_dir(self, owner_id: str, project_id: str, document_id: str) -> Path:
        self.document_record(owner_id, project_id, document_id)
        path = self._project_dir(owner_id, project_id) / "knowledge" / _valid_id(document_id)
        path.mkdir(parents=True, exist_ok=True)
        return path

    def record_extraction(
        self,
        owner_id: str,
        project_id: str,
        document_id: str,
        extraction: dict,
        *,
        status: str = "ready",
    ) -> dict:
        with self._lock:
            normalized_id = _valid_id(document_id)
            payload = self._read_manifest(owner_id, project_id)
            document = next(
                (item for item in (payload.get("documents") or []) if item.get("id") == normalized_id),
                None,
            )
            if not document:
                raise NotFound("Document not found.")
            document["status"] = status
            document["extraction"] = {
                key: value
                for key, value in extraction.items()
                if key in {"method", "pages", "ocr_pages", "truncated", "text_chars", "chunk_count", "vector_provider", "source_provider"}
            }
            payload["updated_at"] = _now()
            self._write_manifest(self._project_dir(owner_id, project_id), payload)
            return self._public_document(document)