#!/usr/bin/env python3
"""
CANBUS OTA Firmware Update Tool

This tool provides a command-line interface to update ESP32 firmware
over CANBUS via the Battery Pack Manager. The host broadcasts firmware
to all sysbus peers simultaneously.

Upload protocol: base64-chunked transfer over serial console
  1. ota begin <size> <sha256_hex>
  2. ota data <offset> <base64_chunk>  (repeated, 512-byte raw chunks)
  3. ota end

The serial port and firmware image are auto-detected when omitted: the port
via the connected ESP32 USB device, and the image from the installed update
package (/usr/share/signalytic-battery-manager/).

Usage:
    signalytic_bmtool check                     # update only if out of date
    signalytic_bmtool upload [firmware.bin]
    signalytic_bmtool update [firmware.bin]
    signalytic_bmtool push
    signalytic_bmtool status [--staging]
    signalytic_bmtool abort
    signalytic_bmtool --port /dev/ttyACM0 ...   # override auto-detected port

(c) 2026 Signalytic, all rights reserved
"""

import argparse
import base64
import hashlib
import json
import sys
import time
from pathlib import Path

try:
    import serial
except ImportError as e:
    print(f"Missing dependency: {e}")
    print("Install with: pip install pyserial")
    sys.exit(1)

import os
for _p in (
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '..', 'scripts', 'signalytic', 'lib'),
    '/usr/lib/signalytic_bm',
):
    if os.path.isfile(os.path.join(_p, 'serial_console.py')):
        sys.path.insert(0, _p)
        break
from serial_console import open_console, discover_port


UPLOAD_CHUNK_SIZE = 512  # Max raw bytes per ota data command

# Default install locations for the update (OTA) package. Note these differ
# from signalytic_bmproduction, which reads from the "-production" share dir.
SHARE_DIR = '/usr/share/signalytic-battery-manager'
DEFAULT_FIRMWARE = os.path.join(SHARE_DIR, 'firmware', 'firmware.bin')
VERSION_FILE = os.path.join(SHARE_DIR, 'VERSION')


def connect(ser):
    """Flush stale data and confirm the console is responsive."""
    ser.reset_input_buffer()
    ser.reset_output_buffer()
    for attempt in range(10):
        if send_command(ser, '?') == 'OK':
            return
    raise ConnectionError("Console did not respond with OK after 10 attempts")


def send_command(ser, cmd: str) -> str:
    """Send a command and read the response (skipping ESP-IDF log lines)."""
    ser.reset_input_buffer()
    ser.write((cmd + '\n').encode('utf-8'))
    ser.flush()

    # Read response lines until we get OK, ERR, or JSON. The per-read timeout
    # is ser.timeout (set when the port was opened); bound the total wait by it.
    lines = []
    start_time = time.time()
    while time.time() - start_time < ser.timeout:
        line = ser.readline().decode('utf-8', errors='replace').strip()
        if not line:
            continue
        # Skip ESP-IDF log lines (e.g. "I (770912) ota_host: ...")
        if len(line) > 1 and line[0] in 'EWIDV' and line[1] == ' ':
            print(f" >> {line}", file=sys.stderr)
            continue
        lines.append(line)
        break

    return '\n'.join(lines)


def upload(ser, filepath: str) -> dict:
    """Upload firmware via base64-chunked protocol."""
    path = Path(filepath)
    if not path.exists():
        raise FileNotFoundError(f"File not found: {filepath}")

    data = path.read_bytes()
    filesize = len(data)
    sha256_hex = hashlib.sha256(data).hexdigest()

    print(f"Uploading {filepath} ({filesize} bytes, sha256={sha256_hex[:16]}...)")

    # Step 1: ota begin
    response = send_command(ser, f'ota begin {filesize} {sha256_hex}')
    if not response.startswith('OK'):
        raise Exception(f"ota begin failed: {response}")

    # Step 2: ota data chunks
    offset = 0
    while offset < filesize:
        chunk = data[offset:offset + UPLOAD_CHUNK_SIZE]
        b64_chunk = base64.b64encode(chunk).decode('ascii')

        response = send_command(ser, f'ota data {offset} {b64_chunk}')
        if not response.startswith('OK'):
            raise Exception(f"ota data failed at offset {offset}: {response}")

        offset += len(chunk)
        pct = offset * 100 // filesize
        print(f"\r  {offset}/{filesize} bytes ({pct}%)", end='', flush=True)
    print()

    # Step 3: ota end
    response = send_command(ser, 'ota end')
    if not response.startswith('OK'):
        raise Exception(f"ota end failed: {response}")

    # Parse size and sha256 from "OK size=... sha256=..."
    result = {'status': 'ok'}
    parts = response.split()
    for part in parts[1:]:
        if '=' in part:
            key, value = part.split('=', 1)
            result[key] = value
    return result


def push(ser) -> bool:
    """Broadcast staged firmware to all sysbus peers."""
    print("Broadcasting firmware to all peers...")

    ser.write(b'otapush\n')
    ser.flush()

    # Read progress updates
    while True:
        line = ser.readline().decode('utf-8', errors='replace').strip()
        if not line:
            continue

        if line.startswith('Block') or 'Block' in line:
            print(f"\r{line}", end='', flush=True)
        elif line.startswith('OK'):
            print()
            return True
        elif line.startswith('ERR'):
            print()
            raise Exception(line)
        elif line.startswith('Broadcasting'):
            print(line)
        else:
            print(line)


def status(ser, staging: bool = False) -> dict:
    """Query OTA status — peer info or staging info."""
    if staging:
        response = send_command(ser, 'otastatus --staging')
    else:
        response = send_command(ser, 'otastatus')

    if response.startswith('ERR'):
        raise Exception(response)

    try:
        return json.loads(response)
    except json.JSONDecodeError:
        return {'raw': response}


def abort(ser) -> bool:
    """Broadcast OTA abort to all targets."""
    response = send_command(ser, 'otaabort')
    return response.startswith('OK')


def available_version() -> str:
    """Read the available firmware version from the package VERSION file."""
    with open(VERSION_FILE) as f:
        return f.read().strip()


def running_version(ser, retries: int = 3, delay: float = 2.0) -> str:
    """Query the running firmware version from the device via 'bmversion'.

    Retries briefly in case the device is still booting. Raises if no valid
    version is obtained.
    """
    version = ''
    for attempt in range(retries):
        response = send_command(ser, 'bmversion')
        if response and not response.startswith('ERR'):
            version = response.strip()
            break
        if attempt < retries - 1:
            time.sleep(delay)
    if not version:
        raise Exception("could not read running version from device (bmversion)")
    return version


def cmd_upload(args, ser):
    """Handle upload command."""
    result = upload(ser, args.firmware)
    print(f"Upload complete:")
    print(f"  Size: {result.get('size', 'unknown')} bytes")
    print(f"  SHA256: {result.get('sha256', 'unknown')}")


def cmd_push(args, ser):
    """Handle push command."""
    push(ser)
    print("Push complete")


def cmd_status(args, ser):
    """Handle status command."""
    result = status(ser, staging=args.staging)
    print(json.dumps(result, indent=2))


def cmd_abort(args, ser):
    """Handle abort command."""
    if abort(ser):
        print("Abort sent")
    else:
        print("Abort failed")


def cmd_update(args, ser):
    """Handle update command (upload then push)."""
    result = upload(ser, args.firmware)
    print(f"Upload complete:")
    print(f"  Size: {result.get('size', 'unknown')} bytes")
    print(f"  SHA256: {result.get('sha256', 'unknown')}")
    push(ser)
    print("Update complete")


def cmd_check(args, ser):
    """Compare available vs running version and update only if they differ."""
    avail = available_version()
    running = running_version(ser)
    print(f"Available version: {avail}")
    print(f"Running version:   {running}")

    if running == avail:
        print("Device is up to date; no update needed")
        return

    print(f"Version mismatch (running {running} != available {avail}); updating...")
    result = upload(ser, args.firmware)
    print(f"Upload complete:")
    print(f"  Size: {result.get('size', 'unknown')} bytes")
    print(f"  SHA256: {result.get('sha256', 'unknown')}")
    push(ser)
    print("Update complete")


def main():
    parser = argparse.ArgumentParser(
        description='CANBUS OTA Firmware Update Tool',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Check the device and update it only if its version differs from the
  # installed package (auto-detects port and firmware image)
  %(prog)s check

  # Upload the default firmware image to staging (auto-detected port)
  %(prog)s upload

  # Upload a specific firmware image to staging on a specific port
  %(prog)s --port /dev/ttyACM0 upload build/signalytic-battery-manager.bin

  # Broadcast firmware to all sysbus peers
  %(prog)s push

  # Upload firmware then broadcast to all peers in one step
  %(prog)s update

  # Show peer firmware versions and OTA status
  %(prog)s status

  # Check staging partition status
  %(prog)s status --staging

  # Abort all transfers
  %(prog)s abort
"""
    )

    parser.add_argument('--port', '-p', default=None,
                        help='Serial port (e.g., /dev/ttyACM0). Auto-detected if omitted.')
    parser.add_argument('--baudrate', '-b', type=int, default=115200,
                        help='Serial baudrate (default: 115200)')
    parser.add_argument('--timeout', '-t', type=float, default=60.0,
                        help='Command timeout in seconds (default: 60)')

    subparsers = parser.add_subparsers(dest='command', required=True)

    # upload command
    upload_parser = subparsers.add_parser('upload', help='Upload firmware to staging')
    upload_parser.add_argument('firmware', nargs='?', default=DEFAULT_FIRMWARE,
                               help=f'Path to firmware binary (default: {DEFAULT_FIRMWARE})')

    # push command (no arguments — broadcasts to all peers)
    subparsers.add_parser('push', help='Broadcast firmware to all sysbus peers')

    # status command
    status_parser = subparsers.add_parser('status', help='Query OTA status')
    status_parser.add_argument('--staging', '-s', action='store_true',
                               help='Show staging partition info instead of peer status')

    # abort command (no arguments — always broadcasts)
    subparsers.add_parser('abort', help='Broadcast abort to all targets')

    # update command (upload then broadcast)
    update_parser = subparsers.add_parser('update', help='Upload firmware then broadcast to all peers')
    update_parser.add_argument('firmware', nargs='?', default=DEFAULT_FIRMWARE,
                               help=f'Path to firmware binary (default: {DEFAULT_FIRMWARE})')

    # check command (update only if running version differs from available)
    check_parser = subparsers.add_parser(
        'check', help='Update the device only if its running version differs from the available version')
    check_parser.add_argument('firmware', nargs='?', default=DEFAULT_FIRMWARE,
                              help=f'Path to firmware binary (default: {DEFAULT_FIRMWARE})')

    args = parser.parse_args()

    handlers = {
        'upload': cmd_upload,
        'push': cmd_push,
        'status': cmd_status,
        'abort': cmd_abort,
        'update': cmd_update,
        'check': cmd_check,
    }

    # Detect and open the serial port once; pass the serial object to handlers.
    port = args.port or discover_port()
    if not port:
        print("Error: no BPM console device found (no /dev/serial/by-id match)")
        sys.exit(1)
    print(f"Using serial port {port}")

    try:
        # open_console blocks until the console is free and closes the port /
        # releases the lock on exit.
        with open_console(port=port, baudrate=args.baudrate,
                          timeout=args.timeout, write_timeout=args.timeout,
                          contention='wait') as ser:
            connect(ser)
            handlers[args.command](args, ser)

    except KeyboardInterrupt:
        print("\nInterrupted")
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()
