import yaml

from pytestlab import AutoInstrument
from pytestlab.instruments.backends import ReplayBackend
from pytestlab.instruments.backends import SessionRecordingBackend
from pytestlab.instruments.backends.errors import ReplayMismatchError

print("PyTestLab Record & Replay Tutorial")
print("=" * 40)

import time


def power_supply_characterization(psu):
    """Example measurement procedure to be recorded."""
    print("🔬 Starting power supply characterization...")
    
    # Get instrument ID
    device_id = psu.id()
    print(f"📋 Device: {device_id}")
    
    # Set current limit for safety
    psu.set_current(1, 0.1)  # 100mA limit
    print("⚡ Current limit set to 100mA")
    
    # Enable output
    psu.output(1, True)
    print("🔛 Output enabled")
    
    # Voltage sweep
    voltages = [1.0, 2.0, 3.0, 4.0, 5.0]
    measurements = []
    
    for voltage in voltages:
        print(f"📊 Setting voltage to {voltage}V")
        psu.set_voltage(1, voltage)
        
        # Wait for settling
        time.sleep(0.1)
        
        # Read back values
        measured_v = psu.read_voltage(1)
        measured_i = psu.read_current(1)
        
        measurements.append({
            'set_voltage': voltage,
            'measured_voltage': measured_v,
            'measured_current': measured_i
        })
        
        print(f"  📈 Measured: {measured_v:.3f}V, {measured_i:.3f}A")
    
    # Clean shutdown
    psu.output(1, False)
    psu.set_voltage(1, 0.0)
    print("🔒 Output disabled and voltage reset")
    
    return measurements

def record_measurement_session():
    """Record the measurement procedure."""
    print("🎬 RECORDING PHASE")
    print("-" * 20)
    
    # Create PSU with simulation (in practice, use real instrument address)
    psu = AutoInstrument.from_config("keysight/EDU36311A", simulate=True)
    
    # Wrap with SessionRecordingBackend to capture interactions
    session_file = "tutorial_session.yaml"
    recording_backend = SessionRecordingBackend(
        psu._backend, 
        session_file,
        instrument_key="psu"
    )
    psu._backend = recording_backend
    
    psu.connect_backend()
    
    try:
        # Run the measurement (will be recorded)
        results = power_supply_characterization(psu)
        print(f"\n✅ Recording complete! Session saved to {session_file}")
        return results
    finally:
        psu.close()


# Record the session
recorded_results = record_measurement_session()

# Load and display the session file
with open("tutorial_session.yaml") as f:
    session_data = yaml.safe_load(f)

print("📄 Session File Structure:")
print(yaml.dump(session_data, default_flow_style=False, indent=2)[:1000] + "...")

# Count recorded commands
command_count = len(session_data['psu']['log'])
print(f"\n📊 Total commands recorded: {command_count}")

# Show command types
command_types = {}
for entry in session_data['psu']['log']:
    cmd_type = entry['type']
    command_types[cmd_type] = command_types.get(cmd_type, 0) + 1

print("📈 Command breakdown:")
for cmd_type, count in command_types.items():
    print(f"  {cmd_type}: {count}")

def replay_measurement_session():
    """Replay the recorded measurement procedure."""
    print("\n🎭 REPLAY PHASE")
    print("-" * 20)
    
    # Create ReplayBackend using the recorded session
    replay_backend = ReplayBackend("tutorial_session.yaml")
    
    # Create PSU with replay backend
    psu = AutoInstrument.from_config(
        "keysight/EDU36311A",
        backend_override=replay_backend
    )
    
    psu.connect_backend()
    
    try:
        # Run the exact same measurement (will replay recorded responses)
        results = power_supply_characterization(psu)
        print("\n✅ Replay successful! All commands matched the recording.")
        return results
    finally:
        psu.close()


# Replay the session
replayed_results = replay_measurement_session()

print("🔍 RESULT COMPARISON")
print("-" * 20)

print(f"Recorded results: {len(recorded_results)} measurements")
print(f"Replayed results: {len(replayed_results)} measurements")

all_match = True
for i, (orig, replay) in enumerate(zip(recorded_results, replayed_results, strict=False)):
    v_match = abs(orig['measured_voltage'] - replay['measured_voltage']) < 1e-6
    i_match = abs(orig['measured_current'] - replay['measured_current']) < 1e-6
    
    if v_match and i_match:
        print(f"✅ Step {i + 1}: Perfect match")
    else:
        print(f"❌ Step {i + 1}: Mismatch detected!")
        all_match = False

if all_match:
    print("\n🎉 Perfect replay! All measurements exactly match the recording.")
else:
    print("\n⚠️  Some measurements don't match - check your replay setup.")

def test_sequence_validation():
    """Demonstrate sequence validation by intentionally deviating."""
    print("\n🚨 VALIDATION TESTING")
    print("-" * 20)
    
    replay_backend = ReplayBackend("tutorial_session.yaml")
    psu = AutoInstrument.from_config(
        "keysight/EDU36311A",
        backend_override=replay_backend
    )
    
    psu.connect_backend()
    
    try:
        # Start normally
        device_id = psu.id()
        print(f"✅ Device ID: {device_id}")
        
        psu.set_current(1, 0.1)
        print("✅ Current limit set correctly")
        
        psu.output(1, True)
        print("✅ Output enabled correctly")
        
        # Now try to deviate from the recorded sequence
        print("\n⚠️  Attempting to deviate from recorded sequence...")
        print("   (Recording expects 1.0V, but we'll try 2.5V)")
        
        # This should trigger a ReplayMismatchError
        psu.set_voltage(1, 2.5)  # Recording expects 1.0V first!
        
    except ReplayMismatchError as e:
        print(f"\n🛡️  VALIDATION SUCCESS: {e}")
        print("   The replay system correctly detected the sequence deviation!")
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}")
    finally:
        psu.close()


test_sequence_validation()

# Create a measurement script file for CLI demonstration
script_content = '''
#!/usr/bin/env python3
"""CLI measurement script for replay demonstration."""

def main(bench):
    """Main measurement function called by PyTestLab CLI."""
    psu = bench.psu
    
    print("Starting CLI measurement...")
    
    # Simple voltage measurement
    psu.set_current(1, 0.05)
    psu.output(1, True)
    psu.set_voltage(1, 3.3)
    
    voltage = psu.read_voltage(1)
    current = psu.read_current(1)
    
    psu.output(1, False)
    
    print(f"Measured: {voltage}V, {current}A")
    return {"voltage": voltage, "current": current}
    
if __name__ == "__main__":
    print("Use: pytestlab replay run cli_measurement.py --session session.yaml")
'''

with open("cli_measurement.py", "w") as f:
    f.write(script_content)

print("📝 Created CLI measurement script: cli_measurement.py")

# Create a simple bench configuration
bench_config = '''
bench_name: "Tutorial Bench"
simulate: true
instruments:
  psu:
    profile: "keysight/EDU36311A"
    address: "TCPIP0::192.168.1.100::inst0::INSTR"
'''

with open("tutorial_bench.yaml", "w") as f:
    f.write(bench_config)

print("📝 Created bench configuration: tutorial_bench.yaml")

print("🖥️  CLI COMMAND EXAMPLES:")
print("=" * 25)
print()
print("1️⃣  Record a measurement session:")
print("   pytestlab replay record cli_measurement.py --bench tutorial_bench.yaml --output cli_session.yaml")
print()
print("2️⃣  Replay the recorded session:")
print("   pytestlab replay run cli_measurement.py --session cli_session.yaml")
print()
print("3️⃣  View replay command help:")
print("   pytestlab replay --help")
print()
print("📂 Files created for CLI workflow:")
print("   - cli_measurement.py (measurement script)")
print("   - tutorial_bench.yaml (bench configuration)")
print("   - tutorial_session.yaml (recorded session)")

print("💡 BEST PRACTICES")
print("=" * 17)
print()
print("🎯 Recording Phase:")
print("   • Always record with real instruments for authentic responses")
print("   • Use descriptive session file names (e.g., 'psu_characterization_v2.1.yaml')")
print("   • Include version information in session metadata")
print("   • Test your recording immediately after creation")
print()
print("🔄 Replay Phase:")
print("   • Use identical instrument profiles during recording and replay")
print("   • Handle ReplayMismatchError gracefully in production code")
print("   • Enable debug logging to troubleshoot sequence issues")
print()
print("📁 File Management:")
print("   • Store session files in version control")
print("   • Organize sessions by measurement type or instrument")
print("   • Include README files documenting session purposes")
print()
print("🧪 Testing:")
print("   • Create regression tests using recorded sessions")
print("   • Validate session integrity periodically")
print("   • Use replay mode for CI/CD pipeline testing")