AgentEvalTool/backend/agenteval/storage/file_repository.py

"""Repository layer for file management (categories + records)."""

import os
from typing import Optional

from sqlmodel import Session, select

from agenteval.storage.db import (
    FILES_DIR,
    FileCategoryDB,
    FileRecordDB,
    get_session,
    new_uuid,
    utc_now,
)


class FileCategoryRepository:
    """Repository for file categories (tree structure)."""

    def __init__(self, session: Optional[Session] = None):
        self.session = session or get_session()

    def list_all(self) -> list[FileCategoryDB]:
        statement = select(FileCategoryDB).order_by(FileCategoryDB.created_at.asc())
        return list(self.session.exec(statement).all())

    def get(self, category_id: str) -> Optional[FileCategoryDB]:
        return self.session.get(FileCategoryDB, category_id)

    def get_tree(self) -> list[dict]:
        """Return categories as a nested tree structure for frontend Tree component."""
        all_cats = self.list_all()
        cat_map: dict[str, dict] = {}
        roots: list[dict] = []

        for cat in all_cats:
            node = {
                "key": cat.id,
                "title": cat.name,
                "parent_id": cat.parent_id,
                "children": [],
            }
            cat_map[cat.id] = node

        for cat in all_cats:
            node = cat_map[cat.id]
            if cat.parent_id and cat.parent_id in cat_map:
                cat_map[cat.parent_id]["children"].append(node)
            else:
                roots.append(node)

        return roots

    def create(self, name: str, parent_id: Optional[str] = None) -> FileCategoryDB:
        db = FileCategoryDB(
            id=new_uuid(),
            name=name,
            parent_id=parent_id,
            created_at=utc_now(),
            updated_at=utc_now(),
        )
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return db

    def update(self, category_id: str, name: str) -> Optional[FileCategoryDB]:
        existing = self.session.get(FileCategoryDB, category_id)
        if not existing:
            return None
        existing.name = name
        existing.updated_at = utc_now()
        self.session.add(existing)
        self.session.commit()
        self.session.refresh(existing)
        return existing

    def delete(self, category_id: str) -> bool:
        """Delete a category and cascade-delete children + files.

        Physical files are cleaned up via FileRecordRepository.delete().
        We must explicitly delete files first to trigger physical cleanup,
        because the DB cascade only removes rows.
        """
        existing = self.session.get(FileCategoryDB, category_id)
        if not existing:
            return False

        # Collect all file IDs to clean up physical files
        file_ids = self._collect_file_ids(existing)

        # Delete physical files
        file_repo = FileRecordRepository(self.session)
        for fid in file_ids:
            file_repo._remove_physical_file(fid)

        self.session.delete(existing)
        self.session.commit()
        return True

    def _collect_file_ids(self, category: FileCategoryDB) -> list[str]:
        """Recursively collect all file IDs under a category and its children."""
        file_ids = [f.id for f in category.files]
        for child in category.children:
            file_ids.extend(self._collect_file_ids(child))
        return file_ids


class FileRecordRepository:
    """Repository for uploaded file records."""

    def __init__(self, session: Optional[Session] = None):
        self.session = session or get_session()

    def list_all(self, category_id: Optional[str] = None) -> list[FileRecordDB]:
        statement = select(FileRecordDB)
        if category_id:
            # Also include files in subcategories
            cat_repo = FileCategoryRepository(self.session)
            cat_ids = self._get_descendant_ids(category_id, cat_repo)
            cat_ids.append(category_id)
            from sqlmodel import col

            statement = statement.where(col(FileRecordDB.category_id).in_(cat_ids))
        else:
            # Only filter when category_id is explicitly provided;
            # None means "all files" (no filter).
            pass
        statement = statement.order_by(FileRecordDB.created_at.desc())
        return list(self.session.exec(statement).all())

    def _get_descendant_ids(self, parent_id: str, cat_repo: FileCategoryRepository) -> list[str]:
        """Recursively collect IDs of all descendant categories."""
        ids: list[str] = []
        all_cats = cat_repo.list_all()
        children = [c for c in all_cats if c.parent_id == parent_id]
        for child in children:
            ids.append(child.id)
            ids.extend(self._get_descendant_ids(child.id, cat_repo))
        return ids

    def get(self, file_id: str) -> Optional[FileRecordDB]:
        return self.session.get(FileRecordDB, file_id)

    def create(
        self,
        original_name: str,
        storage_name: str,
        file_size: int,
        mime_type: str,
        file_ext: str,
        category_id: Optional[str] = None,
    ) -> FileRecordDB:
        db = FileRecordDB(
            id=new_uuid(),
            original_name=original_name,
            storage_name=storage_name,
            category_id=category_id,
            file_size=file_size,
            mime_type=mime_type,
            file_ext=file_ext,
            created_at=utc_now(),
        )
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return db

    def delete(self, file_id: str) -> bool:
        record = self.session.get(FileRecordDB, file_id)
        if not record:
            return False
        self._remove_physical_file(file_id)
        self.session.delete(record)
        self.session.commit()
        return True

    def _remove_physical_file(self, file_id: str) -> None:
        """Remove the physical file from disk if it exists."""
        record = self.session.get(FileRecordDB, file_id)
        if not record:
            return
        file_path = FILES_DIR / record.storage_name
        if file_path.exists():
            os.remove(file_path)