"""
End-to-End Tests for Streaming API

These tests validate the real-time streaming API workflow:
- WebSocket connection establishment
- Audio streaming in chunks
- Real-time emotion predictions
- Connection lifecycle management

Note: These tests make REAL WebSocket connections to the demo server.
Set VALENCE_E2E_TEST=true and VALENCE_API_KEY to run these tests.
"""

import os
import pytest
import time
import wave
import struct
from valenceai import ValenceClient


# Skip E2E tests unless explicitly enabled
pytestmark = pytest.mark.skipif(
    os.getenv('VALENCE_E2E_TEST') != 'true',
    reason="E2E tests only run when VALENCE_E2E_TEST=true"
)


def generate_audio_chunk(duration_ms=100, sample_rate=44100):
    """
    Generate a small audio chunk for streaming.

    Args:
        duration_ms: Duration in milliseconds
        sample_rate: Sample rate in Hz

    Returns:
        bytes: PCM audio data
    """
    num_samples = int((duration_ms / 1000) * sample_rate)
    frequency = 440.0  # A4 note

    audio_data = bytearray()
    for i in range(num_samples):
        value = int(32767.0 * 0.3 * (i / sample_rate) * frequency % 1.0)
        audio_data.extend(struct.pack('h', value))

    return bytes(audio_data)


class TestStreamingE2E:
    """
    End-to-End tests for the Streaming API.
    These tests make REAL WebSocket connections.
    """

    @pytest.fixture(autouse=True)
    def setup(self):
        """Setup test fixtures"""
        # Check for API key
        self.api_key = os.getenv('VALENCE_API_KEY')
        if not self.api_key:
            pytest.skip("VALENCE_API_KEY not set")

        # Use demo WebSocket URL by default
        self.websocket_url = os.getenv('VALENCE_WEBSOCKET_URL', 'https://demo.getvalenceai.com')

        self.client = ValenceClient(
            api_key=self.api_key,
            websocket_url=self.websocket_url
        )

        self.predictions_received = []
        self.errors_received = []
        self.connected = False

    def test_streaming_connection_lifecycle(self):
        """
        Test complete streaming connection lifecycle:
        - Connect to WebSocket
        - Verify connection
        - Disconnect
        """
        print("\n" + "="*70)
        print("Testing Streaming Connection Lifecycle")
        print("="*70)

        # Create streaming connection
        print("\nCreating streaming connection...")
        stream = self.client.streaming.connect(model="4emotions")

        # Define callbacks
        def on_connected(info):
            self.connected = True
            print(f"✓ Connected! Session ID: {info.get('session_id', 'N/A')}")

        def on_error(error):
            self.errors_received.append(error)
            print(f"❌ Error: {error.get('error', str(error))}")

        # Register callbacks
        stream.on_connected(on_connected)
        stream.on_error(on_error)

        # Connect
        print("Connecting to WebSocket...")
        stream.connect()

        # Wait for connection
        time.sleep(2)

        assert self.connected, "Should be connected to WebSocket"
        assert stream.connected, "Stream should report connected state"
        print("✓ Connection established")

        # Disconnect
        print("\nDisconnecting...")
        stream.disconnect()
        time.sleep(1)

        assert not stream.connected, "Stream should report disconnected state"
        print("✓ Disconnected successfully")

        print("\n" + "="*70)
        print("✅ CONNECTION LIFECYCLE TEST PASSED")
        print("="*70)

    def test_streaming_send_audio_and_receive_predictions(self):
        """
        Test sending audio chunks and receiving predictions.
        """
        print("\n" + "="*70)
        print("Testing Streaming Audio and Predictions")
        print("="*70)

        # Create streaming connection
        stream = self.client.streaming.connect(model="4emotions")

        # Define callbacks
        def on_connected(info):
            self.connected = True
            print(f"✓ Connected! Session ID: {info.get('session_id', 'N/A')}")

        def on_prediction(data):
            self.predictions_received.append(data)
            emotion = data.get('main_emotion', 'unknown')
            confidence = data.get('confidence', 0)
            print(f"  Prediction received: {emotion} ({confidence:.2%})")

        def on_error(error):
            self.errors_received.append(error)
            print(f"❌ Error: {error.get('error', str(error))}")

        # Register callbacks
        stream.on_connected(on_connected)
        stream.on_prediction(on_prediction)
        stream.on_error(on_error)

        # Connect
        print("\nConnecting to WebSocket...")
        stream.connect()
        time.sleep(2)

        assert self.connected, "Should be connected"
        print("✓ Connection established")

        # Send audio chunks
        print("\nSending audio chunks...")
        num_chunks = 10

        for i in range(num_chunks):
            audio_chunk = generate_audio_chunk(duration_ms=500)  # 500ms chunks
            stream.send_audio(audio_chunk)
            print(f"  Sent chunk {i+1}/{num_chunks} ({len(audio_chunk)} bytes)")
            time.sleep(0.5)  # Small delay between chunks

        print(f"\n✓ Sent {num_chunks} audio chunks")

        # Wait for predictions
        print("\nWaiting for predictions...")
        time.sleep(3)

        print(f"\n✓ Received {len(self.predictions_received)} predictions")

        # Verify predictions
        assert len(self.predictions_received) > 0, "Should receive at least one prediction"

        # Verify prediction structure
        first_prediction = self.predictions_received[0]
        assert 'main_emotion' in first_prediction, "Prediction should have 'main_emotion'"
        assert 'confidence' in first_prediction, "Prediction should have 'confidence'"
        assert 'all_predictions' in first_prediction, "Prediction should have 'all_predictions'"

        print("\nPrediction structure:")
        print(f"  main_emotion: {first_prediction['main_emotion']}")
        print(f"  confidence: {first_prediction['confidence']:.2%}")
        print(f"  all_predictions: {list(first_prediction['all_predictions'].keys())}")

        # Disconnect
        stream.disconnect()
        time.sleep(1)

        print("\n" + "="*70)
        print("✅ STREAMING PREDICTIONS TEST PASSED")
        print("="*70)

    def test_streaming_multiple_models(self):
        """
        Test streaming with different emotion models.
        """
        print("\n" + "="*70)
        print("Testing Streaming with Multiple Models")
        print("="*70)

        models = ["4emotions", "7emotions"]

        for model in models:
            print(f"\n--- Testing with model: {model} ---")

            self.predictions_received = []
            self.connected = False

            # Create streaming connection
            stream = self.client.streaming.connect(model=model)

            # Define callbacks
            def on_connected(info):
                self.connected = True
                print(f"✓ Connected with {model}")

            def on_prediction(data):
                self.predictions_received.append(data)
                print(f"  Prediction: {data.get('main_emotion')} ({data.get('confidence', 0):.2%})")

            # Register callbacks
            stream.on_connected(on_connected)
            stream.on_prediction(on_prediction)

            # Connect and send audio
            stream.connect()
            time.sleep(2)

            assert self.connected, f"Should connect with {model}"

            # Send a few chunks
            for i in range(3):
                audio_chunk = generate_audio_chunk(duration_ms=500)
                stream.send_audio(audio_chunk)
                time.sleep(0.5)

            time.sleep(2)

            print(f"  Received {len(self.predictions_received)} predictions with {model}")

            # Verify model-specific emotions
            if len(self.predictions_received) > 0:
                emotions = set()
                for pred in self.predictions_received:
                    emotions.add(pred['main_emotion'])
                    emotions.update(pred['all_predictions'].keys())

                print(f"  Emotions detected: {emotions}")

                if model == "4emotions":
                    expected = {"happy", "sad", "angry", "neutral"}
                    assert emotions.issubset(expected) or len(emotions.intersection(expected)) > 0, \
                        f"Emotions should be from 4emotions set"

                elif model == "7emotions":
                    expected = {"happy", "sad", "angry", "neutral", "surprised", "disgusted", "calm"}
                    assert emotions.issubset(expected) or len(emotions.intersection(expected)) > 0, \
                        f"Emotions should be from 7emotions set"

            # Disconnect
            stream.disconnect()
            time.sleep(1)

        print("\n" + "="*70)
        print("✅ MULTIPLE MODELS TEST PASSED")
        print("="*70)

    def test_streaming_continuous_operation(self):
        """
        Test continuous streaming over a longer period.
        """
        print("\n" + "="*70)
        print("Testing Continuous Streaming Operation")
        print("="*70)

        # Create streaming connection
        stream = self.client.streaming.connect(model="4emotions")

        # Define callbacks
        def on_connected(info):
            self.connected = True
            print(f"✓ Connected! Session ID: {info.get('session_id', 'N/A')}")

        def on_prediction(data):
            self.predictions_received.append(data)
            if len(self.predictions_received) % 5 == 0:
                print(f"  Predictions received: {len(self.predictions_received)}")

        def on_error(error):
            self.errors_received.append(error)
            print(f"❌ Error: {error.get('error', str(error))}")

        # Register callbacks
        stream.on_connected(on_connected)
        stream.on_prediction(on_prediction)
        stream.on_error(on_error)

        # Connect
        print("\nConnecting to WebSocket...")
        stream.connect()
        time.sleep(2)

        assert self.connected, "Should be connected"

        # Stream continuously for 30 seconds
        print("\nStreaming audio for 30 seconds...")
        duration = 30
        start_time = time.time()
        chunk_count = 0

        while time.time() - start_time < duration:
            audio_chunk = generate_audio_chunk(duration_ms=500)
            stream.send_audio(audio_chunk)
            chunk_count += 1
            time.sleep(0.5)

        print(f"\n✓ Sent {chunk_count} chunks over {duration} seconds")

        # Wait for final predictions
        time.sleep(3)

        print(f"✓ Received {len(self.predictions_received)} total predictions")
        print(f"  Average: {len(self.predictions_received) / duration:.1f} predictions/second")

        # Verify continuous operation
        assert len(self.predictions_received) > 0, "Should receive predictions"
        assert len(self.errors_received) == 0, "Should not have errors during streaming"
        assert stream.connected, "Should remain connected throughout"

        # Analyze emotion distribution
        emotion_counts = {}
        for pred in self.predictions_received:
            emotion = pred['main_emotion']
            emotion_counts[emotion] = emotion_counts.get(emotion, 0) + 1

        print("\nEmotion distribution:")
        for emotion, count in sorted(emotion_counts.items(), key=lambda x: x[1], reverse=True):
            percentage = (count / len(self.predictions_received)) * 100
            print(f"  {emotion}: {count} ({percentage:.1f}%)")

        # Disconnect
        stream.disconnect()
        time.sleep(1)

        print("\n" + "="*70)
        print("✅ CONTINUOUS OPERATION TEST PASSED")
        print("="*70)

    def test_streaming_error_handling(self):
        """
        Test error handling in streaming API.
        """
        print("\n" + "="*70)
        print("Testing Streaming Error Handling")
        print("="*70)

        # Test 1: Invalid model
        print("\n1. Testing with invalid model...")
        try:
            stream = self.client.streaming.connect(model="invalid_model")

            def on_error(error):
                self.errors_received.append(error)
                print(f"  ✓ Error received as expected: {error.get('error', str(error))}")

            stream.on_error(on_error)
            stream.connect()
            time.sleep(2)

            # Might receive error or connection might fail
            # Either case is acceptable for invalid model
            print("  ✓ Invalid model handled")

            stream.disconnect()
        except Exception as e:
            print(f"  ✓ Exception raised as expected: {e}")

        # Test 2: Send audio before connecting
        print("\n2. Testing send audio before connection...")
        stream = self.client.streaming.connect(model="4emotions")

        # Try to send without connecting
        try:
            audio_chunk = generate_audio_chunk(duration_ms=500)
            stream.send_audio(audio_chunk)
            # May or may not raise exception, depends on implementation
            print("  ✓ Send before connect handled")
        except Exception as e:
            print(f"  ✓ Exception raised as expected: {e}")

        # Test 3: Reconnection after disconnect
        print("\n3. Testing reconnection...")

        self.connected = False

        def on_connected(info):
            self.connected = True
            print(f"  ✓ Reconnected! Session ID: {info.get('session_id', 'N/A')}")

        stream.on_connected(on_connected)
        stream.connect()
        time.sleep(2)

        assert self.connected, "Should reconnect successfully"

        stream.disconnect()
        time.sleep(1)

        print("\n" + "="*70)
        print("✅ ERROR HANDLING TEST PASSED")
        print("="*70)


if __name__ == "__main__":
    print("\n" + "="*70)
    print("Valence SDK - Streaming API E2E Tests")
    print("="*70)
    print("\nTo run these tests:")
    print("  export VALENCE_E2E_TEST=true")
    print("  export VALENCE_API_KEY=your_api_key")
    print("  pytest tests/test_e2e_streaming.py -v -s")
    print("\n" + "="*70)
