#!/bin/bash
# Fix pyrf Python 3 compatibility
# Run this after: pip install pyrf

set -e

VENV_DIR="${1:-venv}"
PYRF_BASE="$VENV_DIR/lib/python3.12/site-packages/pyrf"

if [ ! -d "$PYRF_BASE" ]; then
    echo "❌ pyrf not found at $PYRF_BASE"
    echo "Usage: $0 [venv_directory]"
    exit 1
fi

echo "🔧 Fixing pyrf for Python 3..."

# Backup originals
cp "$PYRF_BASE/devices/thinkrf.py" "$PYRF_BASE/devices/thinkrf.py.bak" 2>/dev/null || true
cp "$PYRF_BASE/devices/thinkrf_properties.py" "$PYRF_BASE/devices/thinkrf_properties.py.bak" 2>/dev/null || true
cp "$PYRF_BASE/connectors/blocking.py" "$PYRF_BASE/connectors/blocking.py.bak" 2>/dev/null || true

# Fix thinkrf.py
echo "  Fixing thinkrf.py..."
sed -i 's/\.iteritems()/.items()/g' "$PYRF_BASE/devices/thinkrf.py"
sed -i 's/raw_input/input/g' "$PYRF_BASE/devices/thinkrf.py"
# Fix print statements (carefully to handle the format string)
sed -i '884s/.*/            print(fmt % (index, wsa["HOST"], modelstring, wsa["SERIAL"]))/' "$PYRF_BASE/devices/thinkrf.py"
sed -i 's/print "r) Refresh"/print("r) Refresh")/g' "$PYRF_BASE/devices/thinkrf.py"
sed -i 's/print "q) Abort"/print("q) Abort")/g' "$PYRF_BASE/devices/thinkrf.py"
sed -i 's/print "error: invalid selection: '\''%s'\''" % choice/print("error: invalid selection: '\''%s'\''" % choice)/g' "$PYRF_BASE/devices/thinkrf.py"

# Fix thinkrf_properties.py
echo "  Fixing thinkrf_properties.py..."
sed -i 's/\.iteritems()/.items()/g' "$PYRF_BASE/devices/thinkrf_properties.py"

# Fix blocking.py (socket bytes issue)
echo "  Fixing blocking.py..."
sed -i '29s/self._sock_scpi.send(cmd)/self._sock_scpi.send(cmd.encode())/' "$PYRF_BASE/connectors/blocking.py"
sed -i '34s/self._sock_scpi.send(cmd)/self._sock_scpi.send(cmd.encode())/' "$PYRF_BASE/connectors/blocking.py"
# Fix line 37 - replace entire line to avoid double decode
sed -i '37s/.*/        return buf.decode()/' "$PYRF_BASE/connectors/blocking.py"

echo "✅ pyrf fixed for Python 3!"
echo "   Backups saved with .bak extension"