#!/usr/bin/env python3
"""
Test script for waveform transfer functionality.
"""

from tektronix_base import TektronixOscilloscopeBase


def main():
    """Test waveform transfer functions on oscilloscope at 192.168.10.105"""

    scope_ip = "192.168.10.105"

    print(f"Connecting to oscilloscope at {scope_ip}...")

    scope = TektronixOscilloscopeBase(resource_name=scope_ip)

    try:
        # Connect to scope
        scope.connect()
        print(f"Connected to {scope.identify()}\n")

        # Save current settings
        print("Saving current waveform transfer settings...")
        current_data_encoding = scope.get_data_encoding()
        current_data_source = scope.get_data_source()
        current_wfmoutpre_encoding = scope.get_wfmoutpre_encoding()
        current_byte_count = scope.get_wfmoutpre_byte_count()
        current_byte_order = scope.get_wfmoutpre_byte_order()
        print(f"  Current data encoding: {current_data_encoding}")
        print(f"  Current data source: {current_data_source}")
        print(f"  Current WFMOutpre encoding: {current_wfmoutpre_encoding}")
        print(f"  Current byte count: {current_byte_count}")
        print(f"  Current byte order: {current_byte_order}\n")

        # Test setting data encoding
        print("Testing data encoding settings:")
        test_encodings = ['ASCIi', 'RIBinary', 'RPBinary']
        for encoding in test_encodings:
            print(f"  Setting data encoding to {encoding}...")
            scope.set_data_encoding(encoding)
            actual = scope.get_data_encoding()
            print(f"  Actual encoding: {actual}")

        # Test setting data source
        print("\nTesting data source selection:")
        for source in ['CH1', 'CH2', 'CH3', 'CH4']:
            print(f"  Setting data source to {source}...")
            scope.set_data_source(source)
            actual = scope.get_data_source()
            print(f"  Actual source: {actual}")

        # Test with integer source
        print("\nTesting data source with integer (2)...")
        scope.set_data_source(2)
        actual = scope.get_data_source()
        print(f"  Actual source: {actual}")

        # Test waveform preamble encoding
        print("\nTesting waveform preamble encoding:")
        for encoding in ['BINary', 'ASCii']:
            print(f"  Setting WFMOutpre encoding to {encoding}...")
            scope.set_wfmoutpre_encoding(encoding)
            actual = scope.get_wfmoutpre_encoding()
            print(f"  Actual encoding: {actual}")

        # Test byte count
        print("\nTesting byte count settings:")
        for byte_count in [1, 2]:
            print(f"  Setting byte count to {byte_count}...")
            scope.set_wfmoutpre_byte_count(byte_count)
            actual = scope.get_wfmoutpre_byte_count()
            print(f"  Actual byte count: {actual}")

        # Test byte order
        print("\nTesting byte order settings:")
        for byte_order in ['MSB', 'LSB']:
            print(f"  Setting byte order to {byte_order}...")
            scope.set_wfmoutpre_byte_order(byte_order)
            actual = scope.get_wfmoutpre_byte_order()
            print(f"  Actual byte order: {actual}")

        # Query complete waveform preamble
        print("\nQuerying complete waveform preamble...")
        preamble = scope.query_wfmoutpre()
        print(f"  Preamble (first 100 chars): {preamble[:100]}...")

        # Set up for binary waveform acquisition
        print("\nConfiguring for binary waveform transfer:")
        print("  Setting data encoding to RIBinary...")
        scope.set_data_encoding('RIBinary')
        print("  Setting WFMOutpre encoding to BINary...")
        scope.set_wfmoutpre_encoding('BINary')
        print("  Setting byte count to 1...")
        scope.set_wfmoutpre_byte_count(1)
        print("  Setting byte order to MSB...")
        scope.set_wfmoutpre_byte_order('MSB')
        print("  Setting data source to CH1...")
        scope.set_data_source('CH1')
        print("  Configuration complete")

        # Transfer curve data
        print("\nTransferring curve data from CH1...")
        curve_data = scope.transfer_curve()
        print(f"  Received {len(curve_data)} bytes of curve data")
        print(f"  First 10 bytes (raw): {list(curve_data[:10])}")

        # Parse the curve data
        print("\nParsing curve data...")
        values = scope.parse_curve_data(curve_data, byte_count=1, signed=True, byte_order='MSB')
        print(f"  Parsed {len(values)} samples")
        print(f"  First 10 values: {values[:10]}")
        print(f"  Min value: {min(values)}")
        print(f"  Max value: {max(values)}")
        print(f"  Average value: {sum(values) / len(values):.2f}")

        # Transfer complete waveform (preamble + curve)
        print("\nTransferring complete waveform (WAVFrm?)...")
        preamble_str, curve_bytes = scope.transfer_waveform()
        print(f"  Received preamble: {preamble_str[:100]}...")
        print(f"  Received {len(curve_bytes)} bytes of curve data")

        # Parse this curve data too
        waveform_values = scope.parse_curve_data(curve_bytes, byte_count=1, signed=True, byte_order='MSB')
        print(f"  Parsed {len(waveform_values)} samples from complete waveform")

        # Test high-level acquire_waveform method
        print("\nTesting high-level acquire_waveform method:")
        print("  Acquiring waveform from CH1...")
        waveform = scope.acquire_waveform('CH1')
        print(f"  Acquired {len(waveform)} samples")
        print(f"  First 10 values: {waveform[:10]}")
        print(f"  Min: {min(waveform)}, Max: {max(waveform)}, Avg: {sum(waveform)/len(waveform):.2f}")

        print("\n  Acquiring waveform from CH2 (using integer)...")
        try:
            waveform2 = scope.acquire_waveform(2)
            print(f"  Acquired {len(waveform2)} samples from CH2")
            print(f"  First 10 values: {waveform2[:10]}")
        except Exception as e:
            print(f"  Could not acquire from CH2 (channel may not be active): {e}")

        # Test FastFrame support
        print("\nTesting FastFrame frame selection:")

        # Check if FastFrame is currently enabled
        try:
            ff_state = scope.get_fastframe_state()
            ff_count = scope.get_fastframe_count()
            print(f"  Current FastFrame state: {ff_state} ({'active' if ff_state else 'off'})")
            print(f"  Current frame count: {ff_count}")
        except Exception as e:
            print(f"  Could not query FastFrame state (skipping FastFrame tests): {e}")
            ff_state = None

        if ff_state is not None and ff_state:
            # FastFrame is active, test frame selection
            print("  FastFrame is active, testing frame selection...")
            current_frame = scope.get_fastframe_selected()
            print(f"  Current selected frame: {current_frame}")

            # Try selecting different frames
            for frame in [1, min(5, ff_count), ff_count]:
                print(f"  Selecting frame {frame}...")
                scope.set_fastframe_selected(frame)
                actual = scope.get_fastframe_selected()
                print(f"    Selected frame: {actual}")

                # Acquire waveform from this frame
                print(f"    Acquiring waveform from frame {frame}...")
                frame_waveform = scope.acquire_waveform('CH1', frame_number=frame)
                print(f"    Acquired {len(frame_waveform)} samples")
                print(f"    First 5 values: {frame_waveform[:5]}")

            # Restore original frame
            scope.set_fastframe_selected(current_frame)
        elif ff_state is not None:
            print("  FastFrame is not active, enabling it temporarily...")
            scope.set_fastframe_state(True)
            scope.set_fastframe_count(10)
            print("  FastFrame enabled with 10 frames")

            # Test frame selection
            for frame in [1, 5, 10]:
                print(f"  Selecting frame {frame}...")
                scope.set_fastframe_selected(frame)
                actual = scope.get_fastframe_selected()
                print(f"    Selected frame: {actual}")

            # Restore FastFrame state
            scope.set_fastframe_state(False)
            print("  FastFrame disabled (restored)")

        # Restore original settings
        print("\nRestoring original waveform transfer settings...")
        try:
            scope.set_data_encoding(current_data_encoding)
            scope.set_data_source(current_data_source)
            scope.set_wfmoutpre_encoding(current_wfmoutpre_encoding)
            scope.set_wfmoutpre_byte_count(current_byte_count)
            scope.set_wfmoutpre_byte_order(current_byte_order)
            print("  Settings restored")
        except Exception as e:
            print(f"  Could not restore settings (connection may be in bad state): {e}")

        # Test error handling
        print("\n=== Testing Error Handling ===\n")

        # Skip error handling tests if connection is already bad
        try:
            # Quick connectivity check
            scope.query("*OPC?")
        except Exception:
            print("Connection appears to be in bad state, skipping error handling tests\n")
            print("=== Waveform transfer tests completed (with some skipped due to connection issues) ===")
            return

        print("Testing invalid data encoding (should raise ValueError)...")
        try:
            scope.set_data_encoding('INVALID')
            print("  ERROR: Should have raised ValueError!")
        except ValueError as e:
            print(f"  Correctly raised ValueError: {e}")

        print("\nTesting invalid data source (should raise ValueError)...")
        try:
            scope.set_data_source('CH5')
            print("  ERROR: Should have raised ValueError!")
        except ValueError as e:
            print(f"  Correctly raised ValueError: {e}")

        print("\nTesting invalid byte count (should raise ValueError)...")
        try:
            scope.set_wfmoutpre_byte_count(3)
            print("  ERROR: Should have raised ValueError!")
        except ValueError as e:
            print(f"  Correctly raised ValueError: {e}")

        print("\nTesting invalid byte order (should raise ValueError)...")
        try:
            scope.set_wfmoutpre_byte_order('INVALID')
            print("  ERROR: Should have raised ValueError!")
        except ValueError as e:
            print(f"  Correctly raised ValueError: {e}")

        print("\nTesting invalid FastFrame frame number (should raise ValueError)...")
        try:
            scope.set_fastframe_selected(-1)
            print("  ERROR: Should have raised ValueError!")
        except ValueError as e:
            print(f"  Correctly raised ValueError: {e}")

        print("\nTesting invalid channel in acquire_waveform (should raise ValueError)...")
        try:
            scope.acquire_waveform(5)
            print("  ERROR: Should have raised ValueError!")
        except ValueError as e:
            print(f"  Correctly raised ValueError: {e}")

        print("\n=== All waveform transfer tests completed successfully! ===")

    except Exception as e:
        print(f"\nError: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
    finally:
        if scope.is_connected:
            scope.disconnect()
            print("\nDisconnected from oscilloscope")


if __name__ == "__main__":
    main()