"""Transcription repository for Trax platform.

This module provides a protocol-based repository for managing transcription jobs
and results in the database.
"""

import logging
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
from uuid import UUID

from ..database.connection import get_db_session
from ..database.models import TranscriptionJob, TranscriptionResult, MediaFile

logger = logging.getLogger(__name__)


@runtime_checkable
class TranscriptionRepositoryProtocol(Protocol):
    """Protocol for transcription repository."""

    async def create_job(
        self, 
        media_file_id: UUID, 
        model_config: Dict[str, Any],
        processing_options: Dict[str, Any],
        priority: int = 0
    ) -> TranscriptionJob:
        """Create a new transcription job."""
        ...

    async def get_job(self, job_id: UUID) -> Optional[TranscriptionJob]:
        """Get a transcription job by ID."""
        ...

    async def get_jobs_by_media_file(self, media_file_id: UUID) -> List[TranscriptionJob]:
        """Get all transcription jobs for a media file."""
        ...

    async def get_pending_jobs(self, limit: int = 10) -> List[TranscriptionJob]:
        """Get pending transcription jobs."""
        ...

    async def update_job_status(
        self, 
        job_id: UUID, 
        status: str, 
        error_message: Optional[str] = None
    ) -> bool:
        """Update transcription job status."""
        ...

    async def update_job_progress(
        self, 
        job_id: UUID, 
        processing_time: float,
        started_at: Optional[str] = None,
        completed_at: Optional[str] = None
    ) -> bool:
        """Update transcription job progress."""
        ...

    async def create_result(
        self,
        job_id: UUID,
        media_file_id: UUID,
        content: Dict[str, Any],
        segments: Optional[List[Dict[str, Any]]] = None,
        confidence_scores: Optional[List[float]] = None,
        accuracy: Optional[float] = None,
        word_count: Optional[int] = None,
        processing_time: Optional[float] = None,
        model_used: Optional[str] = None,
        model_config: Optional[Dict[str, Any]] = None,
        pipeline_version: str = "v1"
    ) -> TranscriptionResult:
        """Create a transcription result."""
        ...

    async def get_result(self, result_id: UUID) -> Optional[TranscriptionResult]:
        """Get a transcription result by ID."""
        ...

    async def get_results_by_media_file(
        self, 
        media_file_id: UUID,
        pipeline_version: Optional[str] = None
    ) -> List[TranscriptionResult]:
        """Get transcription results for a media file."""
        ...

    async def get_latest_result(
        self, 
        media_file_id: UUID,
        pipeline_version: Optional[str] = None
    ) -> Optional[TranscriptionResult]:
        """Get the latest transcription result for a media file."""
        ...


class TranscriptionRepository:
    """Transcription repository implementation."""

    def __init__(self):
        pass

    async def create_job(
        self, 
        media_file_id: UUID, 
        model_config: Dict[str, Any],
        processing_options: Dict[str, Any],
        priority: int = 0
    ) -> TranscriptionJob:
        """Create a new transcription job."""
        try:
            with get_db_session() as session:
                job = TranscriptionJob(
                    media_file_id=media_file_id,
                    status="pending",
                    priority=priority,
                    model_config=model_config,
                    processing_options=processing_options
                )
                
                session.add(job)
                session.commit()
                session.refresh(job)
                
                logger.info(f"Created transcription job {job.id} for media file {media_file_id}")
                return job
        except Exception as e:
            logger.error(f"Error creating transcription job: {e}")
            raise

    async def get_job(self, job_id: UUID) -> Optional[TranscriptionJob]:
        """Get a transcription job by ID."""
        try:
            with get_db_session() as session:
                job = session.query(TranscriptionJob).filter(TranscriptionJob.id == job_id).first()
                return job
        except Exception as e:
            logger.error(f"Error getting transcription job: {e}")
            return None

    async def get_jobs_by_media_file(self, media_file_id: UUID) -> List[TranscriptionJob]:
        """Get all transcription jobs for a media file."""
        try:
            with get_db_session() as session:
                jobs = session.query(TranscriptionJob).filter(
                    TranscriptionJob.media_file_id == media_file_id
                ).order_by(TranscriptionJob.created_at.desc()).all()
                return jobs
        except Exception as e:
            logger.error(f"Error getting jobs by media file: {e}")
            return []

    async def get_pending_jobs(self, limit: int = 10) -> List[TranscriptionJob]:
        """Get pending transcription jobs."""
        try:
            with get_db_session() as session:
                jobs = session.query(TranscriptionJob).filter(
                    TranscriptionJob.status == "pending"
                ).order_by(TranscriptionJob.priority.desc(), TranscriptionJob.created_at.asc()).limit(limit).all()
                return jobs
        except Exception as e:
            logger.error(f"Error getting pending jobs: {e}")
            return []

    async def update_job_status(
        self, 
        job_id: UUID, 
        status: str, 
        error_message: Optional[str] = None
    ) -> bool:
        """Update transcription job status."""
        try:
            with get_db_session() as session:
                job = session.query(TranscriptionJob).filter(TranscriptionJob.id == job_id).first()
                if job:
                    job.status = status
                    if error_message:
                        job.error_message = error_message
                    session.commit()
                    return True
                return False
        except Exception as e:
            logger.error(f"Error updating job status: {e}")
            return False

    async def update_job_progress(
        self, 
        job_id: UUID, 
        processing_time: float,
        started_at: Optional[str] = None,
        completed_at: Optional[str] = None
    ) -> bool:
        """Update transcription job progress."""
        try:
            with get_db_session() as session:
                job = session.query(TranscriptionJob).filter(TranscriptionJob.id == job_id).first()
                if job:
                    job.processing_time = processing_time
                    if started_at:
                        job.started_at = started_at
                    if completed_at:
                        job.completed_at = completed_at
                    session.commit()
                    return True
                return False
        except Exception as e:
            logger.error(f"Error updating job progress: {e}")
            return False

    async def create_result(
        self,
        job_id: UUID,
        media_file_id: UUID,
        content: Dict[str, Any],
        segments: Optional[List[Dict[str, Any]]] = None,
        confidence_scores: Optional[List[float]] = None,
        accuracy: Optional[float] = None,
        word_count: Optional[int] = None,
        processing_time: Optional[float] = None,
        model_used: Optional[str] = None,
        model_config: Optional[Dict[str, Any]] = None,
        pipeline_version: str = "v1"
    ) -> TranscriptionResult:
        """Create a transcription result."""
        try:
            with get_db_session() as session:
                result = TranscriptionResult(
                    job_id=job_id,
                    media_file_id=media_file_id,
                    pipeline_version=pipeline_version,
                    content=content,
                    segments=segments or [],
                    confidence_scores=confidence_scores or [],
                    accuracy=accuracy,
                    word_count=word_count,
                    processing_time=processing_time,
                    model_used=model_used,
                    model_config=model_config or {}
                )
                
                session.add(result)
                session.commit()
                session.refresh(result)
                
                logger.info(f"Created transcription result {result.id} for job {job_id}")
                return result
        except Exception as e:
            logger.error(f"Error creating transcription result: {e}")
            raise

    async def get_result(self, result_id: UUID) -> Optional[TranscriptionResult]:
        """Get a transcription result by ID."""
        try:
            with get_db_session() as session:
                result = session.query(TranscriptionResult).filter(TranscriptionResult.id == result_id).first()
                return result
        except Exception as e:
            logger.error(f"Error getting transcription result: {e}")
            return None

    async def get_results_by_media_file(
        self, 
        media_file_id: UUID,
        pipeline_version: Optional[str] = None
    ) -> List[TranscriptionResult]:
        """Get transcription results for a media file."""
        try:
            with get_db_session() as session:
                query = session.query(TranscriptionResult).filter(
                    TranscriptionResult.media_file_id == media_file_id
                )
                
                if pipeline_version:
                    query = query.filter(TranscriptionResult.pipeline_version == pipeline_version)
                
                results = query.order_by(TranscriptionResult.created_at.desc()).all()
                return results
        except Exception as e:
            logger.error(f"Error getting results by media file: {e}")
            return []

    async def get_latest_result(
        self, 
        media_file_id: UUID,
        pipeline_version: Optional[str] = None
    ) -> Optional[TranscriptionResult]:
        """Get the latest transcription result for a media file."""
        try:
            with get_db_session() as session:
                query = session.query(TranscriptionResult).filter(
                    TranscriptionResult.media_file_id == media_file_id
                )
                
                if pipeline_version:
                    query = query.filter(TranscriptionResult.pipeline_version == pipeline_version)
                
                result = query.order_by(TranscriptionResult.created_at.desc()).first()
                return result
        except Exception as e:
            logger.error(f"Error getting latest result: {e}")
            return None


async def create_transcription_repository() -> TranscriptionRepositoryProtocol:
    """Create a transcription repository instance."""
    return TranscriptionRepository()
