"""SQLAlchemy implementation of SegmentRepository."""

from collections.abc import Sequence
from uuid import UUID

from sqlalchemy import func, select, update

from noteflow.domain.entities import Segment
from noteflow.domain.value_objects import MeetingId
from noteflow.infrastructure.converters import OrmConverter
from noteflow.infrastructure.persistence.models import SegmentModel, WordTimingModel
from noteflow.infrastructure.persistence.repositories._base import BaseRepository


class SqlAlchemySegmentRepository(BaseRepository):
    """SQLAlchemy implementation of SegmentRepository."""

    def _create_segment_model(self, meeting_id: MeetingId, segment: Segment) -> SegmentModel:
        """Create ORM model for a segment without persisting.

        Args:
            meeting_id: Meeting identifier.
            segment: Domain segment.

        Returns:
            SegmentModel ready for persistence.
        """
        model = SegmentModel(
            meeting_id=UUID(str(meeting_id)),
            segment_id=segment.segment_id,
            text=segment.text,
            start_time=segment.start_time,
            end_time=segment.end_time,
            language=segment.language,
            language_confidence=segment.language_confidence,
            avg_logprob=segment.avg_logprob,
            no_speech_prob=segment.no_speech_prob,
            embedding=segment.embedding,
            speaker_id=segment.speaker_id,
            speaker_confidence=segment.speaker_confidence,
        )

        # Add word timings with position index
        for word_index, word in enumerate(segment.words):
            word_kwargs = OrmConverter.word_timing_to_orm_kwargs(word, word_index)
            word_model = WordTimingModel(**word_kwargs)
            model.words.append(word_model)

        return model

    async def add(self, meeting_id: MeetingId, segment: Segment) -> Segment:
        """Add a segment to a meeting.

        Args:
            meeting_id: Meeting identifier.
            segment: Segment to add.

        Returns:
            Added segment with db_id populated.
        """
        model = self._create_segment_model(meeting_id, segment)

        self._session.add(model)
        await self._session.flush()

        # Update segment with db_id
        segment.db_id = model.id
        segment.meeting_id = meeting_id
        return segment

    async def add_batch(
        self,
        meeting_id: MeetingId,
        segments: Sequence[Segment],
    ) -> Sequence[Segment]:
        """Add multiple segments to a meeting in batch.

        Uses add_all() with a single flush for better performance.

        Args:
            meeting_id: Meeting identifier.
            segments: Segments to add.

        Returns:
            Added segments with db_ids populated.
        """
        if not segments:
            return []

        # Build all models upfront
        models: list[SegmentModel] = []
        for segment in segments:
            model = self._create_segment_model(meeting_id, segment)
            models.append(model)

        # Add all models and flush once
        self._session.add_all(models)
        await self._session.flush()

        # Update segments with db_ids
        for segment, model in zip(segments, models, strict=True):
            segment.db_id = model.id
            segment.meeting_id = meeting_id

        return list(segments)

    async def get_by_meeting(
        self,
        meeting_id: MeetingId,
        include_words: bool = True,
    ) -> Sequence[Segment]:
        """Get all segments for a meeting.

        Args:
            meeting_id: Meeting identifier.
            include_words: Include word-level timing.

        Returns:
            List of segments ordered by segment_id.
        """
        stmt = (
            select(SegmentModel)
            .where(SegmentModel.meeting_id == UUID(str(meeting_id)))
            .order_by(SegmentModel.segment_id)
        )

        models = await self._execute_scalars(stmt)

        return [OrmConverter.segment_to_domain(m, include_words) for m in models]

    async def search_semantic(
        self,
        query_embedding: list[float],
        limit: int = 10,
        meeting_id: MeetingId | None = None,
    ) -> Sequence[tuple[Segment, float]]:
        """Search segments by semantic similarity.

        Args:
            query_embedding: Query embedding vector.
            limit: Maximum number of results.
            meeting_id: Optional meeting to restrict search to.

        Returns:
            List of (segment, similarity_score) tuples.
        """
        # Build query with cosine similarity
        similarity = SegmentModel.embedding.cosine_distance(query_embedding)

        stmt = select(SegmentModel, similarity.label("distance")).where(
            SegmentModel.embedding.is_not(None)
        )

        if meeting_id:
            stmt = stmt.where(SegmentModel.meeting_id == UUID(str(meeting_id)))

        stmt = stmt.order_by(similarity).limit(limit)

        result = await self._session.execute(stmt)
        rows = result.all()

        results: list[tuple[Segment, float]] = []
        for row in rows:
            model = row[0]
            distance = row[1]
            # Convert distance to similarity (1 - distance for cosine)
            similarity_score = 1.0 - float(distance)
            segment = OrmConverter.segment_to_domain(model, include_words=False)
            results.append((segment, similarity_score))

        return results

    async def update_embedding(
        self,
        segment_db_id: int,
        embedding: list[float],
    ) -> None:
        """Update the embedding for a segment.

        Uses direct UPDATE statement (single DB round-trip) instead of SELECT+UPDATE.

        Args:
            segment_db_id: Segment database primary key.
            embedding: New embedding vector.
        """
        stmt = (
            update(SegmentModel)
            .where(SegmentModel.id == segment_db_id)
            .values(embedding=embedding)
        )
        await self._session.execute(stmt)

    async def update_speaker(
        self,
        segment_db_id: int,
        speaker_id: str | None,
        speaker_confidence: float,
    ) -> None:
        """Update speaker diarization fields for a segment.

        Uses direct UPDATE statement (single DB round-trip) instead of SELECT+UPDATE.

        Args:
            segment_db_id: Segment database primary key.
            speaker_id: Speaker identifier from diarization.
            speaker_confidence: Confidence of speaker assignment (0.0-1.0).
        """
        stmt = (
            update(SegmentModel)
            .where(SegmentModel.id == segment_db_id)
            .values(speaker_id=speaker_id, speaker_confidence=speaker_confidence)
        )
        await self._session.execute(stmt)

    async def compute_next_segment_id(self, meeting_id: MeetingId) -> int:
        """Compute the next segment_id for a meeting.

        Args:
            meeting_id: Meeting identifier.

        Returns:
            Next segment_id (max + 1), or 0 if no segments exist.
        """
        stmt = select(func.max(SegmentModel.segment_id)).where(
            SegmentModel.meeting_id == UUID(str(meeting_id))
        )
        result = await self._session.execute(stmt)
        max_segment_id = result.scalar_one_or_none()
        return 0 if max_segment_id is None else int(max_segment_id) + 1
