ARIA-AGENT/aria-brain/memory_attachments.py

"""
Anhaenge fuer Memory-Eintraege.

Storage-Layout:
    /shared/memory-attachments/<memory-id>/<original-name>

Eine flache Ordnerstruktur pro Memory — bei Memory-Delete loescht main.py
das ganze Verzeichnis. Anhang-Metadaten (name, mime, size, path) liegen
zusaetzlich im Qdrant-Payload des Memory-Punkts damit die Listen/Suche
sie ohne Filesystem-Lookup zeigen kann.

Anhaenge sind erstmal nur ueber die Diagnostic-UI hochladbar — ARIA
selbst hat in Stufe A kein Tool zum Upload.
"""

from __future__ import annotations

import base64
import logging
import mimetypes
import os
import re
import shutil
from pathlib import Path
from typing import List, Optional

logger = logging.getLogger(__name__)

ROOT = Path(os.environ.get("MEMORY_ATTACHMENTS_DIR", "/shared/memory-attachments"))
MAX_BYTES = int(os.environ.get("MEMORY_ATTACHMENT_MAX_BYTES", str(20 * 1024 * 1024)))  # 20 MB
SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._\-]")


def _safe_filename(name: str) -> str:
    """Macht aus einem User-Namen einen filesystem-sicheren String —
    zerlegt Pfadteile, schneidet Sonderzeichen weg, kuerzt auf 120 Zeichen."""
    base = Path(name).name or "datei"
    base = SAFE_NAME_RE.sub("_", base).strip("._-") or "datei"
    return base[:120]


def memory_dir(memory_id: str) -> Path:
    return ROOT / memory_id


def list_attachments(memory_id: str) -> List[dict]:
    """Liest die Anhaenge fuer eine Memory aus dem Filesystem.
    Returns [{name, mime, size, path}, ...] — leer wenn nichts da.
    Source of Truth ist Qdrant-Payload; diese Funktion ist nur fuer
    Diagnostic-Endpoints wenn Stefan direkt das FS prueft."""
    d = memory_dir(memory_id)
    if not d.is_dir():
        return []
    out = []
    for f in sorted(d.iterdir()):
        if not f.is_file():
            continue
        out.append(_file_meta(memory_id, f))
    return out


def _file_meta(memory_id: str, f: Path) -> dict:
    try:
        size = f.stat().st_size
    except Exception:
        size = 0
    mime = mimetypes.guess_type(f.name)[0] or "application/octet-stream"
    return {
        "name": f.name,
        "mime": mime,
        "size": size,
        "path": str(f),  # absoluter Pfad im Container
    }


def save_attachment(memory_id: str, filename: str, data: bytes) -> dict:
    """Schreibt einen Anhang ins FS und gibt seine Metadaten zurueck.
    Ueberschreibt eine bestehende Datei mit gleichem Namen."""
    if not memory_id:
        raise ValueError("memory_id ist Pflicht")
    if len(data) > MAX_BYTES:
        raise ValueError(f"Anhang zu gross ({len(data)} > {MAX_BYTES} Byte)")
    safe = _safe_filename(filename)
    d = memory_dir(memory_id)
    d.mkdir(parents=True, exist_ok=True)
    target = d / safe
    target.write_bytes(data)
    logger.info("[mem-att] %s -> %s (%d Byte)", memory_id, safe, len(data))
    return _file_meta(memory_id, target)


def save_from_base64(memory_id: str, filename: str, b64: str) -> dict:
    """Convenience fuer Base64-Uploads (Diagnostic schickt Files so)."""
    try:
        data = base64.b64decode(b64, validate=False)
    except Exception as exc:
        raise ValueError(f"Base64-Decode fehlgeschlagen: {exc}") from exc
    return save_attachment(memory_id, filename, data)


def delete_attachment(memory_id: str, filename: str) -> bool:
    """Loescht eine einzelne Anhang-Datei. Returns True wenn was weg ist."""
    safe = _safe_filename(filename)
    target = memory_dir(memory_id) / safe
    if not target.is_file():
        return False
    try:
        target.unlink()
        logger.info("[mem-att] %s/%s geloescht", memory_id, safe)
        return True
    except Exception as exc:
        logger.warning("[mem-att] Loeschen fehlgeschlagen: %s", exc)
        return False


def delete_all(memory_id: str) -> int:
    """Loescht das komplette Memory-Verzeichnis. Wird beim Memory-Delete
    in main.py gerufen damit nichts verwaist."""
    d = memory_dir(memory_id)
    if not d.is_dir():
        return 0
    count = sum(1 for _ in d.iterdir() if _.is_file())
    try:
        shutil.rmtree(d)
        logger.info("[mem-att] %s komplett entfernt (%d Files)", memory_id, count)
    except Exception as exc:
        logger.warning("[mem-att] rmtree fehlgeschlagen: %s", exc)
    return count


def read_bytes(memory_id: str, filename: str) -> Optional[bytes]:
    """Liefert die rohen Bytes einer Datei zurueck — fuer Download/Serve."""
    safe = _safe_filename(filename)
    target = memory_dir(memory_id) / safe
    if not target.is_file():
        return None
    return target.read_bytes()


# /shared/ ist der einzig akzeptable Source-Pfad fuer attach_from_path —
# ARIA bekommt Files vom User immer in /shared/uploads, eigene Files
# generiert sie in /shared/uploads/ als File-Marker. Kein Zugriff auf
# /root, /etc, /tmp, ssh-Keys, etc.
ALLOWED_SOURCE_PREFIXES = ("/shared/uploads/", "/shared/memory-attachments/")


def attach_from_path(memory_id: str, source_path: str) -> dict:
    """Kopiert eine existierende Datei aus /shared/* in das Anhang-Verzeichnis
    des Memories und gibt die neue Metadaten zurueck.

    Verwendung: ARIA bekommt z.B. ein User-Bild als `/shared/uploads/aria_<id>.jpg`.
    Statt das Bild dort liegen zu lassen (kein direkter Memory-Bezug), kopiert
    sie es via `memory_save(..., attach_paths=[<src>])` ins Memory-Verzeichnis.

    Pfadschutz: source_path MUSS unter /shared/ liegen — kein Zugriff auf
    Root-FS, SSH-Keys etc.
    """
    if not memory_id:
        raise ValueError("memory_id ist Pflicht")
    if not source_path or not isinstance(source_path, str):
        raise ValueError("source_path leer")
    if not any(source_path.startswith(p) for p in ALLOWED_SOURCE_PREFIXES):
        raise ValueError(f"source_path muss unter {' oder '.join(ALLOWED_SOURCE_PREFIXES)} liegen")
    src = Path(source_path)
    if not src.is_file():
        raise ValueError(f"Datei nicht gefunden: {source_path}")
    size = src.stat().st_size
    if size > MAX_BYTES:
        raise ValueError(f"Datei zu gross ({size} > {MAX_BYTES} Byte)")
    # Reuse save_attachment damit Filename-Sanitization + Logging konsistent
    data = src.read_bytes()
    return save_attachment(memory_id, src.name, data)