DocHub/backend/app/document_service.py

import base64
import logging
from datetime import datetime, timezone

from app.text_extract import html_to_text, markdown_to_text

logger = logging.getLogger(__name__)

ALLOWED_UPLOAD_EXTENSIONS = {".md", ".html", ".htm", ".png", ".jpg", ".jpeg", ".gif", ".webp"}
MAX_UPLOAD_SIZE_BYTES = 10 * 1024 * 1024


class DocumentNotFoundError(Exception):
    pass


class UnsupportedFileTypeError(Exception):
    def __init__(self, extension: str):
        self.extension = extension
        super().__init__(f"unsupported file type: {extension}")


class FileTooLargeError(Exception):
    def __init__(self, max_size_bytes: int):
        self.max_size_bytes = max_size_bytes
        super().__init__(f"file too large: max {max_size_bytes // (1024 * 1024)}MB")


class InvalidPathError(Exception):
    def __init__(self, path: str):
        self.path = path
        super().__init__(f"invalid path: {path}")


def _validate_path(path: str) -> None:
    if path.startswith("/") or any(segment == ".." for segment in path.split("/")):
        raise InvalidPathError(path)


def _extension_of(key: str) -> str:
    name = key.rsplit("/", 1)[-1]
    if "." not in name:
        return ""
    return "." + name.rsplit(".", 1)[-1].lower()


def _guess_content_type(key: str) -> str:
    if key.endswith(".md"):
        return "text/markdown"
    if key.endswith(".html") or key.endswith(".htm"):
        return "text/html"
    return "application/octet-stream"


class DocumentService:
    def __init__(self, garage_client, search_client):
        self._garage_client = garage_client
        self._search_client = search_client

    def list_documents(self, path: str) -> dict:
        prefix = path if not path or path.endswith("/") else f"{path}/"
        entries = self._garage_client.list_objects(prefix)
        folders = []
        files = []
        for entry in entries:
            key = entry["key"]
            if entry.get("is_dir"):
                name = key[len(prefix):].rstrip("/")
                if name:
                    folders.append(name)
            else:
                name = key[len(prefix):]
                files.append({
                    "key": key,
                    "name": name,
                    "size": entry["size"],
                    "modifiedAt": entry["modifiedAt"],
                })
        return {"folders": folders, "files": files}

    def get_document(self, key: str) -> dict:
        _validate_path(key)
        try:
            data = self._garage_client.get_object(key)
        except FileNotFoundError as exc:
            raise DocumentNotFoundError(key) from exc
        try:
            content = data.decode("utf-8")
            encoding = "utf-8"
        except UnicodeDecodeError:
            content = base64.b64encode(data).decode("ascii")
            encoding = "base64"
        return {
            "key": key,
            "content": content,
            "contentType": _guess_content_type(key),
            "encoding": encoding,
        }

    def _extract_text(self, content: bytes, content_type: str) -> str:
        text = content.decode("utf-8", errors="ignore")
        if content_type == "text/html":
            return html_to_text(text)
        if content_type == "text/markdown":
            return markdown_to_text(text)
        return text

    def save_document(self, key: str, content: bytes, content_type: str) -> dict:
        _validate_path(key)
        self._garage_client.put_object(key, content, content_type)
        try:
            text = self._extract_text(content, content_type)
            title = key.rsplit("/", 1)[-1]
            path = key.rsplit("/", 1)[0] if "/" in key else ""
            self._search_client.index_document(key, path, title, text)
        except Exception:
            logger.exception("Failed to update search index for %s", key)
        return {"key": key, "updatedAt": datetime.now(timezone.utc).isoformat()}

    def delete_document(self, key: str) -> None:
        _validate_path(key)
        self._garage_client.delete_object(key)
        try:
            self._search_client.delete_document(key)
        except Exception:
            logger.exception("Failed to delete search index entry for %s", key)

    def upload_document(self, key: str, data: bytes, content_type: str) -> dict:
        _validate_path(key)
        extension = _extension_of(key)
        if extension not in ALLOWED_UPLOAD_EXTENSIONS:
            raise UnsupportedFileTypeError(extension)
        if len(data) > MAX_UPLOAD_SIZE_BYTES:
            raise FileTooLargeError(MAX_UPLOAD_SIZE_BYTES)
        self.save_document(key, data, content_type)
        name = key.rsplit("/", 1)[-1]
        return {"key": key, "name": name, "size": len(data)}