"""Data access for prompt media uploads."""
from __future__ import annotations
from typing import Optional
from uuid import UUID, uuid4

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from src.domain.prompt_media.models import PromptMediaAsset


class PromptMediaRepository:
    def __init__(self, db: AsyncSession):
        self._db = db

    async def create(
        self,
        *,
        project_id: UUID,
        uploaded_by: Optional[UUID],
        kind: str,
        storage_relative_path: str,
        original_filename: str,
        mime_type: str,
        byte_size: int,
        transcript: Optional[str] = None,
        asset_id: Optional[UUID] = None,
    ) -> PromptMediaAsset:
        row = PromptMediaAsset(
            id=asset_id or uuid4(),
            project_id=project_id,
            uploaded_by=uploaded_by,
            kind=kind,
            storage_relative_path=storage_relative_path,
            original_filename=original_filename,
            mime_type=mime_type,
            byte_size=byte_size,
            transcript=transcript,
        )
        self._db.add(row)
        await self._db.flush()
        await self._db.refresh(row)
        return row

    async def get_by_id(self, asset_id: UUID) -> Optional[PromptMediaAsset]:
        r = await self._db.execute(select(PromptMediaAsset).where(PromptMediaAsset.id == asset_id))
        return r.scalar_one_or_none()

    async def list_by_ids_for_project(
        self, asset_ids: list[UUID], project_id: UUID
    ) -> list[PromptMediaAsset]:
        if not asset_ids:
            return []
        r = await self._db.execute(
            select(PromptMediaAsset)
            .where(PromptMediaAsset.project_id == project_id)
            .where(PromptMediaAsset.id.in_(asset_ids))
        )
        rows = list(r.scalars().all())
        order = {str(i): j for j, i in enumerate(asset_ids)}
        rows.sort(key=lambda x: order.get(str(x.id), 999))
        return rows

    async def attach_to_prompt(
        self, prompt_id: UUID, asset_ids: list[UUID], project_id: UUID
    ) -> int:
        """Set ``prompt_id`` on assets that are unlinked and belong to the project. Returns rows updated."""
        if not asset_ids:
            return 0
        result = await self._db.execute(
            update(PromptMediaAsset)
            .where(PromptMediaAsset.project_id == project_id)
            .where(PromptMediaAsset.id.in_(asset_ids))
            .where(PromptMediaAsset.prompt_id.is_(None))
            .values(prompt_id=prompt_id)
        )
        await self._db.flush()
        return int(result.rowcount or 0)