import importlib
import logging
import multiprocessing
import os
import shutil
import subprocess
import tomllib
from enum import Enum, auto
from io import TextIOWrapper
from pathlib import Path

import click
import platformdirs
import zmq
from zmq import Socket

from flandre import C
from flandre import P
from flandre.BusClient import BusClient
from flandre.kde_pyqt6_mainui import kde_pyqt6_mainui
from flandre.nodes.Broker import Broker
from flandre.nodes.Device import Device, DeviceCmd
from flandre.utils.Msg import KillMsg, NodeOnlineMsg, Msg1, Msg2
from flandre.utils.RfFrame import b2t
from flandre.utils.mi import MiSwitch


class LaunchComponent(Enum):
    MainUI = auto()
    ImageCV = auto()
    Muxer = auto()
    Robot = auto()
    Loader = auto()
    Device = auto()
    Recorder = auto()
    Beamformer = auto()
    ImageFFMPEG = auto()
    ImageQt = auto()
    Midi = auto()
    Mi = auto()
    Web = auto()
    VideoQt = auto()


def launch(arg: dict[LaunchComponent, dict]):
    logging.basicConfig(level=logging.INFO)
    multiprocessing.set_start_method('spawn')
    bp = multiprocessing.Process(target=Broker(broker=True), kwargs=dict(software_config=C))
    bp.start()
    ps = []
    for k, v in arg.items():
        if k == LaunchComponent.MainUI and os.environ.get('XDG_CURRENT_DESKTOP', None) == 'KDE':
            ps.append(kde_pyqt6_mainui)
            continue
        m = importlib.import_module(f"flandre.nodes.{k.name}")
        c = getattr(m, k.name)
        ps.append(c(**v))

    pps = []
    for p in ps:
        pps.append(multiprocessing.Process(target=p, kwargs=dict(software_config=C)))
    for p in pps:
        p.start()

    c = BusClient(KillMsg, NodeOnlineMsg)
    cnt = 0
    ready = False
    dd = dict()
    while True:
        if ready:
            msg = c.recv()
        else:
            c.send(Msg2())
            msg = c.poll(100)
        if msg is None:
            continue
        if isinstance(msg, KillMsg):
            if msg.name == '':
                break
        if not ready:
            if isinstance(msg, NodeOnlineMsg):
                if msg.name not in dd:
                    logging.info(msg)
                dd[msg.name] = 1
                if dd.keys().__len__() == len(ps):
                    logging.info(f'launcher stand by ready')
                    c.send(Msg1())
                    ready = True
    for p in pps:
        p.kill()
    bp.kill()


def launch_from_file(file: Path):
    assert file.exists()
    r = tomllib.loads(file.read_text())
    arg_d = dict()
    for k, v in r.items():
        arg_d[LaunchComponent[k]] = v
    launch(arg_d)


@click.group()
@click.option('--dev/--no-dev', default=True)
def entrypoint(dev):
    if dev:
        C.config_folder = P.DEV_PROJECT_FOLDER / 'config'
        C.log_folder = P.DEV_PROJECT_FOLDER / 'log'
    else:
        C.read_config(C.software_config_file)


@entrypoint.command()
@click.option('--data_folder', default=None)
@click.option('--generate_pyqt', default=True)
@click.option('-p', '--path', type=str,
              default=platformdirs.user_config_path('Flandre', 'Scarlet') / 'launch.toml',
              help='Path to launch.toml'
              )
def gui(data_folder, generate_pyqt, path):
    if (pyuic6 := shutil.which('pyuic6')) is None:
        print('pyuic6 is not installed')
        return
    if generate_pyqt:
        subprocess.run([pyuic6, '-o', P.PYQT / 'Main.py', P.PYQT / 'Main.ui'])
        subprocess.run([pyuic6, '-o', P.PYQT / 'Image.py', P.PYQT / 'Image.ui'])
    if data_folder is not None:
        C.record_folder = Path(data_folder)

    path = Path(path)
    if not path.exists():
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text('[MainUI]\n')
    print('Use launch config', path)
    launch_from_file(path)


mi1: MiSwitch = None
mi2: MiSwitch = None


@entrypoint.group()
def mi():
    global mi1, mi2
    mi1 = MiSwitch(C.switch1_ip, C.switch1_token)
    mi2 = MiSwitch(C.switch2_ip, C.switch2_token)


@mi.command()
@click.argument('name')
def on(name):
    if name == 'c1':
        mi1.on()
    if name == 'c2':
        mi2.on()


@mi.command()
@click.argument('name')
def off(name):
    if name == 'c1':
        mi1.off()
    if name == 'c2':
        mi2.off()


@mi.command()
@click.argument('name')
def status(name):
    if name == 'c1':
        print(mi1.is_on())
        print(mi1.power())
    if name == 'c2':
        print(mi2.is_on())


device_req: Socket = None

dd: Device = None


@entrypoint.group()
def device():
    global device_req
    ctx = zmq.Context()
    device_req = ctx.socket(zmq.REQ)
    device_req.connect(C.live_rep_socket)
    global dd
    dd = Device()
    dd.setup()


@device.command('connect')
def device_connect():
    dd.connect()


@device.command('disconnect')
def device_disconnect():
    dd.disconnect()


@device.command('enable')
def device_enable():
    dd.enable()


@device.command('disable')
def device_disable():
    dd.disable()


@device.command('upload')
@click.argument('name')
@click.argument('file', type=click.File('r'))
def device_upload(name, file: TextIOWrapper):
    dd.set_name_and_file_only(name, file.read())


@device.command('recvtest')
def device_recv_test():
    ctx = zmq.Context()
    pull = ctx.socket(zmq.PULL)
    pull.connect(C.live_push_socket)
    last_ts = 0
    first = True
    while True:
        b = pull.recv()

        ts, sequence_id, encoder, buffer = b2t(b)
        if first:
            last_ts = ts
            first = False
            continue
        print((ts - last_ts)/10**6,flush=True)
        last_ts = ts


if __name__ == '__main__':
    entrypoint()
