# The MIT License (MIT)
# Copyright (c) 2014-2021 Damien P. George
# Copyright (c) 2017 Paul Sokolovsky
# Copyright (c) 2024 yingshaoxo

"""
pyboard interface

This module provides the Pyboard class, used to communicate with and
control the pyboard over a serial USB connection.

Example usage:

    import pyboard
    pyb = pyboard.Pyboard('/dev/ttyACM0')
    pyb.enter_raw_repl()
    pyb.exec('pyb.LED(1).on()')
    print(pyb.exec('print(1+1)'))
    pyb.exit_raw_repl()

To run a script from the local machine on the board and print out the results:

    import pyboard
    pyboard.execfile('test.py', device='/dev/ttyACM0')

This script can also be run directly.  To execute a local script, use:

    python pyboard.py test.py

"""

import time
import os


class Pyboard:
    def __init__(self, serial_device, baudrate=115200):
        # the baudrate is 115200 for micropython project
        try:
            import serial
            self.serial = serial.Serial(serial_device, baudrate=baudrate)
            self._in_waiting = None #Return the number of bytes in the receive buffer.
            if "in_waiting" in dir(self.serial):
                def inWaiting():
                    return self.serial.in_waiting
                self._in_waiting = inWaiting
            else:
                self._in_waiting = self.serial.inWaiting
        except Exception as e:
            #print(e)
            from auto_everything.network_ import Serial
            self.serial = Serial(serial_device, baudrate=baudrate)
            self._in_waiting = self.serial.inWaiting

    def close(self):
        self.serial.close()

    def read_until(self, min_num_bytes, ending, timeout=10):
        data = self.serial.read(min_num_bytes)
        timeout_count = 0
        while True:
            if self._in_waiting() > 0:
                data = data + self.serial.read(self._in_waiting())
                time.sleep(0.01)
                timeout_count = 0
            elif data.endswith(ending):
                break
            else:
                timeout_count += 1
                if timeout_count >= 10 * timeout:
                    break
                time.sleep(0.1)
        return data

    def enter_raw_repl(self):
        # if you do not want to enter the raw shell, you use '\r' to end the input, it treats '\r' as enter key. '\r' == 0x0D, '\n' == 0x0A.
        # but I think you must use raw mode, others won't work well
        # after you get into raw_mode, every input ends with 0x04. but there has no print by default. and for the output, it always starts with '0x4F 0x4B', ends with '0x0D' or '0x04' or '0x3E'. and you should remove all 0x04 from output after first parsing.
        self.serial.write(b'\r\x03') # ctrl-C: interrupt any running program
        self.serial.write(b'\r\x01') # ctrl-A: enter raw REPL
        self.serial.write(b'\x04') # ctrl-D: soft reset
        # the real command is 0x0D, 0x03, 0x0D, 0x01, 0x04
        data = self.read_until(1, b'to exit\r\n>')
        # final result match rule is: starts with '0x4F,0x4B', ends with '0x0D' or '0x04'.
        if not data.endswith(b'raw REPL; CTRL-B to exit\r\n>'):
            print(data)
            raise Exception('could not enter raw repl')

    def exit_raw_repl(self):
        self.serial.write(b'\r\x02') # ctrl-B: enter friendly REPL

    def eval(self, expression):
        ret = self.exec('print({})'.format(expression))
        ret = ret.strip()
        return ret

    def exec(self, command):
        command_bytes = bytes(command, encoding='utf-8')
        for i in range(0, len(command_bytes), 32):
            self.serial.write(command_bytes[i:min(i+32, len(command_bytes))])
            time.sleep(0.01)
        self.serial.write(b'\x04')
        data = self.read_until(2, b'\x04>')
        if not data.endswith(b'\x04>'):
            print(data)
            raise Exception('timeout waiting for EOF reception')
        if data.startswith(b'Traceback') or data.startswith(b'  File '):
            print(data)
            raise Exception('command failed')
        return data[2:-2]

    def exec_without_wait(self, command):
        command_bytes = bytes(command, encoding='utf-8')
        for i in range(0, len(command_bytes), 32):
            self.serial.write(command_bytes[i:min(i+32, len(command_bytes))])
            time.sleep(0.01)
        self.serial.write(b'\x04')

    def execfile(self, filename):
        with open(filename) as f:
            pyfile = f.read()
        return self.exec(pyfile)

    def get_time(self):
        t = str(self.eval('pyb.RTC().datetime()'), encoding='utf-8')[1:-1].split(', ')
        return int(t[4]) * 3600 + int(t[5]) * 60 + int(t[6])

    def read_forever_and_print(self):
        while True:
            if self._in_waiting() > 0:
                data = self.serial.read(self._in_waiting())
                data = data.decode("utf-8", errors="ignore")
                print(data)
            time.sleep(0.1)

    def run(self, code):
        result = self.exec(code)
        if result.endswith(b"\r\n\x04"):
            result = result[:-3]
        return result.decode("utf-8", errors="ignore")

    def list_files_and_folders(self, folder_path="./"):
        script_content = """
import os
print(os.listdir("{folder_path}"))
        """.strip().format(
            folder_path=folder_path,
        )
        data_string = self.run(script_content)
        start_index = data_string.find("[")
        return data_string[start_index:].replace("\',", "\',\n")

    def fs_writefile(self, dest, data, chunk_size=256):
        self.exec("f=open('%s','wb')\nw=f.write" % dest)
        while data:
            chunk = data[:chunk_size]
            self.exec("w(" + repr(chunk) + ")")
            data = data[len(chunk) :]
        self.exec("f.close()\nprint('done')")

    def upload_file(self, source_file_path, target_file_path):
        if not os.path.exists(source_file_path):
            raise Exception("File not exists: {}".format(source_file_path))
        if not os.path.isfile(source_file_path):
            return

        #if os.path.islink(source_file_path):
        #    return

        a_file = open(source_file_path, "rb")
        bytes_data = a_file.read()
        a_file.close()

#try:
#    from machine import freq
#    freq(200000000)
#except Exception as e:
#    pass
        script_content = """
try:
    from gc import collect
    collect()
except Exception as e:
    pass

import os
try:
    os.stat("{folder_path}")
except Exception as e:
    folder_path_splits = "{folder_path}".split("/")
    parent_folder = ""
    for part in folder_path_splits:
        try:
            parent_folder += "/" + part
            os.mkdir(parent_folder)
        except Exception as e:
            pass
        """.strip().format(
            folder_path=os.path.dirname(target_file_path)
        )
        self.exec(script_content)
        self.fs_writefile(target_file_path, bytes_data, chunk_size=256)
        return

        script_content = """
the_bytes_list = [{int_byte_list_string}]

a_file = open("{file_path}", "wb")
a_file.write(bytes(the_bytes_list))
a_file.close()
        """.strip().format(
            int_byte_list_string=",".join([str(one) for one in bytes_data]),
            file_path=target_file_path,
            folder_path=os.path.dirname(target_file_path)
        )
        self.exec(script_content)

    def delete_file_or_folder(self, target_file_path):
        script_content = """
import os

def exists(path):
    try:
        os.stat(path)
        return True
    except Exception as e:
        return False

def file_exists(path):
    try:
        f = open(path, "r")
        f.close()
        return True
    except Exception as e:
        return False

def dir_exists(path):
    try:
        if os.stat(path)[0] & 0x4000:
            return True
        else:
            return False
    except Exception as e:
        return False

def recursive_delete(target_file_path):
    if not exists(target_file_path):
        return

    if file_exists(target_file_path):
        os.remove(target_file_path)
    elif dir_exists(target_file_path):
        sub_list = os.listdir(target_file_path)
        for a_path in sub_list:
            recursive_delete(target_file_path + "/" + a_path)
        os.rmdir(target_file_path)

the_target_file_path = "{target_file_path}"
recursive_delete(the_target_file_path)
        """.strip().format(
            target_file_path=target_file_path.rstrip("/"),
        )
        self.exec(script_content)

    def sync_folder(self, source_folder, target_folder):
        from auto_everything.disk import Disk
        disk = Disk()

        if len(target_folder) != 1:
            target_folder = target_folder.rstrip("/")
        if not target_folder.startswith("/"):
            raise Exception("The target_folder should starts with '/', it is a absolute path")
        if not source_folder.startswith("./"):
            raise Exception("The source_folder should starts with './', it is a relative path")

        if not os.path.exists(source_folder):
            raise Exception("Folder not exists: {}".format(source_folder))
        if os.path.isfile(source_folder):
            return

        try:
            self.delete_file_or_folder(target_folder)
        except Exception as e:
            print(e)

        for path in disk.get_files(source_folder, use_gitignore_file=True):
            if path.startswith("./"):
                if not path.endswith(".swp"):
                    print("In upload:", path)
                    self.upload_file(path, os.path.join(target_folder, path))

    def rsync_folder(self, source_folder, target_folder):
        from auto_everything.disk import Disk
        disk = Disk()

        if len(target_folder) != 1:
            target_folder = target_folder.rstrip("/")
        if not target_folder.startswith("/"):
            raise Exception("The target_folder should starts with '/', it is a absolute path")
        if not source_folder.startswith("./"):
            raise Exception("The source_folder should starts with './', it is a relative path")

        if not os.path.exists(source_folder):
            raise Exception("Folder not exists: {}".format(source_folder))
        if os.path.isfile(source_folder):
            return

        for path in disk.get_files(source_folder, use_gitignore_file=True):
            if path.startswith("./"):
                if not path.endswith(".swp"):
                    target_path = os.path.join(target_folder, path)
                    source_file_size = str(os.stat(path)[6])
                    target_file_size = self.run("import os; print(os.stat('{path}')[6]);".format(path=target_path))
                    if source_file_size != target_file_size:
                        print("In upload:", path)
                        self.upload_file(path, target_path)


def exec_a_file(pyb, filename, device='/dev/ttyACM0'):
    pyb.enter_raw_repl()
    output = pyb.execfile(filename)
    print(str(output, encoding='utf-8'), end='')
    pyb.exit_raw_repl()
    pyb.close()

def run_test_for_pyboard():
    # pi pico is not supported
    device = '/dev/ttyACM0'
    pyb = Pyboard(device)
    pyb.enter_raw_repl()
    print('opened device {}'.format(device))

    print('seconds since boot:', pyb.get_time())

    pyb.exec('def apply(l, f):\r\n for item in l:\r\n  f(item)\r\n')

    pyb.exec('leds=[pyb.LED(l) for l in range(1, 5)]')
    pyb.exec('apply(leds, lambda l:l.off())')

    ## USR switch test
    pyb.exec('switch = pyb.Switch()')

    for i in range(2):
        print("press USR button")
        pyb.exec('while switch(): pyb.delay(10)')
        pyb.exec('while not switch(): pyb.delay(10)')

    print('USR switch passed')

    ## accel test
    if True:
        print("hold level")
        pyb.exec('accel = pyb.Accel()')
        pyb.exec('while abs(accel.x()) > 10 or abs(accel.y()) > 10: pyb.delay(10)')

        print("tilt left")
        pyb.exec('while accel.x() > -10: pyb.delay(10)')
        pyb.exec('leds[0].on()')

        print("tilt forward")
        pyb.exec('while accel.y() < 10: pyb.delay(10)')
        pyb.exec('leds[1].on()')

        print("tilt right")
        pyb.exec('while accel.x() < 10: pyb.delay(10)')
        pyb.exec('leds[2].on()')

        print("tilt backward")
        pyb.exec('while accel.y() > -10: pyb.delay(10)')
        pyb.exec('leds[3].on()')

        print('accel passed')

    print('seconds since boot:', pyb.get_time())

    pyb.exec('apply(leds, lambda l:l.off())')

    pyb.exit_raw_repl()
    pyb.close()

def shell(pyb=None):
    if pyb == None:
        dev_device_id = '/dev/ttyACM0'
        pyb = Pyboard(dev_device_id)
    pyb.enter_raw_repl()

    def get_file_path(input_text):
        _, file_path = input_text.split(" ")
        file_path = file_path.strip("'\"")
        return file_path

    def print_help_function():
        print("""
    list: list files and folders
    monitor: get printed data in real time

    run_main: run main.py in chip
    run "*.py": run a file from computer

    python: enter a mini python

    sync: sync current folder file into pyboard
    upload "*.py": upload a file to pyboard
    delete "*.py": delete a file in pyboard
    rsync: sync current folder changed file into pyboard

    cat "*.py": check a file
    """)

    print_help_function()

    def mini_python():
        print("")
        while True:
            command = input("> ").strip()
            if command == "exit()":
                break
            print(pyb.run(command) + "\n")

    while True:
        command = input("\nyour command (help): ").strip()
        if command == "help":
            print_help_function()
        elif command.startswith("monitor"):
            try:
                pyb.read_forever_and_print()
            except KeyboardInterrupt:
                pass
        elif command.startswith("run_main"):
            pyb.exec_without_wait("import main")
        elif command.startswith("run"):
            file_path = get_file_path(command)
            output = pyb.execfile(file_path)
            print(output)
        elif command == "python":
            try:
                mini_python()
            except KeyboardInterrupt:
                pyb = Pyboard(dev_device_id)
                pyb.exit_raw_repl()
                pyb.enter_raw_repl()
                print_help_function()
        elif command == "list" or command == "ls":
            print(pyb.list_files_and_folders("."))
        elif command == "sync":
            pyb.sync_folder("./", "/")
            print("done")
        elif command == "rsync":
            pyb.rsync_folder("./", "/")
            print("done")
        elif command.startswith("cat "):
            file_path = get_file_path(command)
            print("\n")
            print(pyb.run("""
with open("{name}", "r") as f:
    print(f.read())
    """.format(name=file_path)))
        elif command.startswith("upload "):
            file_path = get_file_path(command)
            pyb.upload_file(file_path, file_path)
            print("done")
        elif command.startswith("delete "):
            file_path = get_file_path(command)
            pyb.delete_file_or_folder(file_path)
            print("done")

    pyb.exit_raw_repl()


if __name__ == "__main__":
    #### there might at least have 2 bugs that causes the fire() function not working
    #from auto_everything.python import Python
    #py = Python()
    #py.fire2(Pyboard)

    shell()

    pyboard = Pyboard("/dev/ttyACM0")
    pyboard.serial.write("print(1+1)".encode("utf-8") + b"\x04")
    while True:
        time.sleep(0.1)
        if pyboard.serial.available() > 0:
            result = pyboard.serial.read(pyboard.serial.available())
            print(result)
            result = result.decode("utf-8")
            print(result)