"""Diarization job management mixin for gRPC service."""

from __future__ import annotations

import asyncio
import contextlib
from datetime import datetime
from typing import TYPE_CHECKING, Protocol, Self

from noteflow.domain.utils.time import utc_now
from noteflow.infrastructure.logging import get_logger

from ..proto import noteflow_pb2
from ._types import GrpcContext
from .errors import ERR_CANCELLED_BY_USER, abort_not_found

if TYPE_CHECKING:
    from noteflow.domain.ports.repositories import DiarizationJobRepository
    from noteflow.domain.ports.unit_of_work import UnitOfWork
    from noteflow.infrastructure.persistence.repositories import DiarizationJob

logger = get_logger(__name__)


# Diarization job TTL default (1 hour in seconds)
_DEFAULT_JOB_TTL_SECONDS: float = 3600.0


class DiarizationJobRepositoryProvider(Protocol):
    supports_diarization_jobs: bool
    diarization_jobs: DiarizationJobRepository

    async def commit(self) -> None: ...

    async def __aenter__(self) -> Self: ...

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> None: ...


class DiarizationJobServicer(Protocol):
    diarization_tasks: dict[str, asyncio.Task[None]]
    diarization_jobs: dict[str, DiarizationJob]

    @property
    def diarization_job_ttl_seconds(self) -> float: ...

    async def prune_diarization_jobs(self) -> None: ...

    def create_repository_provider(self) -> UnitOfWork: ...


class DiarizationJobMixin:
    """Mixin providing diarization job management functionality.

    Handles job status queries, cancellation, and pruning.
    Requires host to implement ServicerHost protocol.
    """

    @property
    def diarization_job_ttl_seconds(self) -> float:
        """Return diarization job TTL from settings.

        Attempts to load from settings, with fallback for test environments.
        """
        try:
            from noteflow.config.constants import SECONDS_PER_HOUR
            from noteflow.config.settings import get_settings

            settings = get_settings()
            return float(settings.diarization_job_ttl_hours * SECONDS_PER_HOUR)
        # INTENTIONAL BROAD HANDLER: Fallback for testing environments
        # - Settings may fail to load in unit tests without full config
        except Exception:
            return _DEFAULT_JOB_TTL_SECONDS

    async def prune_diarization_jobs(self: DiarizationJobServicer) -> None:
        """Remove completed diarization jobs older than retention window.

        Prunes both in-memory task references and database records.
        """
        # Clean up in-memory task references for completed tasks
        completed_tasks = [
            job_id for job_id, task in self.diarization_tasks.items() if task.done()
        ]
        for job_id in completed_tasks:
            self.diarization_tasks.pop(job_id, None)

        # Prune old completed jobs from database
        async with self.create_repository_provider() as repo:
            if not repo.supports_diarization_jobs:
                logger.debug("Job pruning skipped: database required")
                return
            pruned = await repo.diarization_jobs.prune_completed(
                self.diarization_job_ttl_seconds
            )
            await repo.commit()
            if pruned > 0:
                logger.debug("Pruned %d completed diarization jobs", pruned)

    async def GetDiarizationJobStatus(
        self: DiarizationJobServicer,
        request: noteflow_pb2.GetDiarizationJobStatusRequest,
        context: GrpcContext,
    ) -> noteflow_pb2.DiarizationJobStatus:
        """Return current status for a diarization job.

        Queries job state from repository for persistence across restarts.
        """
        await self.prune_diarization_jobs()

        async with self.create_repository_provider() as repo:
            if not repo.supports_diarization_jobs:
                await abort_not_found(context, "Diarization jobs (database required)", "")
                raise  # Unreachable but helps type checker

            job = await repo.diarization_jobs.get(request.job_id)
            if job is None:
                await abort_not_found(context, "Diarization job", request.job_id)
                raise  # Unreachable but helps type checker

            # Calculate progress percentage (time-based for running jobs)
            progress_percent = 0.0
            if job.status == noteflow_pb2.JOB_STATUS_COMPLETED:
                progress_percent = 100.0
            elif job.status == noteflow_pb2.JOB_STATUS_RUNNING and job.started_at is not None:
                # All datetimes should now be timezone-aware UTC.
                now = utc_now()
                # Ensure started_at is also aware; should be UTC from repository.
                started: datetime = job.started_at
                elapsed = (now - started).total_seconds()
                audio_duration = job.audio_duration_seconds
                if audio_duration is not None and audio_duration > 0:
                    # ~10 seconds processing per 60 seconds audio
                    estimated_duration = audio_duration * 0.17
                    progress_percent = min(95.0, (elapsed / estimated_duration) * 100)
                else:
                    # Fallback: assume 2 minutes total
                    progress_percent = min(95.0, (elapsed / 120) * 100)

            return noteflow_pb2.DiarizationJobStatus(
                job_id=job.job_id,
                status=int(job.status),
                segments_updated=job.segments_updated,
                speaker_ids=job.speaker_ids,
                error_message=job.error_message,
                progress_percent=progress_percent,
            )

    async def CancelDiarizationJob(
        self: DiarizationJobServicer,
        request: noteflow_pb2.CancelDiarizationJobRequest,
        context: GrpcContext,
    ) -> noteflow_pb2.CancelDiarizationJobResponse:
        """Cancel a running or queued diarization job.

        Cancels the background asyncio task and updates job status to CANCELLED.
        """
        job_id = request.job_id
        response = noteflow_pb2.CancelDiarizationJobResponse()

        # Cancel the asyncio task if it exists and is still running
        task = self.diarization_tasks.get(job_id)
        if task is not None and not task.done():
            task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await task

        async with self.create_repository_provider() as repo:
            if not repo.supports_diarization_jobs:
                response.success = False
                response.error_message = "Diarization requires database support"
                response.status = noteflow_pb2.JOB_STATUS_UNSPECIFIED
                return response

            job = await repo.diarization_jobs.get(job_id)
            if job is None:
                response.success = False
                response.error_message = "Job not found"
                response.status = noteflow_pb2.JOB_STATUS_UNSPECIFIED
                return response

            # Only cancel if job is in a cancellable state
            if job.status not in (
                noteflow_pb2.JOB_STATUS_QUEUED,
                noteflow_pb2.JOB_STATUS_RUNNING,
            ):
                response.success = False
                response.error_message = "Job already completed or failed"
                response.status = int(job.status)
                return response

            await repo.diarization_jobs.update_status(
                job_id,
                noteflow_pb2.JOB_STATUS_CANCELLED,
                error_message=ERR_CANCELLED_BY_USER,
            )
            await repo.commit()

        logger.info("Cancelled diarization job %s", job_id)
        response.success = True
        response.status = noteflow_pb2.JOB_STATUS_CANCELLED
        return response

    async def GetActiveDiarizationJobs(
        self: DiarizationJobServicer,
        request: noteflow_pb2.GetActiveDiarizationJobsRequest,
        context: GrpcContext,
    ) -> noteflow_pb2.GetActiveDiarizationJobsResponse:
        """Return all active (QUEUED or RUNNING) diarization jobs.

        Used by clients to restore polling state after reconnection or app restart.
        Enables crash recovery without losing track of in-flight jobs.
        """
        response = noteflow_pb2.GetActiveDiarizationJobsResponse()

        async with self.create_repository_provider() as repo:
            if not repo.supports_diarization_jobs:
                # Return empty list if DB not available
                return response

            active_jobs = await repo.diarization_jobs.get_all_active()

            for job in active_jobs:
                # Calculate progress percentage (time-based for running jobs)
                progress_percent = 0.0
                if job.status == noteflow_pb2.JOB_STATUS_RUNNING and job.started_at is not None:
                    now = utc_now()
                    started: datetime = job.started_at
                    elapsed = (now - started).total_seconds()
                    audio_duration = job.audio_duration_seconds
                    if audio_duration is not None and audio_duration > 0:
                        # ~10 seconds processing per 60 seconds audio
                        estimated_duration = audio_duration * 0.17
                        progress_percent = min(95.0, (elapsed / estimated_duration) * 100)
                    else:
                        # Fallback: assume 2 minutes total
                        progress_percent = min(95.0, (elapsed / 120) * 100)

                job_status = noteflow_pb2.DiarizationJobStatus(
                    job_id=job.job_id,
                    status=int(job.status),
                    segments_updated=job.segments_updated,
                    speaker_ids=job.speaker_ids,
                    error_message=job.error_message,
                    progress_percent=progress_percent,
                )
                response.jobs.append(job_status)

        logger.debug("Returning %d active diarization jobs", len(response.jobs))
        return response
