"""Main enhancement service for transcript improvement."""

import asyncio
import logging
import time
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable

from ...base.services import BaseService
from ...config import config
from ...database.models import TranscriptionResult

from .api import DeepSeekAPIClient
from .cache import EnhancementCache, generate_cache_key
from .config import EnhancementConfig
from .errors import EnhancementError
from .models import EnhancementResult
from .prompts import build_enhancement_prompt
from .quality import calculate_confidence, identify_improvements

logger = logging.getLogger(__name__)


@runtime_checkable
class EnhancementServiceProtocol(Protocol):
    """Protocol for enhancement services."""

    async def initialize(self) -> None:
        """Initialize the enhancement service."""
        ...

    async def enhance_transcript(self, transcript: str, **kwargs) -> EnhancementResult:
        """Enhance a transcript."""
        ...

    async def enhance_transcript_batch(self, transcripts: List[str], **kwargs) -> List[EnhancementResult]:
        """Enhance multiple transcripts."""
        ...

    async def enhance_transcription_result(self, transcription_result: TranscriptionResult, **kwargs) -> EnhancementResult:
        """Enhance a transcription result from the database."""
        ...


class DeepSeekEnhancementService(BaseService):
    """DeepSeek AI enhancement service for transcript improvement."""

    def __init__(self, config: EnhancementConfig):
        """Initialize enhancement service."""
        super().__init__("DeepSeekEnhancementService")
        
        self.config = config
        self.config.validate()
        
        self.api_client: Optional[DeepSeekAPIClient] = None
        self.cache = EnhancementCache(
            enable_caching=config.enable_caching,
            cache_ttl=config.cache_ttl
        )
        self.is_initialized = False

    async def _initialize_impl(self) -> None:
        """Initialize the DeepSeek client."""
        if self.is_initialized:
            return

        try:
            if not config.DEEPSEEK_API_KEY:
                raise EnhancementError(
                    "DeepSeek API key not configured",
                    error_type="configuration_error"
                )

            self.api_client = DeepSeekAPIClient(
                api_key=config.DEEPSEEK_API_KEY,
                model=self.config.model
            )

            self.is_initialized = True
            logger.info("DeepSeek Enhancement Service initialized successfully")

        except Exception as e:
            logger.error(f"Failed to initialize enhancement service: {e}")
            raise EnhancementError(
                f"Failed to initialize enhancement service: {e}",
                error_type="initialization_error"
            )

    async def enhance_transcript(self, transcript: str, **kwargs) -> EnhancementResult:
        """Enhance a transcript using DeepSeek AI."""
        if not self.is_initialized:
            await self.initialize()

        start_time = time.time()

        try:
            # Check cache first
            cache_key = generate_cache_key(transcript, self.config.__dict__, kwargs)
            cached_result = self.cache.get(cache_key)
            if cached_result:
                return cached_result

            # Build enhancement prompt
            prompt = build_enhancement_prompt(transcript, self.config.__dict__, **kwargs)

            # Call DeepSeek API
            enhanced_text = await self.api_client.enhance_text(
                prompt,
                temperature=kwargs.get("temperature", self.config.temperature),
                max_tokens=kwargs.get("max_tokens", self.config.max_tokens),
                max_retries=self.config.max_retries,
                retry_delay=self.config.retry_delay
            )

            # Calculate confidence and identify improvements
            confidence_score = calculate_confidence(transcript, enhanced_text)
            improvements = identify_improvements(transcript, enhanced_text)

            # Validate quality threshold
            if confidence_score < self.config.quality_threshold:
                raise EnhancementError(
                    f"Enhancement quality below threshold: {confidence_score:.3f} < {self.config.quality_threshold}",
                    original_text=transcript,
                    error_type="quality_threshold_error"
                )

            # Create result
            result = EnhancementResult(
                original_text=transcript,
                enhanced_text=enhanced_text,
                confidence_score=confidence_score,
                improvements=improvements,
                processing_time=time.time() - start_time,
                model_used=self.config.model,
                metadata={
                    "temperature": kwargs.get("temperature", self.config.temperature),
                    "max_tokens": kwargs.get("max_tokens", self.config.max_tokens),
                    "cache_key": cache_key
                }
            )

            # Cache result
            self.cache.set(cache_key, result)

            logger.info(f"Enhanced transcript successfully: {confidence_score:.3f} confidence")
            return result

        except EnhancementError:
            raise
        except Exception as e:
            logger.error(f"Failed to enhance transcript: {e}")
            raise EnhancementError(
                f"Failed to enhance transcript: {e}",
                original_text=transcript,
                error_type="enhancement_error"
            )

    async def enhance_transcript_batch(self, transcripts: List[str], **kwargs) -> List[EnhancementResult]:
        """Enhance multiple transcripts in batch."""
        if not self.is_initialized:
            await self.initialize()

        # Process transcripts with limited concurrency
        semaphore = asyncio.Semaphore(3)  # Limit concurrent API calls
        
        async def enhance_single(transcript: str) -> EnhancementResult:
            async with semaphore:
                try:
                    return await self.enhance_transcript(transcript, **kwargs)
                except EnhancementError as e:
                    return e  # Return error instead of raising

        # Process all transcripts
        tasks = [enhance_single(transcript) for transcript in transcripts]
        batch_results = await asyncio.gather(*tasks, return_exceptions=True)

        # Convert exceptions to EnhancementError
        results = []
        for i, result in enumerate(batch_results):
            if isinstance(result, Exception) and not isinstance(result, EnhancementError):
                results.append(EnhancementError(
                    f"Batch enhancement failed: {result}",
                    original_text=transcripts[i],
                    error_type="batch_error"
                ))
            else:
                results.append(result)

        successful = len([r for r in results if not isinstance(r, EnhancementError)])
        logger.info(f"Batch enhancement completed: {successful}/{len(transcripts)} successful")
        return results

    async def enhance_transcription_result(self, transcription_result: TranscriptionResult, **kwargs) -> EnhancementResult:
        """Enhance a transcription result from the database."""
        if not self.is_initialized:
            await self.initialize()

        # Extract transcript content
        transcript = transcription_result.content

        # Add context from transcription metadata
        context = kwargs.get("context", "")
        if transcription_result.segments:
            context += f" Transcript has {len(transcription_result.segments)} segments."
        if transcription_result.confidence_scores:
            avg_confidence = sum(transcription_result.confidence_scores) / len(transcription_result.confidence_scores)
            context += f" Average confidence: {avg_confidence:.3f}"

        # Enhance with context
        result = await self.enhance_transcript(
            transcript,
            context=context,
            preserve_timestamps=self.config.preserve_timestamps,
            preserve_speaker_markers=self.config.preserve_speaker_markers,
            **kwargs
        )

        # Add transcription metadata
        result.metadata.update({
            "transcription_result_id": str(transcription_result.id),
            "pipeline_version": transcription_result.pipeline_version,
            "model_used_original": transcription_result.model_used,
            "segments": transcription_result.segments,
            "confidence_scores": transcription_result.confidence_scores,
            "accuracy": transcription_result.accuracy,
            "word_count": transcription_result.word_count
        })

        logger.info(f"Enhanced transcription result: {result.enhanced_text[:50]}...")
        return result

    def clear_cache(self) -> None:
        """Clear the enhancement cache."""
        self.cache.clear()

    def get_cache_stats(self) -> Dict[str, Any]:
        """Get cache statistics."""
        return self.cache.get_stats()


def create_enhancement_service(config: Optional[EnhancementConfig] = None) -> EnhancementServiceProtocol:
    """Create an enhancement service instance."""
    if config is None:
        config = EnhancementConfig()
    
    return DeepSeekEnhancementService(config)
