"""
Egrobots SSA Client

A Python client for the Egrobots Speech Sentiment Analysis service.
Supports both real-time streaming and batch processing modes.
"""
import asyncio
import websockets
import json
import base64
import time
from pathlib import Path
from pydub import AudioSegment
import io
from typing import Dict, Optional, Callable, Union
from dotenv import load_dotenv
import os

# Load environment variables from .env file
# Try multiple locations to support both local and Docker deployments
dotenv_paths = [
    Path(__file__).resolve().parents[2] / '.env',  # Project root (local dev)
    Path(__file__).resolve().parent / '.env',      # Same dir as client (Docker)
    Path.cwd() / '.env',                           # Current working directory
]

for dotenv_path in dotenv_paths:
    if dotenv_path.exists():
        load_dotenv(dotenv_path=dotenv_path)
        break


class EgrobotsSSAClient:
    """
    Client for Egrobots SSA service.

    Supports two modes:
    1. Real-time streaming: Stream audio chunks as they arrive (e.g., from microphone)
    2. Batch processing: Uploads complete audio file via HTTP POST
    """

    def __init__(self):
        """
        Initialize the Egrobots SSA client.
        Reads the websocket_url and http_url from environment variables.
        """
        self.websocket_url = os.getenv("WEBSOCKET_URL")
        self.http_url = os.getenv("BASE_URL")
        if not self.websocket_url or not self.http_url:
            raise ValueError("WEBSOCKET_URL and BASE_URL must be set in the .env file.")
        self.session_id = None
        self.chunk_count = 0
        self.websocket = None
        self.is_streaming = False

    async def start_streaming_session(self, progress_callback: Optional[Callable] = None):
        """
        Start a new streaming session.
        
        Args:
            progress_callback: Optional callback for progress updates
            
        Returns:
            True if session started successfully, False otherwise
        """
        try:
            # Connect to WebSocket
            self.websocket = await websockets.connect(
                self.websocket_url,
                ping_interval=None,
                close_timeout=10
            )
            
            if progress_callback:
                progress_callback('connected', {'url': self.websocket_url})
            
            # Start session
            await self.websocket.send(json.dumps({"action": "start"}))
            response = await self.websocket.recv()
            response_data = json.loads(response)
            
            if response_data.get("status") != "ready":
                if progress_callback:
                    progress_callback('error', {'message': str(response_data)})
                return False
            
            self.session_id = response_data.get("session_id")
            self.is_streaming = True
            self.chunk_count = 0
            
            if progress_callback:
                progress_callback('session_started', {
                    'session_id': self.session_id,
                    'mode': 'real-time streaming'
                })
            
            return True
            
        except Exception as e:
            if progress_callback:
                progress_callback('error', {'message': str(e)})
            return False
    
    async def send_audio_chunk(
        self,
        audio_data: bytes,
        audio_format: str = "mp3",
        progress_callback: Optional[Callable] = None
    ) -> bool:
        """
        Send a single audio chunk to the streaming session.
        
        This is the core method for real-time streaming - call this whenever
        you have a new audio chunk from your microphone or audio source.
        
        Args:
            audio_data: Raw audio bytes (in the specified format)
            audio_format: Audio format (mp3, wav, etc.)
            progress_callback: Optional callback for progress updates
            
        Returns:
            True if chunk was sent successfully, False otherwise
        """
        if not self.is_streaming or not self.websocket:
            if progress_callback:
                progress_callback('error', {'message': 'No active streaming session'})
            return False
        
        try:
            # Encode audio to base64
            chunk_b64 = base64.b64encode(audio_data).decode('utf-8')
            
            # Send chunk
            await self.websocket.send(json.dumps({
                "action": "chunk",
                "data": chunk_b64,
                "format": audio_format
            }))
            
            # Wait for acknowledgment
            response = await self.websocket.recv()
            response_data = json.loads(response)
            
            self.chunk_count += 1
            
            if progress_callback:
                progress_callback('chunk_sent', {
                    'chunk_number': self.chunk_count,
                    'elapsed_seconds': response_data.get('elapsed_time_seconds', 0),
                    'chunk_size_mb': len(audio_data) / (1024 * 1024)
                })
            
            return True
            
        except Exception as e:
            if progress_callback:
                progress_callback('error', {'message': str(e)})
            return False
    
    async def end_streaming_session(
        self,
        progress_callback: Optional[Callable] = None
    ) -> Dict:
        """
        End the streaming session and get analysis results.
        
        Args:
            progress_callback: Optional callback for progress updates
            
        Returns:
            Dict containing sentiment analysis results
        """
        if not self.is_streaming or not self.websocket:
            return {'error': 'No active streaming session'}
        
        try:
            processing_start = time.time()
            
            if progress_callback:
                progress_callback('streaming_complete', {
                    'total_chunks': self.chunk_count
                })
            
            # Send end signal
            await self.websocket.send(json.dumps({"action": "end"}))
            
            # Get processing message
            response = await self.websocket.recv()
            processing_msg = json.loads(response)
            
            if progress_callback:
                progress_callback('processing', {
                    'message': processing_msg.get('message', 'Processing...')
                })
            
            # Get final result
            response = await self.websocket.recv()
            result_data = json.loads(response)
            
            processing_time = time.time() - processing_start
            
            # Close websocket
            await self.websocket.close()
            self.is_streaming = False
            
            if result_data.get("status") == "complete":
                result = result_data.get("result", {})
                result['processing_time_seconds'] = processing_time
                result['session_id'] = self.session_id
                result['total_chunks'] = self.chunk_count
                
                if progress_callback:
                    progress_callback('complete', result)
                
                return result
            else:
                if progress_callback:
                    progress_callback('error', result_data)
                return {'error': result_data}
                
        except Exception as e:
            error = {'error': str(e)}
            if progress_callback:
                progress_callback('error', error)
            return error

    async def stream_from_chunks(
        self,
        chunk_generator,
        progress_callback: Optional[Callable] = None
    ) -> Dict:
        """
        Stream audio chunks from a generator (for real-time audio sources).
        
        This is the recommended method for TRUE real-time streaming where chunks
        arrive over time (e.g., from a microphone or live audio stream).
        
        Args:
            chunk_generator: An async generator that yields (audio_bytes, format) tuples
                            Example: async for chunk_data, fmt in microphone_stream()
            progress_callback: Optional callback for progress updates
            
        Returns:
            Dict containing sentiment analysis results
            
        Example:
            async def microphone_chunks():
                # Your microphone/audio source code here
                while recording:
                    chunk = await get_audio_from_mic()  # Your audio capture logic
                    yield chunk, "mp3"
            
            result = await client.stream_from_chunks(microphone_chunks())
        """
        # Start session
        if not await self.start_streaming_session(progress_callback):
            return {'error': 'Failed to start streaming session'}
        
        # Stream chunks as they arrive
        try:
            async for audio_data, audio_format in chunk_generator:
                success = await self.send_audio_chunk(
                    audio_data,
                    audio_format,
                    progress_callback
                )
                if not success:
                    return {'error': 'Failed to send audio chunk'}
        except Exception as e:
            if progress_callback:
                progress_callback('error', {'message': f'Chunk generation error: {str(e)}'})
            return {'error': str(e)}
        
        # End session and get results
        return await self.end_streaming_session(progress_callback)

    def _process_sentiment_analysis(
        self,
        audio_path: str,
        enable_translation: bool = False,
        progress_callback: Optional[Callable] = None
    ) -> Dict:
        """
        Internal method to process sentiment analysis.
        
        Args:
            audio_path: Path to audio file
            enable_translation: Enable bilingual response (EN + AR)
            progress_callback: Optional callback for progress updates
            
        Returns:
            Dict containing sentiment analysis results
        """
        import requests

        audio_file_path = Path(audio_path)

        if not audio_file_path.exists():
            error = {'error': f'File not found: {audio_path}'}
            if progress_callback:
                progress_callback('error', error)
            return error

        file_size_mb = audio_file_path.stat().st_size / (1024 * 1024)

        if progress_callback:
            progress_callback('loading', {
                'audio_path': audio_path,
                'file_size_mb': file_size_mb,
                'mode': 'batch_sentiment'
            })

        try:
            endpoint_url = f"{self.http_url}/analyze-sentiment"
            upload_start = time.time()

            with open(audio_file_path, 'rb') as audio_file:
                files = {'audio_file': audio_file}
                data = {'enable_translation': str(enable_translation).lower()}

                if progress_callback:
                    progress_callback('uploading', {
                        'url': endpoint_url,
                        'file_size_mb': file_size_mb,
                        'endpoint': 'sentiment'
                    })

                response = requests.post(
                    endpoint_url,
                    files=files,
                    data=data,
                    timeout=3600
                )

            upload_time = time.time() - upload_start

            if response.status_code == 200:
                result = response.json()
                result['upload_time_seconds'] = upload_time
                result['file_size_mb'] = file_size_mb

                if progress_callback:
                    progress_callback('complete', result)

                return result
            else:
                error = {
                    'error': f'HTTP {response.status_code}',
                    'message': response.text
                }
                if progress_callback:
                    progress_callback('error', error)
                return error

        except Exception as e:
            error = {'error': str(e)}
            if progress_callback:
                progress_callback('error', error)
            return error

    def _process_conversation_analysis(
        self,
        audio: Union[str, bytes],
        file_name: str = "audio.wav",
        enable_translation: bool = False,
        progress_callback: Optional[Callable] = None,
        max_retries: int = 3,
        retry_delay: int = 2
    ) -> Dict:
        """
        Internal method to process conversation analysis with retry logic.
        
        Args:
            audio: Path to the audio file or audio bytes.
            file_name: The name of the file (used when audio is bytes).
            enable_translation: Enable bilingual response (EN + AR). Default: False for faster processing.
            progress_callback: Optional callback for progress updates.
            max_retries: Maximum number of retry attempts for failed requests. Default: 3.
            retry_delay: Delay in seconds between retry attempts. Default: 2.
        Returns:
            A dictionary containing the conversation analysis results.
        """
        import requests
        from requests.adapters import HTTPAdapter
        from urllib3.util.retry import Retry

        file_size_mb = 0
        mode = 'batch_conversation'

        if isinstance(audio, str):
            audio_file_path = Path(audio)
            if not audio_file_path.exists():
                error = {'error': f'File not found: {audio}'}
                if progress_callback:
                    progress_callback('error', error)
                return error
            file_size_mb = audio_file_path.stat().st_size / (1024 * 1024)
            if progress_callback:
                progress_callback('loading', {'audio_path': audio, 'file_size_mb': file_size_mb, 'mode': mode})
        elif isinstance(audio, bytes):
            file_size_mb = len(audio) / (1024 * 1024)
            if progress_callback:
                progress_callback('loading', {'audio_path': file_name, 'file_size_mb': file_size_mb, 'mode': mode})

        # Configure retry strategy
        retry_strategy = Retry(
            total=max_retries,
            backoff_factor=retry_delay,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["POST"]
        )
        
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session = requests.Session()
        session.mount("https://", adapter)
        session.mount("http://", adapter)

        for attempt in range(max_retries + 1):
            try:
                upload_start = time.time()
                endpoint = f"{self.http_url}/analyze-conversation"

                if progress_callback:
                    retry_msg = f" (attempt {attempt + 1}/{max_retries + 1})" if attempt > 0 else ""
                    progress_callback('uploading', {
                        'url': endpoint, 
                        'file_size_mb': file_size_mb,
                        'attempt': attempt + 1,
                        'message': f'Uploading{retry_msg}...'
                    })

                if isinstance(audio, str):
                    with open(audio, 'rb') as f:
                        files = {'audio_file': f}
                        data = {'enable_translation': str(enable_translation).lower()}
                        response = session.post(
                            endpoint,
                            files=files,
                            data=data,
                            timeout=3600
                        )
                else:
                    files = {'audio_file': (file_name, audio, 'audio/wav')}
                    data = {'enable_translation': str(enable_translation).lower()}
                    response = session.post(
                        endpoint,
                        files=files,
                        data=data,
                        timeout=3600
                    )

                upload_time = time.time() - upload_start

                if response.status_code == 200:
                    result = response.json()
                    result['upload_time_seconds'] = upload_time
                    result['file_size_mb'] = file_size_mb

                    if progress_callback:
                        progress_callback('complete', result)

                    return result
                else:
                    error = {'error': f'HTTP {response.status_code}', 'message': response.text}
                    if progress_callback:
                        progress_callback('error', error)
                    return error

            except requests.exceptions.SSLError as e:
                error_msg = str(e)
                if attempt < max_retries:
                    if progress_callback:
                        progress_callback('retrying', {
                            'attempt': attempt + 1,
                            'max_retries': max_retries,
                            'error': 'SSL connection error',
                            'retry_delay': retry_delay
                        })
                    time.sleep(retry_delay)
                    continue
                else:
                    error = {
                        'error': 'SSL Connection Error',
                        'message': 'Failed to establish secure connection with the server. Please check your internet connection and try again.',
                        'technical_details': error_msg
                    }
                    if progress_callback:
                        progress_callback('error', error)
                    return error
                    
            except requests.exceptions.ConnectionError as e:
                error_msg = str(e)
                if attempt < max_retries:
                    if progress_callback:
                        progress_callback('retrying', {
                            'attempt': attempt + 1,
                            'max_retries': max_retries,
                            'error': 'Connection error',
                            'retry_delay': retry_delay
                        })
                    time.sleep(retry_delay)
                    continue
                else:
                    error = {
                        'error': 'Connection Error',
                        'message': 'Unable to connect to the server. Please check your internet connection and try again.',
                        'technical_details': error_msg
                    }
                    if progress_callback:
                        progress_callback('error', error)
                    return error
                    
            except requests.exceptions.Timeout as e:
                error = {
                    'error': 'Request Timeout',
                    'message': 'The request took too long to complete. The audio file may be too large.',
                    'technical_details': str(e)
                }
                if progress_callback:
                    progress_callback('error', error)
                return error
                
            except Exception as e:
                error_msg = str(e)
                if attempt < max_retries and ('SSL' in error_msg or 'Connection' in error_msg):
                    if progress_callback:
                        progress_callback('retrying', {
                            'attempt': attempt + 1,
                            'max_retries': max_retries,
                            'error': error_msg,
                            'retry_delay': retry_delay
                        })
                    time.sleep(retry_delay)
                    continue
                else:
                    error = {
                        'error': 'Unexpected Error',
                        'message': error_msg
                    }
                    if progress_callback:
                        progress_callback('error', error)
                    return error
        
        # This should never be reached, but just in case
        return {
            'error': 'Max retries exceeded',
            'message': 'Failed to complete request after multiple attempts'
        }

    def batch_process(
        self,
        audio_path: str,
        endpoint: str = "sentiment",
        enable_translation: bool = False,
        progress_callback: Optional[Callable] = None
    ) -> Dict:
        """
        Gateway method for batch processing audio files.
        Routes to the appropriate analysis endpoint based on the endpoint parameter.

        Args:
            audio_path: Path to audio file
            endpoint: Analysis endpoint - 'sentiment' or 'conversation'. Default: 'sentiment'
            enable_translation: Enable bilingual response (EN + AR). Default: False for faster processing.
            progress_callback: Optional callback function for progress updates

        Returns:
            Dict containing analysis results (sentiment or conversation based on endpoint)
        """
        # Validate endpoint parameter
        valid_endpoints = ['sentiment', 'conversation']
        if endpoint not in valid_endpoints:
            error = {'error': f'Invalid endpoint: {endpoint}. Must be one of {valid_endpoints}'}
            if progress_callback:
                progress_callback('error', error)
            return error

        # Route to appropriate processing method
        if endpoint == 'sentiment':
            return self._process_sentiment_analysis(
                audio_path=audio_path,
                enable_translation=enable_translation,
                progress_callback=progress_callback
            )
        else:  # conversation
            return self._process_conversation_analysis(
                audio=audio_path,
                enable_translation=enable_translation,
                progress_callback=progress_callback
            )

    async def stream_realtime_async(
        self,
        audio_path: str,
        chunk_duration_seconds: int = 30,
        progress_callback: Optional[Callable] = None
    ) -> Dict:
        """Async wrapper for stream_realtime (for consistency)."""
        return await self.stream_realtime(audio_path, chunk_duration_seconds, progress_callback)

    def get_session_info(self) -> Dict:
        """Get current session information."""
        return {
            'session_id': self.session_id,
            'chunk_count': self.chunk_count,
            'is_streaming': self.is_streaming
        }

    @staticmethod
    def extract_language_version(result: Dict, language: str = 'EN') -> Dict:
        """
        Extract a specific language version from a bilingual response.
        
        Args:
            result: Analysis result (may be bilingual with EN/AR keys)
            language: Language to extract ('EN' or 'AR')
            
        Returns:
            Dictionary with the specified language version plus any metadata
        """
        if language.upper() not in result:
            # Not a bilingual response, return as-is
            return result
        
        # Extract the language version
        lang_data = result[language.upper()].copy()
        
        # Add back any metadata or technical fields
        for key, value in result.items():
            if key not in ['EN', 'AR']:
                lang_data[key] = value
        
        return lang_data

    @staticmethod
    def is_bilingual_response(result: Dict) -> bool:
        """
        Check if a response is in bilingual format.
        
        Args:
            result: Analysis result
            
        Returns:
            True if result has EN and AR keys
        """
        return 'EN' in result and 'AR' in result
