kamidso-back/kamidso/storage/layout.py

from __future__ import annotations

import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Optional, Tuple


SUBDIRS = (
    "uploads",
    "runtime",
    "generated",
    "temp",
    "knowledge",
    "vectors",
    "cache",
    "engine-home",
    "logs",
    "meta",
)


def _app_root() -> Path:
    from kamidso.settings.paths import repo_root

    return repo_root()


def resolve_data_path(path: str | os.PathLike[str]) -> Path:
    """Resolve a storage path; relative paths are anchored at the app root."""
    root = Path(path).expanduser()
    if not root.is_absolute():
        root = _app_root() / root
    return root.resolve()


def _default_data_root() -> Path:
    env = os.environ.get("KAMIDSO_STORAGE_PATH") or os.environ.get("KAMIDSO_DATA_ROOT")
    if env:
        return resolve_data_path(env)
    return resolve_data_path("data")


@dataclass(frozen=True)
class DataLayout:
    """Canonical on-disk layout under a single ``data/`` root."""

    root: Path

    @classmethod
    def from_base_path(cls, base_path: str | os.PathLike[str]) -> "DataLayout":
        return cls(root=resolve_data_path(base_path))

    @classmethod
    def current(cls) -> "DataLayout":
        env = os.environ.get("KAMIDSO_DATA_ROOT")
        if env:
            return cls(root=resolve_data_path(env))
        return cls(root=_default_data_root())

    @property
    def uploads(self) -> Path:
        return self.root / "uploads"

    @property
    def runtime(self) -> Path:
        return self.root / "runtime"

    @property
    def generated(self) -> Path:
        return self.root / "generated"

    @property
    def temp(self) -> Path:
        return self.root / "temp"

    @property
    def knowledge(self) -> Path:
        return self.root / "knowledge"

    @property
    def vectors(self) -> Path:
        return self.root / "vectors" / "chromadb"

    @property
    def cache(self) -> Path:
        return self.root / "cache"

    @property
    def engine_home(self) -> Path:
        return self.root / "engine-home"

    @property
    def logs(self) -> Path:
        return self.root / "logs"

    @property
    def meta(self) -> Path:
        return self.root / "meta"

    def iter_dirs(self) -> Iterable[Path]:
        return (
            self.uploads,
            self.runtime,
            self.generated,
            self.temp,
            self.knowledge,
            self.vectors,
            self.cache,
            self.engine_home,
            self.logs,
            self.meta,
            self.engine_home / "workspace",
            self.runtime / "op_snapshots",
        )

    def ensure(self) -> None:
        self.root.mkdir(parents=True, exist_ok=True)
        for path in self.iter_dirs():
            path.mkdir(parents=True, exist_ok=True)

    def bind_process_env(self, *, skills_dir: Optional[Path] = None) -> None:
        """Point the in-process engine at this data tree before it snapshots paths."""
        root = str(self.root)
        os.environ["KAMIDSO_DATA_ROOT"] = root
        os.environ["KAMIDSO_STORAGE_PATH"] = root
        os.environ["KAMIDSO_VECTORS_PATH"] = str(self.vectors)
        os.environ["KAMIDSO_ENGINE_HOME"] = str(self.engine_home)
        os.environ["KAMIDSO_ENGINE_LOG_DIR"] = str(self.logs)
        os.environ["DBGPT_HOME"] = str(self.engine_home)
        os.environ["FILE_SERVER_LOCAL_STORAGE_PATH"] = str(self.cache / "files")
        os.environ["DBGPT_LOG_DIR"] = str(self.logs)
        os.environ["CHROMA_PERSIST_PATH"] = str(self.vectors)
        os.environ["TMPDIR"] = str(self.temp)
        os.environ["TEMP"] = str(self.temp)
        os.environ["TMP"] = str(self.temp)
        os.environ["XDG_CACHE_HOME"] = str(self.cache)
        huggingface = self.cache / "huggingface"
        os.environ["HF_HOME"] = str(huggingface)
        os.environ["HUGGINGFACE_HUB_CACHE"] = str(huggingface / "hub")
        os.environ["TRANSFORMERS_CACHE"] = str(huggingface / "transformers")
        # Connector rows live in the engine metadata DB. Keep that file under
        # data/meta so a process restart does not drop saved connections.
        os.environ.pop("KAMIDSO_SKIP_SQLITE_FILE", None)
        os.environ["DBGPT_OBSERVABILITY_SQLITE_PATH"] = str(
            self.logs / "observability.db"
        )
        os.environ["KAMIDSO_OBSERVABILITY_JSONL"] = str(
            self.logs / "observability.jsonl"
        )
        from kamidso.engine.core.configs.model_config import apply_storage_paths

        apply_storage_paths(str(self.root))
        if skills_dir is not None:
            os.environ.setdefault("KAMIDSO_SKILLS_DIR", str(skills_dir))
            os.environ.setdefault("DBGPT_SKILLS_DIR", str(skills_dir))


def upload_storage_key(workspace_id: str, file_id: str, filename: str) -> str:
    return f"uploads/{workspace_id}/{file_id}/{filename}"


_LEGACY_WS_FILES = re.compile(
    r"^(?:storage/)?(?P<ws>[^/]+)/files/(?P<file_id>[^/]+)/(?P<name>.+)$"
)
_UPLOADS_KEY = re.compile(
    r"^uploads/(?P<ws>[^/]+)/(?P<file_id>[^/]+)/(?P<name>.+)$"
)


def candidate_storage_keys(key: str) -> Tuple[str, ...]:
    """Return lookup order for a stored key, including pre-unification layouts.

    Mongo records created when the base was ``data/storage`` keep keys like
    ``ws_x/files/file_y/name.xlsx``. New writes use ``uploads/ws_x/file_y/name``.
    """
    normalized = key.replace("\\", "/").lstrip("/")
    keys = [normalized]
    match = _UPLOADS_KEY.match(normalized) or _LEGACY_WS_FILES.match(normalized)
    if match:
        workspace_id = match.group("ws")
        file_id = match.group("file_id")
        name = match.group("name")
        keys.extend(
            [
                f"uploads/{workspace_id}/{file_id}/{name}",
                f"storage/{workspace_id}/files/{file_id}/{name}",
                f"{workspace_id}/files/{file_id}/{name}",
            ]
        )
    seen: set[str] = set()
    unique: list[str] = []
    for item in keys:
        if item and item not in seen:
            seen.add(item)
            unique.append(item)
    return tuple(unique)