transcription/src/transcription/services/store.py

from __future__ import annotations

import hashlib
import logging
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from uuid import UUID
from uuid import uuid4

from sqlmodel import select
from sqlmodel.ext.asyncio.session import AsyncSession

from transcription.config import Settings
from transcription.config import get_settings
from transcription.errors import AppError
from transcription.errors import ErrorCategory

from ..db.models import Document
from ..db.models import Job
from ..db.models import JobSource
from ..db.models import JobSourceStatus
from ..db.models import Source
from ..db.session import SessionFactory
from ..db.session import session_scope
from .media_storage import build_stored_filename
from .media_storage import write_media_bytes
from .sources import TranscriptionError
from .sources import build_prompt_execution
from .sources import validate_source_content

logger = logging.getLogger(__name__)


class SourceStorageError(AppError):
    """Raised when Source content cannot be validated or persisted safely."""


@dataclass(frozen=True)
class JobCreateResult:
    """Summary of explicit Job create records."""

    document_id: UUID
    job_id: UUID
    source_ids: tuple[UUID, ...]


@dataclass(frozen=True)
class DocumentJobResult:
    """Summary of a Document, Source, and Job created together."""

    document_id: UUID
    job_id: UUID
    stored_path: Path
    original_filename: str


@dataclass(frozen=True)
class PendingStoredSource:
    """Pre-staged Source artifact tied to a Source id."""

    source_id: UUID
    original_filename: str
    stored_path: Path
    file_hash: str
    file_size_bytes: int


async def create_document_job(
    *,
    filename: str,
    file_bytes: bytes,
    session: AsyncSession | None = None,
    session_factory: SessionFactory | None = None,
    settings: Settings | None = None,
) -> DocumentJobResult:
    """Create a Document, its first Source, and a queued Job.

    Owns its own session when the caller does not supply one, so UI callers
    never have to import a session scope.
    """
    runtime_settings = settings or get_settings()
    prompt_execution = build_prompt_execution(settings=runtime_settings)
    document_id = uuid4()
    source_id = uuid4()
    stored_path = await store_source_file(
        filename=filename,
        file_bytes=file_bytes,
        settings=runtime_settings,
        relative_directory=Path("documents") / str(document_id),
        filename_stem=str(source_id),
    )
    file_hash, file_size_bytes = _compute_file_metadata(file_bytes)
    try:
        async with session_scope(
            session_factory=session_factory,
            session=session,
            settings=runtime_settings,
        ) as _session:
            document, job = await _create_document_job_records(
                session=_session,
                document_id=document_id,
                source_id=source_id,
                original_filename=filename,
                stored_path=stored_path,
                file_hash=file_hash,
                file_size_bytes=file_size_bytes,
                prompt_execution=prompt_execution,
            )
    except Exception as exc:
        _best_effort_delete(stored_path)
        raise SourceStorageError(
            "Failed to create Document, Source, and Job records",
            category=ErrorCategory.INFRA_TRANSIENT,
            suggestion="Retry upload. If this keeps happening, verify database availability.",
            retriable=True,
        ) from exc

    logger.info("Created document job document_id=%s job_id=%s", document.id, job.id)
    return DocumentJobResult(
        document_id=document.id,
        job_id=job.id,
        stored_path=stored_path,
        original_filename=Path(filename).name,
    )


async def create_job_for_document(
    *,
    document_id: UUID,
    source_files: Sequence[tuple[str, bytes]],
    session: AsyncSession | None = None,
    session_factory: SessionFactory | None = None,
    provider: str | None = None,
    model: str | None = None,
    settings: Settings | None = None,
) -> JobCreateResult:
    """Create a queued Job for an existing Document with one or more Sources.

    Owns its own session when the caller does not supply one, so UI callers
    never have to import a session scope.
    """
    if not source_files:
        raise SourceStorageError(
            "At least one Source file is required to create a Job",
            category=ErrorCategory.VALIDATION,
            suggestion="Upload one or more files and try again.",
        )

    runtime_settings = settings or get_settings()
    prompt_execution = build_prompt_execution(settings=runtime_settings)
    sorted_source_files = sorted(source_files, key=lambda item: Path(item[0]).name.casefold())
    stored_sources: list[PendingStoredSource] = []
    for filename, file_bytes in sorted_source_files:
        source_id = uuid4()
        stored_path = await store_source_file(
            filename=filename,
            file_bytes=file_bytes,
            settings=runtime_settings,
            relative_directory=Path("documents") / str(document_id),
            filename_stem=str(source_id),
        )
        stored_sources.append(
            PendingStoredSource(
                source_id=source_id,
                original_filename=filename,
                stored_path=stored_path,
                file_hash=_compute_file_hash(file_bytes),
                file_size_bytes=len(file_bytes),
            )
        )

    try:
        async with session_scope(
            session_factory=session_factory,
            session=session,
            settings=runtime_settings,
        ) as _session:
            job, source_ids = await _create_job_for_document_records(
                session=_session,
                document_id=document_id,
                stored_sources=stored_sources,
                provider=provider,
                model=model,
                prompt_execution=prompt_execution,
            )
    except Exception as exc:
        for source in stored_sources:
            _best_effort_delete(source.stored_path)
        raise SourceStorageError(
            "Failed to create Job records from Source files",
            category=ErrorCategory.INFRA_TRANSIENT,
            suggestion="Retry creation. If this keeps happening, verify database availability.",
            retriable=True,
        ) from exc

    logger.info("Created explicit job document_id=%s job_id=%s sources=%s", document_id, job.id, len(source_ids))
    return JobCreateResult(
        document_id=document_id,
        job_id=job.id,
        source_ids=tuple(source_ids),
    )


async def _create_document_job_records(
    *,
    session: AsyncSession,
    document_id: UUID,
    source_id: UUID,
    original_filename: str,
    stored_path: Path,
    file_hash: str,
    file_size_bytes: int,
    prompt_execution,
) -> tuple[Document, Job]:
    document = Document(
        id=document_id,
        name=Path(original_filename).name,
    )
    session.add(document)
    await session.flush()

    job = Job(
        document_id=document.id,
        prompt_name=prompt_execution.prompt_name,
        prompt_hash=prompt_execution.prompt_hash,
        system_prompt=prompt_execution.system_prompt,
        user_prompt=prompt_execution.user_prompt,
        temperature=prompt_execution.temperature,
        top_p=prompt_execution.top_p,
    )
    session.add(job)
    await session.flush()

    source = Source(
        id=source_id,
        document_id=document.id,
        page_number=1,
        upload_name=Path(original_filename).name,
        filename=stored_path.name,
        file_path=str(stored_path),
        file_hash=file_hash,
        file_size_bytes=file_size_bytes,
    )
    session.add(source)
    await session.flush()

    session.add(
        JobSource(
            job_id=job.id,
            source_id=source.id,
            status=JobSourceStatus.PENDING,
        )
    )

    await session.commit()
    await session.refresh(document)
    await session.refresh(job)
    return document, job


async def _create_job_for_document_records(
    *,
    session: AsyncSession,
    document_id: UUID,
    stored_sources: Sequence[PendingStoredSource],
    provider: str | None,
    model: str | None,
    prompt_execution,
) -> tuple[Job, list[UUID]]:
    document = await session.get(Document, document_id)
    if document is None:
        raise SourceStorageError(
            f"Document with id {document_id} not found",
            category=ErrorCategory.NOT_FOUND,
            suggestion="Select an existing document and retry.",
        )

    existing_sources = (await session.exec(select(Source).where(Source.document_id == document_id))).all()
    next_page_number = max((source.page_number for source in existing_sources), default=0) + 1

    job = Job(
        document_id=document_id,
        provider=(provider or None),
        model=(model or None),
        prompt_name=prompt_execution.prompt_name,
        prompt_hash=prompt_execution.prompt_hash,
        system_prompt=prompt_execution.system_prompt,
        user_prompt=prompt_execution.user_prompt,
        temperature=prompt_execution.temperature,
        top_p=prompt_execution.top_p,
    )
    session.add(job)
    await session.flush()

    source_ids: list[UUID] = []
    for page_offset, stored_source in enumerate(stored_sources):
        source = Source(
            id=stored_source.source_id,
            document_id=document_id,
            page_number=next_page_number + page_offset,
            upload_name=Path(stored_source.original_filename).name,
            filename=stored_source.stored_path.name,
            file_path=str(stored_source.stored_path),
            file_hash=stored_source.file_hash,
            file_size_bytes=stored_source.file_size_bytes,
        )
        session.add(source)
        await session.flush()
        source_ids.append(source.id)

        session.add(
            JobSource(
                job_id=job.id,
                source_id=source.id,
                status=JobSourceStatus.PENDING,
            )
        )

    await session.commit()
    await session.refresh(job)
    return job, source_ids


def _best_effort_delete(path: Path) -> None:
    try:
        if path.exists():
            path.unlink()
    except OSError:
        logger.warning("Failed to clean up Source file after database error: %s", path)


def _compute_file_hash(file_bytes: bytes) -> str:
    return hashlib.sha256(file_bytes).hexdigest()


def _compute_file_metadata(file_bytes: bytes) -> tuple[str, int]:
    return _compute_file_hash(file_bytes), len(file_bytes)


async def store_source_file(
    *,
    filename: str,
    file_bytes: bytes,
    settings: Settings | None = None,
    relative_directory: Path | None = None,
    filename_stem: str | None = None,
) -> Path:
    """Validate and persist a Source file to configured media storage."""
    runtime_settings = settings or get_settings()
    try:
        validate_source_content(filename=filename, content=file_bytes)
    except TranscriptionError as exc:
        raise SourceStorageError(
            exc.message,
            category=exc.category,
            suggestion=exc.suggestion,
            retriable=exc.retriable,
        ) from exc

    upload_dir = runtime_settings.upload_dir
    return await write_media_bytes(
        target_dir=upload_dir if relative_directory is None else upload_dir / relative_directory,
        stored_name=build_stored_filename(filename=filename, filename_stem=filename_stem),
        file_bytes=file_bytes,
        error=SourceStorageError,
        failure_message="Failed to persist Source file",
        failure_suggestion="Check upload directory permissions and available disk space, then retry.",
        log_label="Source file",
    )