#!/usr/bin/env python3
"""
Reliable file transfer over serial console.
Bypasses screen to avoid escaping issues.
"""

import serial
import base64
import hashlib
import time
import sys
import re

def send_command(ser, cmd, wait=0.3):
    """Send a command and wait for output."""
    ser.write(f"{cmd}\n".encode())
    time.sleep(wait)
    response = ser.read(ser.in_waiting or 1).decode('utf-8', errors='replace')
    return response

def wait_for_prompt(ser, timeout=5):
    """Wait for shell prompt."""
    start = time.time()
    buffer = ""
    while time.time() - start < timeout:
        if ser.in_waiting:
            buffer += ser.read(ser.in_waiting).decode('utf-8', errors='replace')
            if buffer.rstrip().endswith('$') or buffer.rstrip().endswith('#'):
                return buffer
        time.sleep(0.1)
    return buffer

def transfer_file(serial_port, local_file, remote_path, baud=115200):
    """Transfer a file reliably over serial."""

    # Read local file
    with open(local_file, 'rb') as f:
        content = f.read()

    # Compute checksums
    local_md5 = hashlib.md5(content).hexdigest()
    b64_content = base64.b64encode(content).decode()
    b64_md5 = hashlib.md5(b64_content.encode()).hexdigest()

    print(f"File size: {len(content)} bytes")
    print(f"Base64 size: {len(b64_content)} bytes")
    print(f"File MD5: {local_md5}")
    print(f"Base64 MD5: {b64_md5}")

    # Open serial port
    print(f"\nOpening {serial_port}...")
    ser = serial.Serial(serial_port, baud, timeout=1)
    time.sleep(0.5)

    # Send initial enter to get prompt
    ser.write(b'\n')
    time.sleep(0.5)
    ser.read(ser.in_waiting)  # Clear buffer

    # Clear target file
    print("Initializing remote file...")
    send_command(ser, f"> /tmp/recv.b64", wait=0.5)

    # Use very small chunks to avoid corruption
    chunk_size = 60
    chunks = [b64_content[i:i+chunk_size] for i in range(0, len(b64_content), chunk_size)]

    print(f"Sending {len(chunks)} chunks...")

    errors = 0
    for i, chunk in enumerate(chunks):
        # Use printf for better escaping than echo
        cmd = f"printf '%s' '{chunk}' >> /tmp/recv.b64"
        ser.write(f"{cmd}\n".encode())
        time.sleep(0.15)  # Small delay between chunks

        # Read any response
        if ser.in_waiting:
            ser.read(ser.in_waiting)

        if (i + 1) % 100 == 0:
            print(f"  Sent {i+1}/{len(chunks)} chunks ({100*(i+1)//len(chunks)}%)")
            # Give Pi time to process
            time.sleep(0.5)

    print("All chunks sent. Waiting for completion...")
    time.sleep(2)

    # Verify base64 file checksum
    print("\nVerifying base64 checksum...")
    ser.read(ser.in_waiting)  # Clear buffer
    send_command(ser, f"md5sum /tmp/recv.b64", wait=1)
    response = wait_for_prompt(ser, timeout=5)
    print(f"Response: {response}")

    if b64_md5 in response:
        print("✓ Base64 checksum matches!")
    else:
        print("✗ Base64 checksum mismatch!")
        # Extract the md5 from response
        match = re.search(r'([a-f0-9]{32})', response)
        if match:
            print(f"  Expected: {b64_md5}")
            print(f"  Got: {match.group(1)}")
        return False

    # Decode to target location
    print(f"\nDecoding to {remote_path}...")
    send_command(ser, f"base64 -d /tmp/recv.b64 > {remote_path}", wait=1)
    time.sleep(1)

    # Verify final file
    print("Verifying final file...")
    ser.read(ser.in_waiting)
    send_command(ser, f"md5sum {remote_path}", wait=1)
    response = wait_for_prompt(ser, timeout=5)
    print(f"Response: {response}")

    if local_md5 in response:
        print("✓ File transfer successful!")
        return True
    else:
        print("✗ Final file checksum mismatch!")
        return False

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: serial_transfer.py <local_file> <remote_path> [serial_port]")
        print("Example: serial_transfer.py wifi_manager.py /opt/dangerous-pi/app/backend/managers/wifi_manager.py")
        sys.exit(1)

    local_file = sys.argv[1]
    remote_path = sys.argv[2]
    serial_port = sys.argv[3] if len(sys.argv) > 3 else "/dev/ttyUSB1"

    success = transfer_file(serial_port, local_file, remote_path)
    sys.exit(0 if success else 1)