import socket
import msgpack
import numpy as np
import cmd
import sys

class PicoClient(cmd.Cmd):
    intro = 'Welcome to Pico ADC client. Type help or ? to list commands.\n'
    prompt = '(pico) '

    def __init__(self, host='localhost', port=5003):
        super().__init__()
        self.host = host
        self.port = port
        self.sock = None
        self.magic = 0xAA
        self.connected = False
        self.pk = msgpack.Packer()

    def connect(self):
        if self.connected:
            print("Already connected.")
            return
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((self.host, self.port))
            self.connected = True
            print(f"Connected to {self.host}:{self.port}")
        except Exception as e:
            print(f"Connection failed: {e}")

    def disconnect(self):
        if self.sock:
            self.sock.close()
            self.sock = None
            self.connected = False
            print("Disconnected.")

    def send_command(self, cmd_code, *args):
        if not self.connected:
            print("Not connected. Use 'connect' first.")
            return None
        try:
            buf = bytearray()
            buf.extend(self.pk.pack(0xAA))
            buf.extend(self.pk.pack(cmd_code))
            for arg in args:
                buf.extend(self.pk.pack(arg))
            self.sock.sendall(buf)
            return self.receive_response()
        except Exception as e:
            print(f"Send failed: {e}")
            return None

    def receive_response(self):
        try:
            unpacker = msgpack.Unpacker()
            while True:
                chunk = self.sock.recv(8192)
                if not chunk:
                    break
                unpacker.feed(chunk)
                
                objects = []
                for obj in unpacker:
                    objects.append(obj)
                print(f"Code: {objects[1]}")

                return objects[1], objects[2:]
        except Exception as e:
            print(f"Receive failed: {e}")
            return None

    def do_connect(self, arg):
        """Connect to the server."""
        self.connect()

    def do_disconnect(self, arg):
        """Disconnect from the server."""
        self.disconnect()

    def do_open(self, arg):
        """Open the Pico device."""
        resp = self.send_command(0x01)
        if resp:
            code, data = resp
            if code == 0xCC:
                print("Device opened successfully.")
            else:
                print(f"Error: {code}, {data}")

    def do_close(self, arg):
        """Close the Pico device."""
        resp = self.send_command(0x03)
        if resp:
            code, data = resp
            if code == 0xCC:
                print("Device closed successfully.")
            else:
                print(f"Error: {code}, {data}")

    def do_status(self, arg):
        """Get status."""
        resp = self.send_command(0x30)
        if resp:
            code, data = resp
            if code == 0xFF:
                print(f"Error: {code}, {data}")
            else:
                print(f"Status: {code}")

    def do_get_params(self, arg):
        """Get current parameters."""
        resp = self.send_command(0x04)
        if resp:
            code, data = resp
            if code == 0xCD and data:
                params = data
                sample_rate, measurement_times, number_channels, trig_channel, threshold, th_direction, trig_autoTrigger_ms = params
                print(f"Sample rate: {sample_rate}")
                print(f"Measurement times: {measurement_times}")
                print(f"Number channels: {number_channels}")
                print(f"Trig channel: {trig_channel}")
                print(f"Threshold: {threshold}")
                print(f"TH direction: {th_direction}")
                print(f"Trig auto trigger ms: {trig_autoTrigger_ms}")
            else:
                print(f"Error: {code}, {data}")

    def do_set_rate(self, arg):
        """Set sample rate. Usage: set_rate <rate>"""
        try:
            rate = int(arg)
            resp = self.send_command(0x07, rate)
            if resp:
                code, data = resp
                if code == 0xCC:
                    print("Sample rate set.")
                else:
                    print(f"Error: {code}, {data}")
        except ValueError:
            print("Invalid rate.")

    def do_set_points(self, arg):
        """Set points vector. Usage: set_points <p1> <p2> ..."""
        try:
            points = [int(x) for x in arg.split()]
            resp = self.send_command(0x06, points)
            if resp:
                code, data = resp
                if code == 0xCC:
                    print("Points set.")
                else:
                    print(f"Error: {code}, {data}")
        except ValueError:
            print("Invalid points.")

    def do_set_times(self, arg):
        """Set measurement times. Usage: set_times <times>"""
        try:
            times = int(arg)
            resp = self.send_command(0x08, times)
            if resp:
                code, data = resp
                if code == 0xCC:
                    print("Times set.")
                else:
                    print(f"Error: {code}, {data}")
        except ValueError:
            print("Invalid times.")

    def do_channel_configure(self, arg):
        """Configure channels. Usage: channel_configure <num_channels> <range1> <range2> ... <trig_channel>"""
        try:
            parts = arg.split()
            num_channels = int(parts[0])
            ranges = [int(x) for x in parts[1:1+num_channels]]
            trig_channel = int(parts[1+num_channels])
            resp = self.send_command(0x09, num_channels, ranges, trig_channel)
            if resp:
                code, data = resp
                if code == 0xCC:
                    print("Channels configured.")
                else:
                    print(f"Error: {code}, {data}")
        except (ValueError, IndexError):
            print("Invalid arguments. Usage: channel_configure <num> <ranges...> <trig>")

    def do_trigger_configure(self, arg):
        """Configure trigger. Usage: trigger_configure <th_direction> <threshold> <auto_trigger_ms>"""
        try:
            th_dir, thresh, auto_ms = [int(x) for x in arg.split()]
            resp = self.send_command(0x0A, th_dir, thresh, auto_ms)  # Assuming 0x0A
            if resp:
                code, data = resp
                if code == 0xCC:
                    print("Trigger configured.")
                else:
                    print(f"Error: {code}, {data}")
        except ValueError:
            print("Invalid arguments.")

    def do_set_premeasurement(self, arg):
        """Set premeasurement percentage. Usage: set_premeasurement <percentage>"""
        try:
            perc = int(arg)
            resp = self.send_command(0x0B, perc)  # Assuming 0x0B
            if resp:
                code, data = resp
                if code == 0xCC:
                    print("Premeasurement set.")
                else:
                    print(f"Error: {code}, {data}")
        except ValueError:
            print("Invalid percentage.")

    def do_measurement(self, arg):
        """Start measurement."""
        resp = self.send_command(0x20)
        if resp:
            code, data = resp
            if code == 0xCC:
                print("Measurement started.")
            else:
                print(f"Error: {code}, {data}")

    def do_stop(self, arg):
        """Stop measurement."""
        resp = self.send_command(0x0E)
        if resp:
            code, data = resp
            if code == 0xCC:
                print("Measurement stopped.")
            else:
                print(f"Error: {code}, {data}")

    def do_send_points(self, arg):
        """Send points (after measurement)."""
        if not self.connected:
            print("Not connected.")
            return
        try:
            # Receive data packets
            unpacker = msgpack.Unpacker()
            data_collected = {}
            buf = bytearray()
            buf.extend(self.pk.pack(0xAA))
            buf.extend(self.pk.pack(0xBB))
            self.sock.sendall(buf)
            while True:
                resp_code, data = self.receive_response()
                print("Chunk recieved!")
                print(f'Resp: {resp_code} (need {0xCD} or {0xCE})')
                if resp_code == 0xCD:
                    i, j, packet, size, packet_data = data[0], data[1], data[2], data[3], data[4]
                    key = (i, j)
                    print("Received packet params.")
                    if key not in data_collected:
                        data_collected[key] = []
                        print("Key created!")
                    data_collected[key].extend(packet_data)  # Assuming packet is list of int16
                    print("Appended!")
                    buf = bytearray()
                    buf.extend(self.pk.pack(0xAA))
                    buf.extend(self.pk.pack(0x3D))
                    self.sock.sendall(buf)
                elif resp_code == 0xCE:
                    # Assemble data
                    for key, full_data in data_collected.items():
                        arr = np.array(full_data, dtype=np.int16)
                        print(f"Data for measurement {key[0]}, channel {key[1]}: shape {arr.shape}, mean {arr.mean():.2f}")
                        # Optionally save to file
                        # np.save(f"data_{key[0]}_{key[1]}.npy", arr)
                    # Send continue
                    buf = bytearray()
                    buf.extend(self.pk.pack(0xAA))
                    buf.extend(self.pk.pack(0x3E))
                    self.sock.sendall(buf)
                    print("Data received and acknowledged.")
                    return
                else:
                    print(f"Unexpected response: {resp_code}")
                    return
        except Exception as e:
            print(f"Send points failed: {e}")

    def do_quit(self, arg):
        """Quit the client."""
        self.disconnect()
        return True

if __name__ == '__main__':
    if len(sys.argv) > 1:
        host = sys.argv[1]
        port = int(sys.argv[2]) if len(sys.argv) > 2 else 12345
        client = PicoClient(host, port)
    else:
        client = PicoClient()
    client.cmdloop()