"""Diarization job management (start, status, cancel)."""

from __future__ import annotations

import asyncio
from collections.abc import Callable
from typing import TYPE_CHECKING, cast
from uuid import UUID, uuid4

import grpc

from noteflow.domain.ports.unit_of_work import UnitOfWork
from noteflow.domain.utils import utc_now
from noteflow.domain.value_objects import MeetingState
from noteflow.infrastructure.logging import get_logger, log_state_transition
from noteflow.infrastructure.persistence.repositories import DiarizationJob

from ...proto import noteflow_pb2
from .._types import GrpcStatusContext
from ..converters import parse_meeting_id
from ._status import JobStatusMixin
from ._types import DIARIZATION_TIMEOUT_SECONDS

if TYPE_CHECKING:
    from ..protocols import ServicerHost

logger = get_logger(__name__)


def _job_status_name(status: int) -> str:
    name_fn = cast(Callable[[int], str], noteflow_pb2.JobStatus.Name)
    return name_fn(status)


def create_diarization_error_response(
    error_message: str,
    status: noteflow_pb2.JobStatus | str = noteflow_pb2.JOB_STATUS_FAILED,
    *,
    context: GrpcStatusContext | None = None,
    grpc_code: grpc.StatusCode | None = None,
    job_id: str = "",
) -> noteflow_pb2.RefineSpeakerDiarizationResponse:
    """Create error response for RefineSpeakerDiarization.

    Consolidates the 6+ duplicated response construction patterns.

    Args:
        error_message: Error message describing the failure.
        status: Job status code (default: JOB_STATUS_FAILED).
        context: Optional gRPC context for setting status code.
        grpc_code: Optional gRPC status code to set.
        job_id: Optional job ID to include in response.

    Returns:
        Populated RefineSpeakerDiarizationResponse with error state.
    """
    if context is not None and grpc_code is not None:
        context.set_code(grpc_code)
        context.set_details(error_message)
    return noteflow_pb2.RefineSpeakerDiarizationResponse(
        segments_updated=0,
        speaker_ids=[],
        error_message=error_message,
        job_id=job_id,
        status=status,
    )


def _validate_diarization_preconditions(
    servicer: ServicerHost,
    request: noteflow_pb2.RefineSpeakerDiarizationRequest,
    context: GrpcStatusContext,
) -> noteflow_pb2.RefineSpeakerDiarizationResponse | None:
    """Validate preconditions before starting diarization job.

    Returns error response if validation fails, None if valid.
    """
    if not servicer.diarization_refinement_enabled:
        return create_diarization_error_response("Diarization refinement disabled on server")

    if servicer.diarization_engine is None:
        return create_diarization_error_response(
            "Diarization not enabled on server",
            context=context,
            grpc_code=grpc.StatusCode.UNAVAILABLE,
        )

    try:
        UUID(request.meeting_id)
    except ValueError:
        return create_diarization_error_response("Invalid meeting_id")
    return None


async def _create_and_persist_job(
    job_id: str,
    meeting_id: str,
    audio_duration: float | None,
    repo: UnitOfWork,
) -> bool:
    """Create job record and persist to database.

    Returns:
        True if job was persisted, False if database doesn't support diarization jobs.
    """
    if not repo.supports_diarization_jobs:
        return False

    job = DiarizationJob(
        job_id=job_id,
        meeting_id=meeting_id,
        status=noteflow_pb2.JOB_STATUS_QUEUED,
        audio_duration_seconds=audio_duration,
    )

    await repo.diarization_jobs.create(job)
    await repo.commit()
    logger.debug(
        "Diarization job %s persisted to database",
        job_id,
        extra={"meeting_id": meeting_id},
    )
    return True


class JobsMixin(JobStatusMixin):
    """Mixin providing diarization job management."""

    async def start_diarization_job(
        self: ServicerHost,
        request: noteflow_pb2.RefineSpeakerDiarizationRequest,
        context: GrpcStatusContext,
    ) -> noteflow_pb2.RefineSpeakerDiarizationResponse:
        """Start a new diarization refinement job.

        Validates the request, creates a job record, and launches the background task.
        """
        if error := _validate_diarization_preconditions(self, request, context):
            return error

        async with self.create_repository_provider() as repo:
            meeting = await repo.meetings.get(parse_meeting_id(request.meeting_id))
            if meeting is None:
                return create_diarization_error_response("Meeting not found")

            valid_states = (MeetingState.STOPPED, MeetingState.COMPLETED, MeetingState.ERROR)
            if meeting.state not in valid_states:
                return create_diarization_error_response(
                    f"Meeting must be stopped before refinement (state: {meeting.state.name.lower()})"
                )

            # Concurrency guard: check for existing active job
            if repo.supports_diarization_jobs:
                active_job = await repo.diarization_jobs.get_active_for_meeting(request.meeting_id)
                if active_job is not None:
                    return create_diarization_error_response(
                        f"Diarization already in progress (job: {active_job.job_id})",
                        status=_job_status_name(active_job.status),
                        context=context,
                        grpc_code=grpc.StatusCode.ALREADY_EXISTS,
                        job_id=active_job.job_id,
                    )

            job_id = str(uuid4())
            persisted = await _create_and_persist_job(
                job_id, request.meeting_id, meeting.duration_seconds, repo
            )
            if not persisted:
                return create_diarization_error_response(
                    "Diarization requires database support",
                    context=context,
                    grpc_code=grpc.StatusCode.FAILED_PRECONDITION,
                )

        num_speakers = request.num_speakers or None
        task = asyncio.create_task(self.run_diarization_job(job_id, num_speakers))
        self.diarization_tasks[job_id] = task

        return noteflow_pb2.RefineSpeakerDiarizationResponse(
            segments_updated=0, job_id=job_id, status=noteflow_pb2.JOB_STATUS_QUEUED
        )

    async def run_diarization_job(
        self: ServicerHost,
        job_id: str,
        num_speakers: int | None,
    ) -> None:
        """Run background diarization job.

        Updates job status in repository as the job progresses.
        """
        # Get meeting_id and update status to RUNNING
        meeting_id: str | None = None
        job: DiarizationJob | None = None
        async with self.create_repository_provider() as repo:
            if not repo.supports_diarization_jobs:
                logger.error("Diarization job %s cannot run: database required", job_id)
                return

            job = await repo.diarization_jobs.get(job_id)
            if job is None:
                logger.warning("Diarization job %s not found in database", job_id)
                return
            meeting_id = job.meeting_id
            old_status = job.status
            await repo.diarization_jobs.update_status(
                job_id,
                noteflow_pb2.JOB_STATUS_RUNNING,
                started_at=utc_now(),
            )
            await repo.commit()
            log_state_transition(
                "diarization_job",
                job_id,
                _job_status_name(old_status),
                _job_status_name(int(noteflow_pb2.JOB_STATUS_RUNNING)),
                meeting_id=meeting_id,
            )
        try:
            async with asyncio.timeout(DIARIZATION_TIMEOUT_SECONDS):
                updated_count = await self.refine_speaker_diarization(
                    meeting_id=meeting_id,
                    num_speakers=num_speakers,
                )
                speaker_ids = await self.collect_speaker_ids(meeting_id)

            # Update status to COMPLETED
            await self.update_job_completed(job_id, job, updated_count, speaker_ids)

        except TimeoutError:
            await self.handle_job_timeout(job_id, job, meeting_id)

        except asyncio.CancelledError:
            await self.handle_job_cancelled(job_id, job, meeting_id)
            raise  # Re-raise to propagate cancellation

        # INTENTIONAL BROAD HANDLER: Job error boundary
        # - Diarization can fail in many ways (model errors, audio issues, etc.)
        # - Must capture any failure and update job status
        except Exception as exc:
            await self.handle_job_failed(job_id, job, meeting_id, exc)
