"""SQLAlchemy implementation of DiarizationJobRepository."""

from collections.abc import Sequence
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Final
from uuid import UUID

from sqlalchemy import delete, select, update
from sqlalchemy.exc import IntegrityError

from noteflow.config.constants import ERR_SERVER_RESTARTED
from noteflow.domain.utils.time import utc_now
from noteflow.infrastructure.persistence.models import (
    DiarizationJobModel,
    StreamingDiarizationTurnModel,
)
from noteflow.infrastructure.persistence.repositories._base import BaseRepository

# Job status constants (mirrors proto enum)
JOB_STATUS_UNSPECIFIED: Final[int] = 0
JOB_STATUS_QUEUED: Final[int] = 1
JOB_STATUS_RUNNING: Final[int] = 2
JOB_STATUS_COMPLETED: Final[int] = 3
JOB_STATUS_FAILED: Final[int] = 4
JOB_STATUS_CANCELLED: Final[int] = 5


@dataclass
class DiarizationJob:
    """Data transfer object for diarization job state.

    Separate from ORM model to allow easy passing between layers.
    """

    job_id: str
    meeting_id: str
    status: int
    segments_updated: int = 0
    speaker_ids: list[str] = field(default_factory=list)
    error_message: str = ""
    created_at: datetime = field(default_factory=utc_now)
    updated_at: datetime = field(default_factory=utc_now)
    audio_duration_seconds: float | None = None
    started_at: datetime | None = None


@dataclass
class StreamingTurn:
    """Data transfer object for streaming diarization turn.

    Represents a speaker turn collected during real-time streaming diarization.
    """

    speaker: str
    start_time: float
    end_time: float
    confidence: float = 0.0


class SqlAlchemyDiarizationJobRepository(BaseRepository):
    """SQLAlchemy implementation of DiarizationJobRepository."""

    @staticmethod
    def _to_domain(model: DiarizationJobModel) -> DiarizationJob:
        """Convert ORM model to domain object."""
        return DiarizationJob(
            job_id=model.id,
            meeting_id=str(model.meeting_id),
            status=model.status,
            segments_updated=model.segments_updated,
            speaker_ids=list(model.speaker_ids),
            error_message=model.error_message,
            created_at=model.created_at,
            updated_at=model.updated_at,
            audio_duration_seconds=model.audio_duration_seconds,
            started_at=model.started_at,
        )

    async def create(self, job: DiarizationJob) -> DiarizationJob:
        """Persist a new diarization job.

        Args:
            job: Job to create.

        Returns:
            Created job.
        """
        model = DiarizationJobModel()
        model.id = job.job_id
        model.meeting_id = UUID(job.meeting_id)
        model.status = job.status
        model.segments_updated = job.segments_updated
        model.speaker_ids = job.speaker_ids
        model.error_message = job.error_message
        model.created_at = job.created_at
        model.updated_at = job.updated_at
        model.audio_duration_seconds = job.audio_duration_seconds
        model.started_at = job.started_at
        self._session.add(model)
        try:
            await self._session.flush()
        except IntegrityError as exc:
            msg = f"DiarizationJob {job.job_id} already exists"
            raise ValueError(msg) from exc
        return job

    async def get(self, job_id: str) -> DiarizationJob | None:
        """Retrieve a job by ID.

        Args:
            job_id: Job identifier.

        Returns:
            Job if found, None otherwise.
        """
        stmt = select(DiarizationJobModel).where(DiarizationJobModel.id == job_id)
        model = await self._execute_scalar(stmt)

        return None if model is None else self._to_domain(model)

    async def update_status(
        self,
        job_id: str,
        status: int,
        *,
        segments_updated: int | None = None,
        speaker_ids: list[str] | None = None,
        error_message: str | None = None,
        started_at: datetime | None = None,
    ) -> bool:
        """Update job status and optional fields.

        Args:
            job_id: Job identifier.
            status: New status value.
            segments_updated: Optional segments count.
            speaker_ids: Optional speaker IDs list.
            error_message: Optional error message.
            started_at: Optional job start timestamp.

        Returns:
            True if job was updated, False if not found.
        """
        values: dict[str, int | list[str] | str | datetime] = {
            "status": status,
            "updated_at": utc_now(),
        }
        if segments_updated is not None:
            values["segments_updated"] = segments_updated
        if speaker_ids is not None:
            values["speaker_ids"] = speaker_ids
        if error_message is not None:
            values["error_message"] = error_message
        if started_at is not None:
            values["started_at"] = started_at

        stmt = update(DiarizationJobModel).where(DiarizationJobModel.id == job_id).values(**values)
        rowcount = await self._execute_update(stmt)
        return rowcount > 0

    async def list_for_meeting(self, meeting_id: str) -> Sequence[DiarizationJob]:
        """List all jobs for a meeting.

        Args:
            meeting_id: Meeting identifier.

        Returns:
            List of jobs ordered by creation time (newest first).
        """
        stmt = (
            select(DiarizationJobModel)
            .where(DiarizationJobModel.meeting_id == UUID(meeting_id))
            .order_by(DiarizationJobModel.created_at.desc())
        )
        models = await self._execute_scalars(stmt)
        return [self._to_domain(model) for model in models]

    async def get_active_for_meeting(self, meeting_id: str) -> DiarizationJob | None:
        """Get active (QUEUED or RUNNING) job for a meeting.

        Used to prevent duplicate concurrent jobs for the same meeting.

        Args:
            meeting_id: Meeting identifier.

        Returns:
            Active job if found, None otherwise.
        """
        stmt = (
            select(DiarizationJobModel)
            .where(
                DiarizationJobModel.meeting_id == UUID(meeting_id),
                DiarizationJobModel.status.in_([JOB_STATUS_QUEUED, JOB_STATUS_RUNNING]),
            )
            .order_by(DiarizationJobModel.created_at.desc())
            .limit(1)
        )
        model = await self._execute_scalar(stmt)
        return self._to_domain(model) if model else None

    async def get_all_active(self) -> list[DiarizationJob]:
        """Get all active (QUEUED or RUNNING) jobs.

        Used by client to restore polling after reconnection.

        Returns:
            List of active jobs across all meetings.
        """
        stmt = (
            select(DiarizationJobModel)
            .where(DiarizationJobModel.status.in_([JOB_STATUS_QUEUED, JOB_STATUS_RUNNING]))
            .order_by(DiarizationJobModel.created_at.desc())
        )
        models = await self._execute_scalars(stmt)
        return [self._to_domain(model) for model in models]

    async def mark_running_as_failed(self, error_message: str = ERR_SERVER_RESTARTED) -> int:
        """Mark all QUEUED or RUNNING jobs as FAILED.

        Used during crash recovery to mark orphaned jobs.

        Args:
            error_message: Error message to set on failed jobs.

        Returns:
            Number of jobs marked as failed.
        """
        stmt = (
            update(DiarizationJobModel)
            .where(DiarizationJobModel.status.in_([JOB_STATUS_QUEUED, JOB_STATUS_RUNNING]))
            .values(
                status=JOB_STATUS_FAILED,
                error_message=error_message,
                updated_at=utc_now(),
            )
        )
        return await self._execute_update(stmt)

    async def prune_completed(self, ttl_seconds: float) -> int:
        """Delete completed/failed jobs older than TTL.

        Args:
            ttl_seconds: Time-to-live in seconds.

        Returns:
            Number of jobs deleted.
        """
        cutoff = utc_now().timestamp() - ttl_seconds
        cutoff_dt = datetime.fromtimestamp(cutoff, tz=UTC)

        stmt = delete(DiarizationJobModel).where(
            DiarizationJobModel.status.in_(
                [JOB_STATUS_COMPLETED, JOB_STATUS_FAILED, JOB_STATUS_CANCELLED]
            ),
            DiarizationJobModel.updated_at < cutoff_dt,
        )
        return await self._execute_delete(stmt)

    # Streaming diarization turn methods

    async def add_streaming_turns(self, meeting_id: str, turns: Sequence[StreamingTurn]) -> int:
        """Persist streaming diarization turns for a meeting.

        Immediately stores speaker turns as they arrive during streaming.
        Used for crash resilience.

        Args:
            meeting_id: Meeting identifier.
            turns: Speaker turns to persist.

        Returns:
            Number of turns added.
        """
        if not turns:
            return 0

        meeting_uuid = UUID(meeting_id)
        for turn in turns:
            model = StreamingDiarizationTurnModel()
            model.meeting_id = meeting_uuid
            model.speaker = turn.speaker
            model.start_time = turn.start_time
            model.end_time = turn.end_time
            model.confidence = turn.confidence
            self._session.add(model)

        await self._session.flush()
        return len(turns)

    async def get_streaming_turns(self, meeting_id: str) -> list[StreamingTurn]:
        """Retrieve streaming diarization turns for a meeting.

        Used to recover streaming state after server restart.

        Args:
            meeting_id: Meeting identifier.

        Returns:
            List of streaming turns ordered by start time.
        """
        stmt = (
            select(StreamingDiarizationTurnModel)
            .where(StreamingDiarizationTurnModel.meeting_id == UUID(meeting_id))
            .order_by(StreamingDiarizationTurnModel.start_time)
        )
        models = await self._execute_scalars(stmt)
        return [
            StreamingTurn(
                speaker=model.speaker,
                start_time=model.start_time,
                end_time=model.end_time,
                confidence=model.confidence,
            )
            for model in models
        ]

    async def clear_streaming_turns(self, meeting_id: str) -> int:
        """Delete streaming diarization turns for a meeting.

        Called when a meeting stops recording to clean up temporary turns.

        Args:
            meeting_id: Meeting identifier.

        Returns:
            Number of turns deleted.
        """
        stmt = delete(StreamingDiarizationTurnModel).where(
            StreamingDiarizationTurnModel.meeting_id == UUID(meeting_id)
        )
        return await self._execute_delete(stmt)
