"""
NOMAZRA WebSocket Consumer — Real-time AI Voice Call.
Architecture:
  1. Client connects via WebSocket with JWT token
  2. Client sends audio chunks (base64 PCM 16-bit 16kHz mono)
  3. On silence/stop → Whisper transcription → GPT response → TTS audio
  4. Server streams back transcription + text + audio
  5. On disconnect → call summary generated and saved
"""
import os
import io
import json
import time
import uuid
import wave
import base64
import logging
import tempfile
import asyncio
from datetime import timedelta

from channels.generic.websocket import AsyncJsonWebsocketConsumer
from channels.db import database_sync_to_async
from django.utils import timezone

logger = logging.getLogger(__name__)

# Audio config: PCM 16-bit, 16kHz, mono
SAMPLE_RATE = 16000
SAMPLE_WIDTH = 2  # 16-bit
CHANNELS = 1
MAX_CALL_DURATION = 1800  # 30 minutes max
MAX_AUDIO_BUFFER = 60 * SAMPLE_RATE * SAMPLE_WIDTH  # 60s of audio max per turn


class VoiceCallConsumer(AsyncJsonWebsocketConsumer):
    """
    WebSocket consumer for real-time voice calls.
    Protocol: JSON messages with 'type' field.
    """

    async def connect(self):
        self.conversation_id = self.scope['url_route']['kwargs']['conversation_id']
        self.user = None
        self.session_id = None
        self.audio_buffer = bytearray()
        self.transcript_history = []
        self.call_start_time = None
        self.voice = 'nova'
        self.language = 'km'
        self.is_call_active = False
        self._response_task = None

        # Authenticate via query param
        query_string = self.scope.get('query_string', b'').decode()
        token = self._parse_token(query_string)
        if not token:
            await self.close(code=4001)
            return

        user = await self._authenticate_token(token)
        if not user:
            await self.close(code=4001)
            return

        self.user = user
        # Verify conversation belongs to user
        conv = await self._get_conversation()
        if not conv:
            await self.close(code=4004)
            return

        await self.accept()

        # Create voice session
        self.session_id = await self._create_voice_session()
        self.call_start_time = time.time()
        self.is_call_active = True

        await self.send_json({
            'type': 'call.connected',
            'data': {
                'session_id': str(self.session_id),
                'message': 'Connected. You can start speaking.',
            }
        })

    async def disconnect(self, close_code):
        if self.is_call_active and self.session_id:
            await self._end_call('disconnected')
        if self._response_task and not self._response_task.done():
            self._response_task.cancel()

    async def receive_json(self, content, **kwargs):
        msg_type = content.get('type', '')
        data = content.get('data', {})

        try:
            if msg_type == 'call.start':
                self.language = data.get('language', 'km')
                self.voice = data.get('voice', 'nova')
                context = data.get('context', '')
                if context:
                    self.transcript_history.append({
                        'role': 'user',
                        'text': f'[Context: {context[:500]}]',
                        'timestamp': self._elapsed(),
                    })

            elif msg_type == 'audio.chunk':
                await self._handle_audio_chunk(data)

            elif msg_type == 'audio.stop':
                await self._handle_audio_stop()

            elif msg_type == 'call.interrupt':
                if self._response_task and not self._response_task.done():
                    self._response_task.cancel()

            elif msg_type == 'call.end':
                reason = data.get('reason', 'user_ended')
                await self._end_call(reason)

            else:
                await self.send_json({
                    'type': 'error',
                    'data': {'code': 'unknown_type', 'message': f'Unknown message type: {msg_type}'}
                })

        except Exception as e:
            logger.error("VoiceCallConsumer error: %s", e)
            await self.send_json({
                'type': 'error',
                'data': {'code': 'server_error', 'message': 'Internal error processing message'}
            })

    # ── Audio handling ──

    async def _handle_audio_chunk(self, data):
        """Buffer incoming audio chunk."""
        if not self.is_call_active:
            return

        # Check call duration
        if time.time() - self.call_start_time > MAX_CALL_DURATION:
            await self._end_call('max_duration')
            return

        audio_b64 = data.get('audio', '')
        if not audio_b64:
            return

        try:
            audio_bytes = base64.b64decode(audio_b64)
        except Exception:
            return

        # Prevent buffer overflow
        if len(self.audio_buffer) + len(audio_bytes) > MAX_AUDIO_BUFFER:
            # Auto-process if buffer is getting too large
            await self._handle_audio_stop()
            return

        self.audio_buffer.extend(audio_bytes)

    async def _handle_audio_stop(self):
        """User stopped speaking — process buffered audio."""
        if not self.audio_buffer:
            return

        # Grab buffer and reset
        audio_data = bytes(self.audio_buffer)
        self.audio_buffer = bytearray()

        # Process in background task so we don't block
        self._response_task = asyncio.create_task(
            self._process_turn(audio_data)
        )

    async def _process_turn(self, audio_data):
        """Full turn: transcribe → GPT → TTS → stream back."""
        try:
            # 1. Transcribe with Whisper
            transcription = await self._transcribe_audio(audio_data)
            if not transcription or not transcription.strip():
                return

            # Send final transcription to client
            await self.send_json({
                'type': 'transcription.final',
                'data': {
                    'text': transcription,
                    'language': self.language,
                }
            })

            # Record in history
            self.transcript_history.append({
                'role': 'user',
                'text': transcription,
                'timestamp': self._elapsed(),
            })

            # 2. Get AI response
            ai_response, tokens_used = await self._get_ai_response(transcription)
            if not ai_response:
                return

            # Send complete text response
            await self.send_json({
                'type': 'response.text.done',
                'data': {
                    'text': ai_response,
                    'tokens_used': tokens_used,
                }
            })

            # Record in history
            self.transcript_history.append({
                'role': 'assistant',
                'text': ai_response,
                'timestamp': self._elapsed(),
            })

            # 3. Generate TTS audio and stream back
            await self._generate_and_stream_tts(ai_response)

        except asyncio.CancelledError:
            # Call was interrupted
            logger.info("Voice call turn cancelled (interrupted)")
        except Exception as e:
            logger.error("Voice call turn error: %s", e)
            await self.send_json({
                'type': 'error',
                'data': {'code': 'ai_error', 'message': 'Error processing your voice input'}
            })

    # ── Whisper transcription ──

    async def _transcribe_audio(self, pcm_data):
        """Convert PCM audio to WAV and transcribe with Whisper."""
        try:
            return await asyncio.to_thread(self._sync_transcribe, pcm_data)
        except Exception as e:
            logger.error("Transcription error: %s", e)
            await self.send_json({
                'type': 'error',
                'data': {'code': 'transcription_error', 'message': 'Could not transcribe audio'}
            })
            return None

    def _sync_transcribe(self, pcm_data):
        """Synchronous Whisper transcription."""
        import openai

        # Convert raw PCM to WAV
        wav_buffer = io.BytesIO()
        with wave.open(wav_buffer, 'wb') as wav:
            wav.setnchannels(CHANNELS)
            wav.setsampwidth(SAMPLE_WIDTH)
            wav.setframerate(SAMPLE_RATE)
            wav.writeframes(pcm_data)
        wav_buffer.seek(0)
        wav_buffer.name = 'audio.wav'

        client = openai.OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
        transcript = client.audio.transcriptions.create(
            model='whisper-1',
            file=wav_buffer,
            response_format='text',
            language=self.language if self.language != 'km' else None,
        )
        return transcript.strip()

    # ── AI response ──

    async def _get_ai_response(self, user_text):
        """Get GPT response with conversation context."""
        try:
            return await asyncio.to_thread(self._sync_ai_response, user_text)
        except Exception as e:
            logger.error("AI response error: %s", e)
            await self.send_json({
                'type': 'error',
                'data': {'code': 'ai_error', 'message': 'Error generating response'}
            })
            return None, 0

    def _sync_ai_response(self, user_text):
        """Synchronous GPT call."""
        from . import ai_service

        # Build history from transcript
        history = []
        for entry in self.transcript_history[-10:]:  # Last 10 exchanges
            if entry['role'] in ('user', 'assistant'):
                history.append({
                    'role': entry['role'],
                    'content': entry['text'],
                })

        result = ai_service.chat(
            user_message=user_text,
            conversation_history=history,
        )
        tokens = result.get('input_tokens', 0) + result.get('output_tokens', 0)
        return result['response'], tokens

    # ── TTS ──

    async def _generate_and_stream_tts(self, text):
        """Generate TTS audio and stream chunks to client."""
        try:
            audio_data = await asyncio.to_thread(self._sync_tts, text)
            if not audio_data:
                return

            # Stream in ~32KB chunks
            chunk_size = 32768
            sequence = 0
            for i in range(0, len(audio_data), chunk_size):
                chunk = audio_data[i:i + chunk_size]
                await self.send_json({
                    'type': 'response.audio.delta',
                    'data': {
                        'audio': base64.b64encode(chunk).decode('ascii'),
                        'sequence': sequence,
                    }
                })
                sequence += 1

            # Calculate approximate duration (MP3 ~16kbps for speech)
            duration_estimate = len(audio_data) / 16000 if audio_data else 0

            await self.send_json({
                'type': 'response.audio.done',
                'data': {
                    'duration_seconds': round(duration_estimate, 1),
                    'sequence_count': sequence,
                }
            })

        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.error("TTS streaming error: %s", e)
            await self.send_json({
                'type': 'error',
                'data': {'code': 'tts_error', 'message': 'Error generating voice response'}
            })

    def _sync_tts(self, text):
        """Synchronous OpenAI TTS call."""
        import openai

        # TTS has a 4096 char limit
        text = text[:4096]

        client = openai.OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
        response = client.audio.speech.create(
            model='tts-1',
            voice=self.voice,
            input=text,
            response_format='mp3',
        )
        return response.content

    # ── Call lifecycle ──

    async def _end_call(self, reason):
        """End the call, generate summary, save session."""
        if not self.is_call_active:
            return
        self.is_call_active = False

        duration = int(time.time() - self.call_start_time) if self.call_start_time else 0

        # Generate summary if there's a transcript
        summary = ''
        advice = []
        total_tokens = 0
        if self.transcript_history:
            summary, advice, total_tokens = await asyncio.to_thread(
                self._sync_generate_summary
            )

        # Save session to database
        await self._save_voice_session(duration, summary)

        # Save messages to conversation
        await self._save_messages_to_conversation()

        # Send summary to client
        await self.send_json({
            'type': 'call.summary',
            'data': {
                'session_id': str(self.session_id),
                'duration_seconds': duration,
                'transcript': self.transcript_history,
                'summary': summary,
                'tokens_used': total_tokens,
                'advice': advice,
            }
        })

        await self.send_json({
            'type': 'call.ended',
            'data': {
                'reason': reason,
                'session_id': str(self.session_id),
            }
        })

        await self.close()

    def _sync_generate_summary(self):
        """Generate call summary and action items."""
        from . import ai_service

        # Build transcript text
        transcript_text = '\n'.join(
            f"{'User' if e['role'] == 'user' else 'AI'}: {e['text']}"
            for e in self.transcript_history
            if e['role'] in ('user', 'assistant')
        )

        prompt = (
            "Summarize this legal voice consultation. Provide:\n"
            "1. A brief summary (2-3 sentences)\n"
            "2. Key legal advice points (as a JSON array of strings)\n\n"
            "Format your response as:\n"
            "SUMMARY: <your summary>\n"
            "ADVICE: [\"advice 1\", \"advice 2\", ...]\n\n"
            f"Transcript:\n{transcript_text[:8000]}"
        )

        result = ai_service.chat(prompt)
        response = result['response']
        tokens = result.get('input_tokens', 0) + result.get('output_tokens', 0)

        # Parse summary and advice
        summary = response
        advice = []
        if 'SUMMARY:' in response:
            parts = response.split('ADVICE:')
            summary = parts[0].replace('SUMMARY:', '').strip()
            if len(parts) > 1:
                try:
                    advice = json.loads(parts[1].strip())
                except (json.JSONDecodeError, ValueError):
                    advice = [parts[1].strip()]

        return summary, advice, tokens

    # ── Auth helpers ──

    def _parse_token(self, query_string):
        """Extract token from query string."""
        params = dict(p.split('=', 1) for p in query_string.split('&') if '=' in p)
        return params.get('token', '')

    @database_sync_to_async
    def _authenticate_token(self, token):
        """Validate JWT token and return user."""
        try:
            from rest_framework_simplejwt.tokens import AccessToken
            from django.contrib.auth import get_user_model
            User = get_user_model()

            access = AccessToken(token)
            user_id = access['user_id']
            return User.objects.get(id=user_id, is_active=True, is_deleted=False)
        except Exception as e:
            logger.warning("WebSocket auth failed: %s", e)
            return None

    @database_sync_to_async
    def _get_conversation(self):
        """Verify conversation exists and belongs to user."""
        try:
            from .models import Conversation
            return Conversation.objects.get(
                pk=self.conversation_id,
                user=self.user,
            )
        except Exception:
            return None

    @database_sync_to_async
    def _create_voice_session(self):
        """Create a VoiceSession record."""
        from .models import VoiceSession, Conversation, UserStats
        try:
            conv = Conversation.objects.get(pk=self.conversation_id)
            session = VoiceSession.objects.create(
                user=self.user,
                conversation=conv,
            )
            stats, _ = UserStats.objects.get_or_create(user=self.user)
            stats.total_voice_sessions += 1
            stats.save()
            return session.id
        except Exception as e:
            logger.error("Create voice session failed: %s", e)
            return uuid.uuid4()

    @database_sync_to_async
    def _save_voice_session(self, duration, summary):
        """Update voice session with duration, transcript, summary."""
        from .models import VoiceSession
        try:
            session = VoiceSession.objects.get(pk=self.session_id)
            session.duration_seconds = duration
            session.transcript = json.dumps(self.transcript_history)
            session.summary = summary
            session.save()
        except Exception as e:
            logger.error("Save voice session failed: %s", e)

    @database_sync_to_async
    def _save_messages_to_conversation(self):
        """Save transcript entries as conversation messages."""
        from .models import Message, Conversation
        try:
            conv = Conversation.objects.get(pk=self.conversation_id)
            for entry in self.transcript_history:
                if entry['role'] in ('user', 'assistant'):
                    Message.objects.create(
                        conversation=conv,
                        role=entry['role'],
                        content=entry['text'],
                        input_type='call',
                    )
        except Exception as e:
            logger.error("Save call messages failed: %s", e)

    # ── Utilities ──

    def _elapsed(self):
        """Get elapsed time string since call start."""
        if not self.call_start_time:
            return '00:00'
        elapsed = int(time.time() - self.call_start_time)
        mins = elapsed // 60
        secs = elapsed % 60
        return f'{mins:02d}:{secs:02d}'
