ArcaSuite/services/mcp-docugen/src/mcp_docugen/template_store.py

from __future__ import annotations

import base64
import re
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import aiofiles
import yaml
from pydantic import ValidationError

from mcp_docugen.models import TemplateAsset, TemplateFrontmatter, TemplateSummary

_SLUG_RE = re.compile(r"^[a-z0-9][a-z0-9-]*$")
_ASSET_FILENAME_RE = re.compile(r"^[A-Za-z0-9._-]+$")
_FRONTMATTER_DELIM = "---"


class TemplateNotFound(Exception):
    pass


class TemplateAlreadyExists(Exception):
    pass


class InvalidFrontmatter(Exception):
    pass


class InvalidTemplateName(ValueError):
    pass


@dataclass
class LoadedTemplate:
    frontmatter: TemplateFrontmatter
    body: str
    assets: list[TemplateAsset]
    updated_at: datetime


class TemplateStore:
    def __init__(self, base_dir: Path) -> None:
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(parents=True, exist_ok=True)

    async def create(
        self,
        name: str,
        frontmatter: TemplateFrontmatter,
        body: str,
        assets: list[dict] | None = None,
    ) -> None:
        self._validate_name(name)
        tdir = self.base_dir / name
        if tdir.exists():
            raise TemplateAlreadyExists(name)
        tdir.mkdir(parents=True)
        (tdir / "assets").mkdir()
        if assets:
            self._validate_assets(assets)
        await self._write_template_file(tdir, frontmatter, body)
        if assets:
            await self._write_assets(tdir, assets)

    async def update(
        self,
        name: str,
        frontmatter: TemplateFrontmatter,
        body: str,
        assets: list[dict] | None = None,
    ) -> None:
        self._validate_name(name)
        tdir = self.base_dir / name
        if not tdir.exists():
            raise TemplateNotFound(name)
        await self._write_template_file(tdir, frontmatter, body)
        if assets is not None:
            assets_dir = tdir / "assets"
            for f in assets_dir.iterdir():
                f.unlink()
            await self._write_assets(tdir, assets)

    async def delete(self, name: str) -> None:
        self._validate_name(name)
        tdir = self.base_dir / name
        if not tdir.exists():
            raise TemplateNotFound(name)
        for f in (tdir / "assets").iterdir():
            f.unlink()
        (tdir / "assets").rmdir()
        (tdir / "template.md").unlink()
        tdir.rmdir()

    async def get(self, name: str) -> LoadedTemplate:
        self._validate_name(name)
        tdir = self.base_dir / name
        tpath = tdir / "template.md"
        if not tpath.exists():
            raise TemplateNotFound(name)
        frontmatter, body = await self._read_template_file(tpath)
        assets = self._list_assets(tdir)
        stat = tpath.stat()
        updated_at = datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc)
        return LoadedTemplate(
            frontmatter=frontmatter, body=body, assets=assets, updated_at=updated_at
        )

    async def list(self) -> list[TemplateSummary]:
        summaries = []
        for tdir in sorted(self.base_dir.iterdir()):
            if not tdir.is_dir():
                continue
            tpath = tdir / "template.md"
            if not tpath.exists():
                continue
            try:
                frontmatter, _ = await self._read_template_file(tpath)
            except InvalidFrontmatter:
                continue
            stat = tpath.stat()
            summaries.append(
                TemplateSummary(
                    name=frontmatter.name,
                    description=frontmatter.description,
                    updated_at=datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc),
                )
            )
        return summaries

    async def read_asset(self, template_name: str, filename: str) -> bytes:
        self._validate_name(template_name)
        self._validate_asset_filename(filename)
        path = self.base_dir / template_name / "assets" / filename
        if not path.exists():
            raise FileNotFoundError(filename)
        async with aiofiles.open(path, "rb") as f:
            return await f.read()

    def asset_path(self, template_name: str, filename: str) -> Path:
        self._validate_name(template_name)
        self._validate_asset_filename(filename)
        return self.base_dir / template_name / "assets" / filename

    def _validate_name(self, name: str) -> None:
        if not _SLUG_RE.match(name):
            raise InvalidTemplateName(name)

    def _validate_asset_filename(self, filename: str) -> None:
        if not _ASSET_FILENAME_RE.match(filename):
            raise ValueError(f"invalid asset filename: {filename!r}")

    def _validate_assets(self, assets: list[dict]) -> None:
        for asset in assets:
            self._validate_asset_filename(asset["filename"])

    async def _write_template_file(
        self, tdir: Path, frontmatter: TemplateFrontmatter, body: str
    ) -> None:
        yaml_text = yaml.safe_dump(
            frontmatter.model_dump(exclude_none=True), sort_keys=False
        )
        content = f"{_FRONTMATTER_DELIM}\n{yaml_text}{_FRONTMATTER_DELIM}\n{body}"
        async with aiofiles.open(tdir / "template.md", "w") as f:
            await f.write(content)

    async def _read_template_file(
        self, path: Path
    ) -> tuple[TemplateFrontmatter, str]:
        async with aiofiles.open(path) as f:
            raw = await f.read()
        if not raw.startswith(_FRONTMATTER_DELIM):
            raise InvalidFrontmatter("missing opening '---'")
        parts = raw.split(_FRONTMATTER_DELIM, 2)
        if len(parts) < 3:
            raise InvalidFrontmatter("missing closing '---'")
        yaml_text = parts[1]
        body = parts[2].lstrip("\n")
        try:
            data = yaml.safe_load(yaml_text) or {}
            frontmatter = TemplateFrontmatter(**data)
        except (yaml.YAMLError, ValidationError) as exc:
            raise InvalidFrontmatter(str(exc)) from exc
        return frontmatter, body

    async def _write_assets(self, tdir: Path, assets: list[dict]) -> None:
        for asset in assets:
            filename = asset["filename"]
            self._validate_asset_filename(filename)
            data = base64.b64decode(asset["data_b64"])
            async with aiofiles.open(tdir / "assets" / filename, "wb") as f:
                await f.write(data)

    def _list_assets(self, tdir: Path) -> list[TemplateAsset]:
        assets_dir = tdir / "assets"
        if not assets_dir.exists():
            return []
        out = []
        for f in sorted(assets_dir.iterdir()):
            if not f.is_file():
                continue
            out.append(
                TemplateAsset(
                    filename=f.name,
                    mime=_guess_mime(f.name),
                    size_bytes=f.stat().st_size,
                )
            )
        return out


def _guess_mime(filename: str) -> str:
    ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    return {
        "png": "image/png",
        "jpg": "image/jpeg",
        "jpeg": "image/jpeg",
        "webp": "image/webp",
    }.get(ext, "application/octet-stream")