"""NoteFlow gRPC service implementation (async with UoW)."""

from __future__ import annotations

import asyncio
import contextlib
import time
from collections import deque
from pathlib import Path
from typing import TYPE_CHECKING, ClassVar, Final

from noteflow import __version__
from noteflow.config.constants import APP_DIR_NAME
from noteflow.config.constants import DEFAULT_SAMPLE_RATE as _DEFAULT_SAMPLE_RATE
from noteflow.domain.entities import Meeting
from noteflow.domain.value_objects import MeetingState
from noteflow.grpc.meeting_store import MeetingStore
from noteflow.infrastructure.asr import Segmenter, SegmenterConfig, StreamingVad
from noteflow.infrastructure.audio.partial_buffer import PartialAudioBuffer
from noteflow.infrastructure.audio.writer import MeetingAudioWriter
from noteflow.infrastructure.diarization import DiarizationSession
from noteflow.infrastructure.logging import get_logger
from noteflow.infrastructure.persistence.memory import MemoryUnitOfWork
from noteflow.infrastructure.persistence.repositories import DiarizationJob
from noteflow.infrastructure.persistence.unit_of_work import SqlAlchemyUnitOfWork
from noteflow.infrastructure.security.crypto import AesGcmCryptoBox
from noteflow.infrastructure.security.keystore import KeyringKeyStore

from ._config import ServicesConfig
from ._mixins import (
    AnnotationMixin,
    CalendarMixin,
    DiarizationJobMixin,
    DiarizationMixin,
    EntitiesMixin,
    ExportMixin,
    MeetingMixin,
    ObservabilityMixin,
    OidcMixin,
    PreferencesMixin,
    ProjectMembershipMixin,
    ProjectMixin,
    StreamingMixin,
    SummarizationMixin,
    SyncMixin,
    WebhooksMixin,
)
from .proto import noteflow_pb2, noteflow_pb2_grpc
from .stream_state import MeetingStreamState

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

    from noteflow.infrastructure.asr import FasterWhisperEngine
    from noteflow.infrastructure.auth.oidc_registry import OidcAuthService
    from noteflow.infrastructure.diarization import SpeakerTurn

    from ._mixins._types import GrpcContext

logger = get_logger(__name__)

if TYPE_CHECKING:
    _GrpcBaseServicer = object
else:
    _GrpcBaseServicer = noteflow_pb2_grpc.NoteFlowServiceServicer


class NoteFlowServicer(
    StreamingMixin,
    DiarizationMixin,
    DiarizationJobMixin,
    MeetingMixin,
    SummarizationMixin,
    AnnotationMixin,
    ExportMixin,
    EntitiesMixin,
    CalendarMixin,
    WebhooksMixin,
    SyncMixin,
    ObservabilityMixin,
    PreferencesMixin,
    OidcMixin,
    ProjectMixin,
    ProjectMembershipMixin,
    _GrpcBaseServicer,
):
    """Async gRPC service implementation for NoteFlow with PostgreSQL persistence.

    Type stubs for mixin methods are defined to fix type inference when mixins
    use `self: Protocol` annotations.
    """

    # Type stubs for mixin methods (fixes type inference when mixins use `self: Protocol`)
    if TYPE_CHECKING:
        # StreamingMixin (test_streaming_real_pipeline.py, test_e2e_streaming.py)
        def StreamTranscription(
            self,
            request_iterator: AsyncIterator[noteflow_pb2.AudioChunk],
            context: GrpcContext,
        ) -> AsyncIterator[noteflow_pb2.TranscriptUpdate]: ...

        # CalendarMixin (test_oauth.py)
        async def GetCalendarProviders(
            self,
            request: noteflow_pb2.GetCalendarProvidersRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.GetCalendarProvidersResponse: ...

        async def InitiateOAuth(
            self,
            request: noteflow_pb2.InitiateOAuthRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.InitiateOAuthResponse: ...

        async def CompleteOAuth(
            self,
            request: noteflow_pb2.CompleteOAuthRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.CompleteOAuthResponse: ...

        async def GetOAuthConnectionStatus(
            self,
            request: noteflow_pb2.GetOAuthConnectionStatusRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.GetOAuthConnectionStatusResponse: ...

        async def DisconnectOAuth(
            self,
            request: noteflow_pb2.DisconnectOAuthRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.DisconnectOAuthResponse: ...

        # Type stubs for SummarizationMixin methods (test_cloud_consent.py, test_generate_summary.py)
        async def GetCloudConsentStatus(
            self,
            request: noteflow_pb2.GetCloudConsentStatusRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.GetCloudConsentStatusResponse: ...

        async def GrantCloudConsent(
            self,
            request: noteflow_pb2.GrantCloudConsentRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.GrantCloudConsentResponse: ...

        async def RevokeCloudConsent(
            self,
            request: noteflow_pb2.RevokeCloudConsentRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.RevokeCloudConsentResponse: ...

        async def GenerateSummary(
            self,
            request: noteflow_pb2.GenerateSummaryRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.Summary: ...

        # Type stubs for SyncMixin methods (test_sync_orchestration.py)
        async def StartIntegrationSync(
            self,
            request: noteflow_pb2.StartIntegrationSyncRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.StartIntegrationSyncResponse: ...

        async def GetSyncStatus(
            self,
            request: noteflow_pb2.GetSyncStatusRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.GetSyncStatusResponse: ...

        async def ListSyncHistory(
            self,
            request: noteflow_pb2.ListSyncHistoryRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.ListSyncHistoryResponse: ...

        async def GetUserIntegrations(
            self,
            request: noteflow_pb2.GetUserIntegrationsRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.GetUserIntegrationsResponse: ...

        # Type stubs for DiarizationMixin methods (test_diarization_mixin.py, test_diarization_refine.py)
        async def RefineSpeakerDiarization(
            self,
            request: noteflow_pb2.RefineSpeakerDiarizationRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.RefineSpeakerDiarizationResponse: ...

        # Type stubs for SpeakerMixin methods (test_diarization_mixin.py)
        async def RenameSpeaker(
            self,
            request: noteflow_pb2.RenameSpeakerRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.RenameSpeakerResponse: ...

        # Type stubs for DiarizationJobMixin methods (test_diarization_mixin.py, test_diarization_cancel.py)
        async def GetDiarizationJobStatus(
            self,
            request: noteflow_pb2.GetDiarizationJobStatusRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.DiarizationJobStatus: ...

        async def CancelDiarizationJob(
            self,
            request: noteflow_pb2.CancelDiarizationJobRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.CancelDiarizationJobResponse: ...

        # Type stubs for WebhooksMixin methods (test_webhooks_mixin.py)
        async def RegisterWebhook(
            self,
            request: noteflow_pb2.RegisterWebhookRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.WebhookConfigProto: ...

        async def ListWebhooks(
            self,
            request: noteflow_pb2.ListWebhooksRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.ListWebhooksResponse: ...

        async def UpdateWebhook(
            self,
            request: noteflow_pb2.UpdateWebhookRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.WebhookConfigProto: ...

        async def DeleteWebhook(
            self,
            request: noteflow_pb2.DeleteWebhookRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.DeleteWebhookResponse: ...

        async def GetWebhookDeliveries(
            self,
            request: noteflow_pb2.GetWebhookDeliveriesRequest,
            context: GrpcContext,
        ) -> noteflow_pb2.GetWebhookDeliveriesResponse: ...

    VERSION: Final[str] = __version__
    MAX_CHUNK_SIZE: Final[int] = 1024 * 1024  # 1MB
    DEFAULT_SAMPLE_RATE: Final[int] = _DEFAULT_SAMPLE_RATE
    SUPPORTED_SAMPLE_RATES: ClassVar[list[int]] = [16000, 44100, 48000]
    PARTIAL_CADENCE_SECONDS: Final[float] = 2.0  # Emit partials every 2 seconds
    MIN_PARTIAL_AUDIO_SECONDS: Final[float] = 0.5  # Minimum audio for partial inference

    # Server state version for cache invalidation (Sprint GAP-002: State Synchronization).
    # Increment this when making changes that require clients to invalidate their caches.
    STATE_VERSION: Final[int] = 1

    def __init__(
        self,
        asr_engine: FasterWhisperEngine | None = None,
        session_factory: async_sessionmaker[AsyncSession] | None = None,
        meetings_dir: Path | None = None,
        services: ServicesConfig | None = None,
    ) -> None:
        """Initialize the service.

        Args:
            asr_engine: Optional ASR engine.
            session_factory: Optional async session factory for database persistence.
                           If not provided, falls back to in-memory MeetingStore.
            meetings_dir: Optional directory for meeting audio storage.
                         Defaults to ~/.noteflow/meetings.
            services: Optional services configuration grouping all optional services.
        """
        # Injected services
        self.asr_engine = asr_engine
        self.session_factory = session_factory
        services = services or ServicesConfig()
        self.summarization_service = services.summarization_service
        self.diarization_engine = services.diarization_engine
        self.diarization_refinement_enabled = services.diarization_refinement_enabled
        self.ner_service = services.ner_service
        self.calendar_service = services.calendar_service
        self.webhook_service = services.webhook_service
        self.project_service = services.project_service
        self._start_time = time.time()
        self.memory_store: MeetingStore | None = MeetingStore() if session_factory is None else None
        # Audio infrastructure
        self.meetings_dir = meetings_dir or (Path.home() / APP_DIR_NAME / "meetings")
        self._keystore = KeyringKeyStore()
        self.crypto = AesGcmCryptoBox(self._keystore)
        self.audio_writers: dict[str, MeetingAudioWriter] = {}
        # Per-meeting streaming state
        self.vad_instances: dict[str, StreamingVad] = {}
        self.segmenters: dict[str, Segmenter] = {}
        self.was_speaking: dict[str, bool] = {}
        self.segment_counters: dict[str, int] = {}
        self.stream_formats: dict[str, tuple[int, int]] = {}
        self.active_streams: set[str] = set()
        self.stop_requested: set[str] = set()
        self.chunk_sequences: dict[str, int] = {}
        self.chunk_counts: dict[str, int] = {}
        self.chunk_receipt_times: dict[str, deque[float]] = {}
        self.pending_chunks: dict[str, int] = {}
        self.partial_buffers: dict[str, PartialAudioBuffer] = {}
        self.last_partial_time: dict[str, float] = {}
        self.last_partial_text: dict[str, str] = {}
        self.audio_write_failed: set[str] = set()
        self.stream_states: dict[str, MeetingStreamState] = {}
        # Diarization state
        self.diarization_turns: dict[str, list[SpeakerTurn]] = {}
        self.diarization_stream_time: dict[str, float] = {}
        self.diarization_streaming_failed: set[str] = set()
        self.diarization_sessions: dict[str, DiarizationSession] = {}
        self.diarization_jobs: dict[str, DiarizationJob] = {}
        self.diarization_tasks: dict[str, asyncio.Task[None]] = {}
        self.diarization_lock = asyncio.Lock()
        self.stream_init_lock = asyncio.Lock()
        self.oidc_service: OidcAuthService | None = None

    def use_database(self) -> bool:
        """Check if database persistence is configured."""
        return self.session_factory is not None

    def get_memory_store(self) -> MeetingStore:
        """Get the in-memory store, raising if not configured."""
        if self.memory_store is None:
            raise RuntimeError("Memory store not configured")
        return self.memory_store

    def create_uow(self) -> SqlAlchemyUnitOfWork:
        """Create a new Unit of Work (database-backed)."""
        if self.session_factory is None:
            raise RuntimeError("Database not configured")
        return SqlAlchemyUnitOfWork(self.session_factory, self.meetings_dir)

    def create_repository_provider(self) -> SqlAlchemyUnitOfWork | MemoryUnitOfWork:
        """Create a repository provider (database or memory backed).

        Returns a UnitOfWork implementation appropriate for the current
        configuration. Use this for operations that can work with either
        backend, eliminating the need for if/else branching.

        Returns:
            SqlAlchemyUnitOfWork if database configured, MemoryUnitOfWork otherwise.
        """
        if self.session_factory is not None:
            return SqlAlchemyUnitOfWork(self.session_factory, self.meetings_dir)
        return MemoryUnitOfWork(self.get_memory_store())

    def init_streaming_state(self, meeting_id: str, next_segment_id: int) -> None:
        """Initialize VAD, Segmenter, speaking state, and partial buffers for a meeting."""
        # Create core components
        vad = StreamingVad()
        segmenter = Segmenter(config=SegmenterConfig(sample_rate=self.DEFAULT_SAMPLE_RATE))
        partial_buffer = PartialAudioBuffer(sample_rate=self.DEFAULT_SAMPLE_RATE)
        current_time = time.time()

        # Create consolidated state (single lookup in hot paths)
        state = MeetingStreamState(
            vad=vad,
            segmenter=segmenter,
            partial_buffer=partial_buffer,
            sample_rate=self.DEFAULT_SAMPLE_RATE,
            channels=1,
            next_segment_id=next_segment_id,
            was_speaking=False,
            last_partial_time=current_time,
            last_partial_text="",
            diarization_session=None,
            diarization_turns=[],
            diarization_stream_time=0.0,
            diarization_streaming_failed=False,
            is_active=True,
            stop_requested=False,
            audio_write_failed=False,
        )
        self.stream_states[meeting_id] = state

        # Also populate legacy dicts for backward compatibility during migration
        self.vad_instances[meeting_id] = vad
        self.segmenters[meeting_id] = segmenter
        self.was_speaking[meeting_id] = False
        self.segment_counters[meeting_id] = next_segment_id
        self.partial_buffers[meeting_id] = partial_buffer
        self.last_partial_time[meeting_id] = current_time
        self.last_partial_text[meeting_id] = ""
        self.diarization_turns[meeting_id] = state.diarization_turns  # Share reference
        self.diarization_stream_time[meeting_id] = 0.0
        self.diarization_streaming_failed.discard(meeting_id)
        # NOTE: Per-meeting diarization sessions are created lazily in
        # _process_streaming_diarization() to avoid blocking on model load

    def cleanup_streaming_state(self, meeting_id: str) -> None:
        """Clean up VAD, Segmenter, speaking state, and partial buffers for a meeting."""
        # Clean up consolidated state
        if (state := self.stream_states.pop(meeting_id, None)) and state.diarization_session is not None:
            state.diarization_session.close()

        # Clean up legacy dicts (backward compatibility)
        self.vad_instances.pop(meeting_id, None)
        self.segmenters.pop(meeting_id, None)
        self.was_speaking.pop(meeting_id, None)
        self.segment_counters.pop(meeting_id, None)
        self.stream_formats.pop(meeting_id, None)
        self.partial_buffers.pop(meeting_id, None)
        self.last_partial_time.pop(meeting_id, None)
        self.last_partial_text.pop(meeting_id, None)
        self.diarization_turns.pop(meeting_id, None)
        self.diarization_stream_time.pop(meeting_id, None)
        self.diarization_streaming_failed.discard(meeting_id)

        # Clean up chunk sequence tracking
        self.chunk_sequences.pop(meeting_id, None)
        self.chunk_counts.pop(meeting_id, None)

        # Clean up congestion tracking (Phase 3)
        if hasattr(self, "_chunk_receipt_times"):
            self.chunk_receipt_times.pop(meeting_id, None)
        if hasattr(self, "_pending_chunks"):
            self.pending_chunks.pop(meeting_id, None)

        # Clean up per-meeting diarization session (legacy path)
        if session := self.diarization_sessions.pop(meeting_id, None):
            session.close()

    def get_stream_state(self, meeting_id: str) -> MeetingStreamState | None:
        """Get consolidated streaming state for a meeting.

        Returns None if meeting has no active stream state.
        Single lookup replaces 13+ dict accesses in hot paths.
        """
        return self.stream_states.get(meeting_id)

    def ensure_meeting_dek(self, meeting: Meeting) -> tuple[bytes, bytes, bool]:
        """Ensure meeting has a DEK, generating one if needed.

        Args:
            meeting: Meeting entity.

        Returns:
            Tuple of (dek, wrapped_dek, needs_update).
        """
        if meeting.wrapped_dek is None:
            dek = self.crypto.generate_dek()
            wrapped_dek = self.crypto.wrap_dek(dek)
            meeting.wrapped_dek = wrapped_dek
            return dek, wrapped_dek, True
        wrapped_dek = meeting.wrapped_dek
        dek = self.crypto.unwrap_dek(wrapped_dek)
        return dek, wrapped_dek, False

    def start_meeting_if_needed(self, meeting: Meeting) -> tuple[bool, str | None]:
        """Start recording on meeting if not already recording.

        Args:
            meeting: Meeting entity.

        Returns:
            Tuple of (needs_update, error_message).
        """
        if meeting.state == MeetingState.RECORDING:
            return False, None
        try:
            meeting.start_recording()
            return True, None
        except ValueError as e:
            return False, str(e)

    def open_meeting_audio_writer(
        self,
        meeting_id: str,
        dek: bytes,
        wrapped_dek: bytes,
        asset_path: str | None = None,
    ) -> None:
        """Open audio writer for a meeting.

        Args:
            meeting_id: Meeting ID string.
            dek: Data encryption key.
            wrapped_dek: Wrapped DEK.
            asset_path: Relative path for audio storage (defaults to meeting_id).
        """
        writer = MeetingAudioWriter(self.crypto, self.meetings_dir)
        writer.open(
            meeting_id=meeting_id,
            dek=dek,
            wrapped_dek=wrapped_dek,
            sample_rate=self.DEFAULT_SAMPLE_RATE,
            asset_path=asset_path,
        )
        self.audio_writers[meeting_id] = writer
        logger.info("Audio writer opened for meeting %s", meeting_id)

    def close_audio_writer(self, meeting_id: str) -> None:
        """Close and remove the audio writer for a meeting."""
        # Clean up write failure tracking
        self.audio_write_failed.discard(meeting_id)

        if meeting_id not in self.audio_writers:
            return

        try:
            writer = self.audio_writers.pop(meeting_id)
            writer.close()
            logger.info(
                "Audio writer closed for meeting %s: %d bytes written",
                meeting_id,
                writer.bytes_written,
            )
        except (OSError, RuntimeError) as e:
            logger.error(
                "Failed to close audio writer for meeting %s: %s",
                meeting_id,
                e,
            )

    def next_segment_id(self, meeting_id: str, fallback: int = 0) -> int:
        """Get and increment the next segment id for a meeting."""
        next_id = self.segment_counters.get(meeting_id)
        if next_id is None:
            next_id = fallback
        self.segment_counters[meeting_id] = next_id + 1
        return next_id

    async def _count_active_meetings_db(self) -> int:
        """Count active meetings using database state."""
        async with self.create_uow() as uow:
            total = 0
            for state in (MeetingState.RECORDING, MeetingState.STOPPING):
                total += await uow.meetings.count_by_state(state)
            return total

    async def GetServerInfo(
        self,
        request: noteflow_pb2.ServerInfoRequest,
        context: GrpcContext,
    ) -> noteflow_pb2.ServerInfo:
        """Get server information."""
        asr_model = ""
        asr_ready = False
        if self.asr_engine:
            asr_ready = self.asr_engine.is_loaded
            asr_model = self.asr_engine.model_size or ""

        diarization_enabled = self.diarization_engine is not None
        diarization_ready = self.diarization_engine is not None and (
            self.diarization_engine.is_streaming_loaded
            or self.diarization_engine.is_offline_loaded
        )

        if self.session_factory is not None:
            active = await self._count_active_meetings_db()
        else:
            active = self.get_memory_store().active_count

        return noteflow_pb2.ServerInfo(
            version=self.VERSION,
            asr_model=asr_model,
            asr_ready=asr_ready,
            supported_sample_rates=self.SUPPORTED_SAMPLE_RATES,
            max_chunk_size=self.MAX_CHUNK_SIZE,
            uptime_seconds=time.time() - self._start_time,
            active_meetings=active,
            diarization_enabled=diarization_enabled,
            diarization_ready=diarization_ready,
            state_version=self.STATE_VERSION,
        )

    async def shutdown(self) -> None:
        """Clean up servicer state before server stops.

        Cancel in-flight diarization tasks, close audio writers, and mark
        any running jobs as failed in the database.
        """
        logger.info("Shutting down servicer...")

        # Cancel in-flight diarization tasks
        for job_id, task in list(self.diarization_tasks.items()):
            if not task.done():
                logger.debug("Cancelling diarization task %s", job_id)
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task

        self.diarization_tasks.clear()

        # Close all diarization sessions
        for meeting_id, session in list(self.diarization_sessions.items()):
            logger.debug("Closing diarization session for meeting %s", meeting_id)
            session.close()
        self.diarization_sessions.clear()

        # Close all audio writers
        for meeting_id in list(self.audio_writers.keys()):
            logger.debug("Closing audio writer for meeting %s", meeting_id)
            self.close_audio_writer(meeting_id)

        # Mark running jobs as FAILED in database
        if self.session_factory is not None:
            async with self.create_uow() as uow:
                failed_count = await uow.diarization_jobs.mark_running_as_failed()
                await uow.commit()
                if failed_count > 0:
                    logger.warning(
                        "Marked %d running diarization jobs as failed on shutdown",
                        failed_count,
                    )

        # Close webhook service HTTP client
        if self.webhook_service is not None:
            logger.debug("Closing webhook service HTTP client")
            await self.webhook_service.close()

        logger.info("Servicer shutdown complete")
