iron/app/services/upload_service.py

from __future__ import annotations

import base64
import hashlib
from dataclasses import dataclass
from pathlib import Path

import anyio
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.config import get_settings
from app.core.exceptions import UploadConflictError, UploadNotFoundError, ValidationError
from app.core.ids import new_id
from app.models.entities import UploadSession
from app.repositories.file_repository import FileRepository
from app.repositories.upload_repository import UploadRepository
from app.services.file_service import CreateFileMetadataInput, FileService
from app.services.job_service import JobService
from app.services.local_storage_service import LocalStorageService
from app.services.storage_service import StorageService


@dataclass
class CreateUploadInput:
    upload_length: int
    upload_metadata: str | None = None


@dataclass
class FinalizeUploadInput:
    directory_id: str
    filename: str


class UploadService:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session
        self.repository = UploadRepository(session)
        self.file_service = FileService(session)
        self.file_repository = FileRepository(session)
        self.local_storage_service = LocalStorageService(session)
        self.storage_service = StorageService(session)
        self.job_service = JobService(session)
        self.settings = get_settings()

    async def ensure_temp_dir(self) -> Path:
        temp_dir = Path(self.settings.upload_temp_dir)
        await anyio.to_thread.run_sync(lambda: temp_dir.mkdir(parents=True, exist_ok=True))
        return temp_dir

    async def create_upload(self, payload: CreateUploadInput) -> UploadSession:
        if payload.upload_length <= 0:
            raise ValidationError("Upload-Length must be greater than zero.")

        temp_dir = await self.ensure_temp_dir()
        upload_id = new_id("upl")
        temp_path = temp_dir / f"{upload_id}.bin"
        await anyio.Path(temp_path).write_bytes(b"")

        upload_session = UploadSession(
            id=upload_id,
            total_size_bytes=payload.upload_length,
            received_size_bytes=0,
            status="created",
            temp_path=str(temp_path),
            tus_upload_url=f"/files/{upload_id}",
            upload_metadata_json=payload.upload_metadata,
        )
        created = await self.repository.create(upload_session)
        await self.session.commit()
        return created

    async def get_upload(self, upload_id: str) -> UploadSession:
        upload_session = await self.repository.get_by_id(upload_id)
        if upload_session is None:
            raise UploadNotFoundError()
        return upload_session

    async def append_upload_bytes(self, upload_id: str, offset: int, chunk: bytes) -> UploadSession:
        upload_session = await self.get_upload(upload_id)

        if upload_session.status in {"finalized", "finalizing", "expired"}:
            raise UploadConflictError("Upload session is not writable.")
        if offset != upload_session.received_size_bytes:
            raise UploadConflictError("Upload offset does not match the current stored offset.")
        if upload_session.received_size_bytes + len(chunk) > upload_session.total_size_bytes:
            raise UploadConflictError("Upload would exceed declared length.")

        temp_path = anyio.Path(upload_session.temp_path)
        async with await anyio.open_file(temp_path, "ab") as file_handle:
            await file_handle.write(chunk)

        upload_session.received_size_bytes += len(chunk)
        upload_session.status = (
            "uploaded"
            if upload_session.received_size_bytes == upload_session.total_size_bytes
            else "uploading"
        )
        saved = await self.repository.save(upload_session)
        await self.session.commit()
        return saved

    async def finalize_upload(self, upload_id: str, payload: FinalizeUploadInput):
        upload_session = await self.get_upload(upload_id)
        if upload_session.received_size_bytes != upload_session.total_size_bytes:
            raise UploadConflictError("Upload is not complete.")
        if upload_session.status == "finalized":
            raise UploadConflictError("Upload has already been finalized.")

        upload_session.status = "finalizing"
        await self.repository.save(upload_session)
        await self.session.flush()

        content_hash = await self._compute_sha256(upload_session.temp_path)
        mime_type = self._guess_mime_type(payload.filename)
        file_entry = await self.file_service.create_file_metadata(
            CreateFileMetadataInput(
                directory_id=payload.directory_id,
                name=payload.filename,
                mime_type=mime_type,
                size_bytes=upload_session.total_size_bytes,
                content_hash=content_hash,
            )
        )
        blob = await self.file_repository.get_primary_blob_for_file(file_entry.id)
        if blob is None:
            raise UploadConflictError("Primary blob was not created for finalized upload.")

        await self.local_storage_service.persist_blob_from_temp(blob, upload_session.temp_path)
        await self.job_service.enqueue_file_reconcile(file_entry.id)
        if mime_type is not None:
            await self.job_service.enqueue_preview_generation(file_entry.id)

        upload_session.directory_id = payload.directory_id
        upload_session.filename = payload.filename
        upload_session.status = "finalized"
        await self.repository.save(upload_session)
        await self.session.commit()
        await anyio.Path(upload_session.temp_path).unlink(missing_ok=True)

        return upload_session, file_entry

    async def _compute_sha256(self, temp_path: str) -> str:
        def _hash_file() -> str:
            digest = hashlib.sha256()
            with open(temp_path, "rb") as file_handle:
                while True:
                    chunk = file_handle.read(1024 * 1024)
                    if not chunk:
                        break
                    digest.update(chunk)
            return digest.hexdigest()

        return await anyio.to_thread.run_sync(_hash_file)

    @staticmethod
    def parse_upload_metadata(raw_metadata: str | None) -> dict[str, str]:
        if not raw_metadata:
            return {}

        parsed: dict[str, str] = {}
        for item in raw_metadata.split(","):
            item = item.strip()
            if not item:
                continue
            parts = item.split(" ", 1)
            key = parts[0]
            value = ""
            if len(parts) == 2 and parts[1]:
                value = base64.b64decode(parts[1]).decode("utf-8")
            parsed[key] = value
        return parsed

    @staticmethod
    def _guess_mime_type(filename: str) -> str | None:
        lowered = filename.lower()
        if lowered.endswith(".mp4"):
            return "video/mp4"
        if lowered.endswith(".jpg") or lowered.endswith(".jpeg"):
            return "image/jpeg"
        if lowered.endswith(".png"):
            return "image/png"
        if lowered.endswith(".pdf"):
            return "application/pdf"
        return None