"""Trax-specific service base classes built on AI Assistant Library."""

import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field

# Import from the library - for now, use simplified base classes
# The full library integration will need proper packaging
from enum import Enum
from pathlib import Path
from typing import Any, Dict, Optional, Protocol, runtime_checkable


class ServiceStatus(Enum):
    """Service status values."""

    INITIALIZING = "initializing"
    HEALTHY = "healthy"
    DEGRADED = "degraded"
    UNHEALTHY = "unhealthy"
    SHUTTING_DOWN = "shutting_down"
    STOPPED = "stopped"


class BaseService(ABC):
    """Simplified base service class."""

    def __init__(self, name: str, config: Optional[Dict[str, Any]] = None):
        self.name = name
        self.config = config or {}
        self._status = ServiceStatus.STOPPED

    @property
    def status(self) -> ServiceStatus:
        return self._status

    async def initialize(self) -> None:
        self._status = ServiceStatus.INITIALIZING
        await self._initialize_impl()
        self._status = ServiceStatus.HEALTHY

    async def shutdown(self) -> None:
        self._status = ServiceStatus.SHUTTING_DOWN
        await self._shutdown_impl()
        self._status = ServiceStatus.STOPPED

    @abstractmethod
    async def _initialize_impl(self) -> None:
        pass

    async def _shutdown_impl(self) -> None:
        pass

    def get_health_status(self) -> Dict[str, Any]:
        return {
            "status": self._status.value,
            "is_healthy": self._status == ServiceStatus.HEALTHY,
            "service": self.name,
        }


@dataclass
class AIModelConfig:
    """AI model configuration."""

    model_name: str
    api_key: Optional[str] = None
    temperature: float = 0.0
    max_tokens: int = 4096


@dataclass
class AIRequest:
    """AI service request."""

    prompt: str
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None


@dataclass
class AIResponse:
    """AI service response."""

    content: str
    model: str
    metadata: Dict[str, Any]


class BaseAIService(BaseService):
    """Simplified AI service base class."""

    def __init__(
        self, name: str, model_config: AIModelConfig, config: Optional[Dict[str, Any]] = None
    ):
        super().__init__(name, config)
        self.model_config = model_config

    async def process_request(self, request: AIRequest) -> AIResponse:
        """Process AI request - to be implemented by subclasses."""
        raise NotImplementedError


@dataclass
class RetryConfig:
    """Retry configuration."""

    max_attempts: int = 3
    initial_delay: float = 1.0
    backoff_factor: float = 2.0
    max_delay: float = 60.0
    jitter: bool = True
    retryable_exceptions: tuple = field(default_factory=lambda: (Exception,))
    non_retryable_exceptions: tuple = field(default_factory=tuple)


# Simplified retry decorator
def async_retry(max_attempts: int = 3, backoff_factor: float = 2.0):
    """Simplified async retry decorator."""

    def decorator(func):
        async def wrapper(*args, **kwargs):
            import asyncio

            for attempt in range(max_attempts):
                try:
                    return await func(*args, **kwargs)
                except Exception:
                    if attempt == max_attempts - 1:
                        raise
                    await asyncio.sleep(backoff_factor**attempt)
            return None

        return wrapper

    return decorator


retry = async_retry  # Alias for compatibility

logger = logging.getLogger(__name__)


class TraxService(BaseService):
    """Base service for all Trax services with protocol support.

    Extends the AI Assistant Library's BaseService with Trax-specific
    functionality for media processing pipelines.
    """

    def __init__(self, name: str, config: Optional[Dict[str, Any]] = None):
        """Initialize Trax service.

        Args:
            name: Service name for identification
            config: Service configuration dictionary

        """
        super().__init__(name, config)
        self.pipeline_version = config.get("pipeline_version", "v1") if config else "v1"

    async def _initialize_impl(self) -> None:
        """Trax-specific initialization."""
        logger.info(f"Initializing {self.name} service (pipeline {self.pipeline_version})")
        # Override in subclasses for specific initialization

    def get_pipeline_info(self) -> Dict[str, Any]:
        """Get information about the pipeline version this service supports."""
        return {
            "service": self.name,
            "version": self.pipeline_version,
            "status": self.status.value,
            "capabilities": self._get_capabilities(),
        }

    def _get_capabilities(self) -> list:
        """Get service capabilities. Override in subclasses."""
        return []


@runtime_checkable
class TranscriptionProtocol(Protocol):
    """Protocol for transcription services."""

    async def transcribe(self, audio_path: Path) -> Dict[str, Any]:
        """Transcribe audio file to text."""
        ...

    def can_handle(self, audio_path: Path) -> bool:
        """Check if this service can handle the given audio file."""
        ...


class TranscriptionService(TraxService):
    """Base transcription service with protocol support.

    Provides common functionality for all transcription implementations
    (Whisper, cloud services, etc.).
    """

    def __init__(self, name: str = "TranscriptionService", config: Optional[Dict[str, Any]] = None):
        super().__init__(name, config)
        self.supported_formats = (
            config.get("supported_formats", [".mp3", ".wav", ".flac", ".aac", ".ogg", ".m4a"])
            if config
            else [".mp3", ".wav", ".flac", ".aac", ".ogg", ".m4a"]
        )

        # Retry configuration for transcription
        self.retry_config = RetryConfig(
            max_attempts=config.get("max_retries", 3) if config else 3,
            initial_delay=1.0,
            backoff_factor=2.0,
            retryable_exceptions=(Exception,),  # Will be refined in implementations
        )

    @async_retry(max_attempts=3)
    async def transcribe(self, audio_path: Path) -> Dict[str, Any]:
        """Transcribe audio file with retry logic.

        Args:
            audio_path: Path to audio file

        Returns:
            Transcript dictionary with text, segments, metadata

        Raises:
            TranscriptionError: If transcription fails after retries

        """
        if not self.can_handle(audio_path):
            raise ValueError(f"Unsupported audio format: {audio_path.suffix}")

        return await self._transcribe_impl(audio_path)

    async def _transcribe_impl(self, audio_path: Path) -> Dict[str, Any]:
        """Implementation-specific transcription logic.

        Override this in concrete implementations.
        """
        raise NotImplementedError("Transcription implementation required")

    def can_handle(self, audio_path: Path) -> bool:
        """Check if this service can handle the audio file.

        Args:
            audio_path: Path to audio file

        Returns:
            True if the file format is supported

        """
        return audio_path.suffix.lower() in self.supported_formats

    def _get_capabilities(self) -> list:
        """Get transcription service capabilities."""
        return [
            f"formats:{','.join(self.supported_formats)}",
            f"pipeline:{self.pipeline_version}",
            f"retry:{self.retry_config.max_attempts}",
        ]


class EnhancementService(BaseAIService):
    """AI enhancement service for transcript improvement.

    Uses the AI Assistant Library's BaseAIService for robust AI interactions
    with built-in retry logic and error handling.
    """

    def __init__(self, name: str = "EnhancementService", config: Optional[Dict[str, Any]] = None):
        """Initialize enhancement service.

        Args:
            name: Service name
            config: Service configuration including AI model settings

        """
        # Configure AI model
        model_config = AIModelConfig(
            model_name=config.get("model", "deepseek-chat") if config else "deepseek-chat",
            api_key=config.get("api_key") if config else None,
            temperature=config.get("temperature", 0.0) if config else 0.0,
            max_tokens=config.get("max_tokens", 4096) if config else 4096,
        )

        super().__init__(name, model_config, config)

        # Enhancement-specific settings
        self.enhancement_prompt = config.get("enhancement_prompt") if config else None
        self.quality_threshold = config.get("quality_threshold", 0.7) if config else 0.7

    async def enhance_transcript(self, transcript: str, **kwargs) -> Dict[str, Any]:
        """Enhance transcript using AI.

        Args:
            transcript: Raw transcript text
            **kwargs: Additional enhancement parameters

        Returns:
            Enhanced transcript with metadata

        """
        # Build enhancement request
        prompt = self._build_enhancement_prompt(transcript, **kwargs)

        request = AIRequest(
            prompt=prompt,
            temperature=kwargs.get("temperature", self.model_config.temperature),
            max_tokens=kwargs.get("max_tokens", self.model_config.max_tokens),
        )

        # Process with retry logic from base class
        response = await self.process_request(request)

        return {
            "enhanced_text": response.content,
            "model": response.model,
            "confidence": self._calculate_confidence(transcript, response.content),
            "improvements": self._identify_improvements(transcript, response.content),
            "metadata": response.metadata,
        }

    def _build_enhancement_prompt(self, transcript: str, **kwargs) -> str:
        """Build the enhancement prompt.

        Args:
            transcript: Raw transcript
            **kwargs: Additional context

        Returns:
            Formatted prompt for AI enhancement

        """
        if self.enhancement_prompt:
            return self.enhancement_prompt.format(transcript=transcript, **kwargs)

        # Default enhancement prompt
        return f"""Fix the following transcript issues:
1. Correct punctuation and capitalization
2. Fix technical terms and proper nouns
3. Format into readable paragraphs
4. Preserve timestamps and speaker markers if present
5. Maintain original meaning

Transcript:
{transcript}

Return ONLY the enhanced transcript without any explanations."""

    def _calculate_confidence(self, original: str, enhanced: str) -> float:
        """Calculate confidence score for enhancement.

        Args:
            original: Original transcript
            enhanced: Enhanced transcript

        Returns:
            Confidence score between 0 and 1

        """
        # Simple length-based confidence (can be improved)
        if not enhanced:
            return 0.0

        length_ratio = len(enhanced) / len(original) if original else 1.0

        # Good enhancements should be similar in length
        if 0.8 <= length_ratio <= 1.3:
            return 0.95
        elif 0.5 <= length_ratio <= 2.0:
            return 0.7
        else:
            return 0.5

    def _identify_improvements(self, original: str, enhanced: str) -> list:
        """Identify what improvements were made.

        Args:
            original: Original transcript
            enhanced: Enhanced transcript

        Returns:
            List of improvement types detected

        """
        improvements = []

        # Check for punctuation improvements
        if enhanced.count(".") > original.count("."):
            improvements.append("punctuation")

        # Check for capitalization
        if enhanced[0].isupper() and not original[0].isupper():
            improvements.append("capitalization")

        # Check for paragraph formatting
        if "\n\n" in enhanced and "\n\n" not in original:
            improvements.append("formatting")

        return improvements


class BatchProcessingService(TraxService):
    """Service for batch processing of media files.

    Handles parallel processing, progress tracking, and error recovery.
    """

    def __init__(
        self, name: str = "BatchProcessingService", config: Optional[Dict[str, Any]] = None
    ):
        super().__init__(name, config)
        self.max_parallel = config.get("max_parallel", 4) if config else 4
        self.batch_size = config.get("batch_size", 10) if config else 10

    async def process_batch(
        self, files: list[Path], processor: TranscriptionProtocol
    ) -> Dict[str, Any]:
        """Process a batch of files.

        Args:
            files: List of file paths to process
            processor: Transcription service to use

        Returns:
            Batch processing results with statistics

        """
        results = {"total": len(files), "successful": 0, "failed": 0, "results": [], "errors": []}

        # Process files (simplified - real implementation would use asyncio.gather)
        for file_path in files:
            try:
                if processor.can_handle(file_path):
                    result = await processor.transcribe(file_path)
                    results["successful"] += 1
                    results["results"].append({"file": str(file_path), "result": result})
                else:
                    results["failed"] += 1
                    results["errors"].append(
                        {"file": str(file_path), "error": "Unsupported format"}
                    )
            except Exception as e:
                results["failed"] += 1
                results["errors"].append({"file": str(file_path), "error": str(e)})

        return results

    def _get_capabilities(self) -> list:
        """Get batch processing capabilities."""
        return [
            f"max_parallel:{self.max_parallel}",
            f"batch_size:{self.batch_size}",
            f"pipeline:{self.pipeline_version}",
        ]
