ForcePilot/backend/package/yuxi/channel/message/media_store.py

from __future__ import annotations

import asyncio
import json
import logging
import os
import uuid
from datetime import UTC, datetime
from pathlib import Path

logger = logging.getLogger(__name__)

_META_SUFFIX = ".meta.json"


def _now_iso() -> str:
    return datetime.now(UTC).isoformat()


def _date_dir_name() -> str:
    return datetime.now(UTC).strftime("%Y-%m-%d")


class MediaStore:
    def __init__(
        self,
        base_dir: Path,
        max_total_bytes: int,
        max_file_bytes: int,
        ttl_seconds: int,
    ):
        self.base_dir = base_dir
        self.max_total_bytes = max_total_bytes
        self.max_file_bytes = max_file_bytes
        self.ttl_seconds = ttl_seconds
        self._size_initialized = False
        self._running_total: int = 0
        self._size_lock = asyncio.Lock()

    def _ensure_dir(self, dir_path: Path) -> None:
        dir_path.mkdir(parents=True, exist_ok=True)

    def _make_file_id(self, date_dir: str, uid: str, ext: str) -> str:
        return f"{date_dir}/{uid}{ext}"

    def _get_file_path(self, file_id: str) -> Path:
        return self.base_dir / file_id

    def _get_meta_path(self, file_id: str) -> Path:
        return self.base_dir / f"{file_id}{_META_SUFFIX}"

    def _dir_size(self, dir_path: Path) -> int:
        total = 0
        try:
            for entry in dir_path.rglob("*"):
                if entry.is_file():
                    try:
                        total += entry.stat().st_size
                    except OSError:
                        logger.debug("Failed to stat file during size calculation: %s", entry)
        except OSError:
            logger.debug("Failed to calculate directory size: %s", dir_path)
        return total

    async def _check_quota(self, incoming_bytes: int) -> bool:
        if incoming_bytes > self.max_file_bytes:
            logger.warning(
                "File size %d exceeds per-file limit %d",
                incoming_bytes,
                self.max_file_bytes,
            )
            return False
        current_total = await self._get_running_total()
        if current_total + incoming_bytes > self.max_total_bytes:
            logger.warning(
                "Storage quota exceeded: current=%d, incoming=%d, max=%d",
                current_total,
                incoming_bytes,
                self.max_total_bytes,
            )
            return False
        return True

    async def _get_running_total(self) -> int:
        async with self._size_lock:
            if not self._size_initialized:
                self._running_total = self._dir_size(self.base_dir)
                self._size_initialized = True
            return self._running_total

    def _infer_ext(self, content_type: str | None, filename: str | None) -> str:
        if filename:
            suffix = Path(filename).suffix
            if suffix:
                return suffix.lower()
        if content_type:
            ct = content_type.split(";")[0].strip().lower()
            mapping = {
                "image/jpeg": ".jpg",
                "image/png": ".png",
                "image/gif": ".gif",
                "image/webp": ".webp",
                "image/bmp": ".bmp",
                "image/svg+xml": ".svg",
                "audio/mpeg": ".mp3",
                "audio/wav": ".wav",
                "audio/ogg": ".ogg",
                "video/mp4": ".mp4",
                "application/pdf": ".pdf",
            }
            if ct in mapping:
                return mapping[ct]
        return ".bin"

    async def store(
        self,
        data: bytes,
        content_type: str | None = None,
        filename: str | None = None,
        source_url: str | None = None,
    ) -> str | None:
        if not await self._check_quota(len(data)):
            return None

        date_dir = _date_dir_name()
        uid = uuid.uuid4().hex
        ext = self._infer_ext(content_type, filename)
        file_id = self._make_file_id(date_dir, uid, ext)

        file_path = self._get_file_path(file_id)
        meta_path = self._get_meta_path(file_id)

        self._ensure_dir(file_path.parent)

        try:
            file_path.write_bytes(data)
        except OSError:
            logger.exception("Failed to write media file: %s", file_path)
            return None

        meta = {
            "original_filename": filename,
            "content_type": content_type,
            "source_url": source_url,
            "stored_at": _now_iso(),
            "ttl_seconds": self.ttl_seconds,
            "size": len(data),
        }
        try:
            meta_path.write_text(json.dumps(meta, ensure_ascii=False), encoding="utf-8")
        except OSError:
            logger.exception("Failed to write media metadata: %s", meta_path)
            try:
                file_path.unlink(missing_ok=True)
            except OSError:
                logger.debug("Failed to clean up orphaned media file: %s", file_path)
            return None

        logger.info("Stored media: file_id=%s, size=%d", file_id, len(data))
        async with self._size_lock:
            self._running_total += len(data)
        return file_id

    def resolve_path(self, file_id: str) -> Path | None:
        file_path = self._get_file_path(file_id)
        if file_path.exists():
            return file_path
        return None

    def read_meta(self, file_id: str) -> dict | None:
        meta_path = self._get_meta_path(file_id)
        try:
            return json.loads(meta_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            return None

    async def delete(self, file_id: str) -> bool:
        file_path = self._get_file_path(file_id)
        meta_path = self._get_meta_path(file_id)
        deleted = False
        for p in (file_path, meta_path):
            if not p.exists():
                continue
            try:
                p.unlink()
                deleted = True
            except OSError:
                logger.debug("Failed to delete file during media cleanup: %s", p)
        if deleted:
            logger.info("Deleted media: file_id=%s", file_id)
        return deleted

    def get_total_size(self) -> int:
        return self._dir_size(self.base_dir)

    async def cleanup_expired(self) -> int:
        deleted_count = 0
        now = datetime.now(UTC)
        try:
            for entry in sorted(self.base_dir.rglob(f"*{_META_SUFFIX}")):
                try:
                    meta = json.loads(entry.read_text(encoding="utf-8"))
                except (OSError, json.JSONDecodeError):
                    continue

                stored_at_str = meta.get("stored_at")
                ttl_s = meta.get("ttl_seconds", self.ttl_seconds)
                if not stored_at_str:
                    continue

                try:
                    stored_at = datetime.fromisoformat(stored_at_str)
                except ValueError:
                    continue

                if (now - stored_at).total_seconds() <= ttl_s:
                    continue

                file_id = str(entry.relative_to(self.base_dir)).removesuffix(_META_SUFFIX)
                if await self.delete(file_id):
                    deleted_count += 1

            self._remove_empty_dirs()
        except OSError:
            logger.exception("Error during media cleanup scan")

        if deleted_count:
            logger.info("Media cleanup: deleted %d expired files", deleted_count)
        return deleted_count

    def _remove_empty_dirs(self) -> None:
        try:
            for dirpath, dirnames, filenames in os.walk(self.base_dir, topdown=False):
                if dirpath == str(self.base_dir):
                    continue
                if not dirnames and not filenames:
                    try:
                        os.rmdir(dirpath)
                    except OSError:
                        logger.debug("Failed to remove empty directory: %s", dirpath)
        except OSError:
            logger.debug("Failed to walk directory for cleanup: %s", self.base_dir)