#!/usr/bin/env python3
import time
import os
from dataclasses import dataclass
from typing import Optional, Callable, List
from datetime import datetime
from threading import Thread
from enum import Enum
from steamvr.paths import LH_CONSOLE_PATH
from steamvr.unbuffered.process import UnbufferedProcess


def static_vars(**kwargs):
    def decorate(func):
        for k in kwargs:
            setattr(func, k, kwargs[k])
        return func

    return decorate


@static_vars(lines=[], buffer='')
def readline(pipe):
    if readline.lines:
        return readline.lines.pop(0)

    while True:
        data = os.read(pipe, 1024)
        if not data:
            break
        readline.buffer += data.decode('utf-8')
        if '\r\n' in readline.buffer:
            readline.lines = readline.buffer.split('\r\n')
            readline.buffer = readline.lines.pop()
            return readline.lines.pop(0)


class LhProcess(Enum):
    SENSORCHECK = b"sensorcheck\r\n"
    IMU = b"imu\r\n"
    MONITOR = b""


class LighthouseProcess(UnbufferedProcess):
    def __init__(self, lh_process: LhProcess):
        super().__init__(LH_CONSOLE_PATH)
        self.lh_process = lh_process

    def run(self):
        super().run()
        try:
            os.write(self.stdin, self.lh_process.value)
            os.write(self.stdin, b"dump\r\n")
        except OSError:
            print("Device not connected")
            self.run()

    def write_cmd(self, str: str):
        os.write(self.stdin, str.encode('utf-8'))

    def switch_device(self, serial: str):
        self.write_cmd("serial " + serial + "\r\n")


@dataclass(frozen=True)
class Sensor:
    SensorID: str
    HitCount: str
    MaxWidthTicks: str
    SensorIndex: str
    DemodCount: str
    Example: str


@dataclass
class SensorcheckResult:
    sensors: List[Optional[Sensor]]
    timestamp: datetime

    def seen_sensor_count(self) -> int:
        return sum(1 for s in self.sensors if s is not None)


def read_sensor_lines(pipe):
    sensors = SensorcheckResult(sensors=[], timestamp=datetime.now())
    for i in range(32):
        line = readline(pipe)
        split = line.split('\t\t')
        if len(split) == 6:
            sensors.sensors.append(Sensor(*split))
        else:
            sensors.sensors.append(None)
    return sensors


class ConnectionMonitor(Thread):
    def __init__(self):
        super().__init__()
        self.on_connect = None
        self.on_disconnect = None
        self._p = LighthouseProcess(LhProcess.MONITOR)
    def run(self):
        self._p.run()
        while True:
            line = readline(self._p.stdout)

            if 'Successfully' in line:
                self.on_connect(line.split()[0][:-1])
            elif "Lighthouse IMU HID device error" in line:
                if readline(self._p.stdout) == "Lighthouse Optical HID device error":
                    if readline(self._p.stdout) == "Lighthouse VrController HID device error":
                        self.on_disconnect()

    def register_on_connect(self, on_connect: Callable[[str], None]):
        self.on_connect = on_connect

    def register_on_disconnect(self, on_disconnect: Callable[[str], None]):
        self.on_disconnect = on_disconnect


class SensorcheckThread(Thread):
    def __init__(self,
                 on_sensorcheck: Callable[[SensorcheckResult], None],
                 on_connect: Callable[[str], None],
                 on_disconnect: Callable[[], None],
                 on_error: Callable[[Exception], None]):
        super().__init__()
        self.on_sensorcheck = on_sensorcheck
        self.on_connect = on_connect
        self.on_disconnect = on_disconnect
        self.on_error = on_error
        self.process = LighthouseProcess(LhProcess.SENSORCHECK)
        self.process.run()
        self.is_disconnected = False

    def run(self):
        while True:
            line = readline(self.process.stdout)
            if line is None:
                self.process.run()
                continue
            if 'Successfully' in line:
                self.on_connect(line.split()[0][:-1])
            if 'SensorID' in line:
                self.on_sensorcheck(read_sensor_lines(self.process.stdout))
            elif line == "Lighthouse IMU HID device error":
                if readline(self.process.stdout) == "Lighthouse Optical HID device error":
                    if readline(self.process.stdout) == "Lighthouse VrController HID device error":
                        self.on_disconnect()
            elif line == "No connected lighthouse device found.":
                pass
            # else:
            #    self.on_error(Exception("Unexpected output" + line))


@dataclass(frozen=True)
class ImuResult:
    timestamp: float
    unknown: int
    gyro: List[float]
    accel: List[float]
    host_timestamp: datetime


def read_imu_line(line) -> Optional[ImuResult]:
    assert 'accel' in line
    split = line.split()
    if len(split) == 10:
        try:
            imu_result = ImuResult(timestamp=float(split[0]), unknown=int(split[1]),
                                   gyro=[float(x) for x in split[3:6]],
                                   accel=[float(x) for x in split[7:10]],
                                   host_timestamp=datetime.now())
        except ValueError:
            print(split)
            return None
        return imu_result
    return None


class ImuThread(Thread):
    def __init__(self,
            on_imu: Callable[[ImuResult], None],
            on_error: Callable[[Exception], None]):
        super().__init__()
        self.on_imu = on_imu
        self.on_error = on_error

        self.process = LighthouseProcess(LhProcess.IMU)
        self.process.run()

    def run(self):
        while True:
            line = readline(self.process.stdout)
            if line == 'Attached lighthouse receiver devices: 0' or \
                    line == 'No connected lighthouse device found.' or \
                    line == 'No devices connected':
                print("No devices connected")
                time.sleep(1)
                self.process.run()
            if 'accel' in line:
                imu_result = read_imu_line(line)
                if imu_result is not None:
                    self.on_imu(imu_result)
            else:
                print(line)
            # self.process.run()


# if a device is unplugged it prints the following, lighthouse needs to be restarted
# Lighthouse IMU HID device error
# Lighthouse Optical HID device error
# Lighthouse VrController HID device error


def on_sensorcheck_cb(result: SensorcheckResult):
    print(result.timestamp, result.seen_sensor_count())


def on_imu_cb(result: ImuResult):
    print(result.host_timestamp, result.accel)


def main():
    SensorcheckThread(on_sensorcheck_cb, lambda: print("Device Disconnected"),
                      lambda: print("Device Connected"),
                      lambda e: print(e)).start()

    # ImuThread(on_imu_cb, lambda e: print(e)).start()


if __name__ == '__main__':
    main()