DocHub/backend/app/document_service.py

import base64
import logging
import re
from datetime import datetime, timezone

from app.text_extract import html_to_text, markdown_to_text

logger = logging.getLogger(__name__)


def _natural_sort_key(name: str) -> list:
    return [int(part) if part.isdigit() else part for part in re.split(r"(\d+)", name)]

ALLOWED_UPLOAD_EXTENSIONS = {
    ".md", ".html", ".htm",
    ".js", ".ts", ".css",
    ".json", ".yaml", ".yml", ".toml", ".xml", ".csv", ".txt",
    ".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()


_IMAGE_CONTENT_TYPES = {
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".gif": "image/gif",
    ".webp": "image/webp",
}


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"
    extension = _extension_of(key)
    if extension in _IMAGE_CONTENT_TYPES:
        return _IMAGE_CONTENT_TYPES[extension]
    return "application/octet-stream"


class DocumentService:
    def __init__(self, file_client, search_client):
        self._file_client = file_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._file_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"],
                })
        folders.sort(key=_natural_sort_key)
        files.sort(key=lambda f: _natural_sort_key(f["name"]))
        return {"folders": folders, "files": files}

    def get_document(self, key: str) -> dict:
        _validate_path(key)
        try:
            data = self._file_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._file_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._file_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)}

    def reindex_all(self) -> dict:
        self._search_client.delete_all_documents()
        keys = self._file_client.walk_files()
        indexed = 0
        failed = 0
        for key in keys:
            content_type = _guess_content_type(key)
            if content_type not in ("text/markdown", "text/html"):
                continue
            try:
                data = self._file_client.get_object(key)
                text = self._extract_text(data, 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)
                indexed += 1
            except Exception:
                logger.exception("Failed to reindex %s", key)
                failed += 1
        return {"indexed": indexed, "failed": failed, "total": len(keys)}