from __future__ import annotations

from repo_core.base import BaseRepository

from .model import CodalFile


class CodalFileRepository(BaseRepository[CodalFile]):
    TABLE = "codal_files"
    _backends = {"db"}

    async def upsert_file(self, data: dict) -> CodalFile:
        """Insert or update by statement_id + filename."""
        return await self.save(self.TABLE, data)

    async def find_by_statement_id(self, statement_id: str) -> list[CodalFile]:
        return await self.find(self.TABLE, {"statement_id": statement_id})

    async def find_pending(self, limit: int = 100) -> list[CodalFile]:
        return await self.find(
            self.TABLE,
            {"download_status": "pending"},
            limit=limit,
        )

    async def mark_downloaded(
        self,
        file_id: str,
        *,
        size_bytes: int,
        content_type: str,
        sha256: str | None = None,
        rustfs_bucket: str,
        rustfs_key: str,
    ) -> CodalFile | None:
        """Mark a file as successfully downloaded with RustFS location."""
        return await self.update_by_id(
            self.TABLE,
            file_id,
            {
                "download_status": "downloaded",
                "size_bytes": size_bytes,
                "content_type": content_type,
                "sha256": sha256,
                "rustfs_bucket": rustfs_bucket,
                "rustfs_key": rustfs_key,
            },
        )

    async def mark_failed(self, file_id: str, error: str) -> CodalFile | None:
        """Mark a file download as failed."""
        return await self.update_by_id(
            self.TABLE,
            file_id,
            {
                "download_status": "failed",
                "download_error": error,
            },
        )

    async def count_by_status(self, status: str) -> int:
        return await self.count(self.TABLE, {"download_status": status})
