#!/bin/python3

# Simple utility script that enables reading/writing files from the filesystem
# of SB2.

import socket
import json
import argparse
import readline
  
parser = argparse.ArgumentParser(description="Simple SB2 Shell")
parser.add_argument("--ip", default="192.168.77.100", help="IP address to the SB2 to connect to")
parser.add_argument("--port", type=int, default=5, help="Port number to connect to")
parser.add_argument("--debug", action='store_true', help="Enable debug printing")
parser.add_argument("--cmd", help="A single command to run")
args = parser.parse_args()

readline.parse_and_bind("tab: complete")

def print_usage():
  print("Supported commands: <read/write/ls/rm/update>")

class Ret:
  def __init__(self, ok, msg):
    self.ok = ok
    self.msg = msg

def send_json_udp(ip_address, port, fs_path, json_path):
  try:
    # Load JSON data from the file
    with open(json_path, 'r') as json_file:
        json_data = json.load(json_file)
    
    # Convert JSON data to a string
    json_string = json.dumps(json_data)

    # Create a UDP socket
    udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    # Send the JSON data as a UDP packet
    udp_socket.sendto(f"write {fs_path} {json_string}".encode('utf-8'), (ip_address, port))
    data, _ = udp_socket.recvfrom(8192)
    return Ret(True, data.decode('utf-8'))

    
  except Exception as e:
    return Ret(False, e)


def send_config_update(ip_address, port, json_path):
  try:
    # Load JSON data from the file
    with open(json_path, 'r') as json_file:
        json_data = json.load(json_file)
    
    # Convert JSON data to a string
    json_string = json.dumps(json_data)

    # Create a UDP socket
    udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    # Send the JSON data as a UDP packet
    udp_socket.sendto(f"update config {json_string}".encode('utf-8'), (ip_address, port))
    data, _ = udp_socket.recvfrom(8192)
    return Ret(True, data.decode('utf-8'))

    
  except Exception as e:
    return Ret(False, e)

def read_json_udp(ip_address, port, path):
  # Create a UDP socket
  udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

  # Send the JSON data as a UDP packet
  udp_socket.sendto(f"read {path}".encode('utf-8'), (ip_address, port))
  data, _ = udp_socket.recvfrom(8192)
  try:
    parsed = json.loads(data.decode('utf-8'))
    return Ret(True, json.dumps(parsed, indent=2))
  except Exception as e:
    return Ret(False, data.decode('utf-8'))


def ls(ip_address, port, path):
  # Create a UDP socket
  udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

  # Send the JSON data as a UDP packet
  udp_socket.sendto(f"ls {path}".encode('utf-8'), (ip_address, port))
  data, _ = udp_socket.recvfrom(8192)
  return Ret(True, data.decode())

def rm(ip_address, port, path):
  udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

  udp_socket.sendto(f"rm {path}".encode('utf-8'), (ip_address, port))
  data, _ = udp_socket.recvfrom(8192)
  return Ret(True, data.decode())

def handle_cmd(argv):
  # Check if the correct number of command-line arguments is provided
  if args.debug:
    print(f"Got cmd {argv}")

  if len(argv) == 2:
    cmd = argv[0]
    if cmd == "ls":
      path = argv[1]
      return ls(args.ip, args.port, path)
    elif cmd == "read":
      path = argv[1]
      return read_json_udp(args.ip, args.port, path)
    elif cmd == "rm":
      path = argv[1]
      return rm(args.ip, args.port, path)
    else:
      return Ret(False, "Did not recognize command")

  elif len(argv) == 3:
    cmd = argv[0]
    if cmd == "write":
      fsPath = argv[1]
      jsonPath = argv[2]
      return send_json_udp(args.ip, args.port, fsPath, jsonPath)
    
    elif cmd == "update":
      second = argv[1]
      if second == "config":
        jsonPath = argv[2]
        return send_config_update(args.ip, args.port, jsonPath)
      else:
        return Ret(False, "Did not recognize command")
  return Ret(False, "Unrecognized command")


def shell():
  print(f"Connecting to SB2 @ {args.ip}:{args.port}")
  ret = handle_cmd(["ls", "/"])
  if not ret.ok:
    print("Could not connect ...")
    return

  try:
    while True:
      user_input = input("Senti> ")
        
      if user_input.lower() == 'exit':
        break

      ret = handle_cmd(user_input.split(" "))
      if ret.ok:
        print(ret.msg)
        # Add the command to history
        readline.add_history(user_input)
      else:
        print(f"[1] {ret.msg}")

  except KeyboardInterrupt:
    print("\nShell canceled by user (Ctrl+C)")

if __name__ == "__main__":
  if args.cmd:
    ret = handle_cmd(args.cmd.split(" "))
    if ret.ok:
      print(ret.msg)
    else:
      print(f"[1] {ret.msg}")
  else:
    shell()