import gradio as gr
import asyncio
import json
import simplepyble
import datetime
import time
import struct
from functools import partial
from apscheduler.schedulers.background import BackgroundScheduler
import hashlib
from pylsl import StreamInfo, StreamOutlet, cf_double64
import os, sys
import numpy as np
import logging

yams_dir = "yams-data"

def session_manager_interface():
    sub_list = gr.Text("sub-Test")
    ses_list = gr.Text("ses-Demo")

class MsenseOutlet(StreamOutlet):
    def __init__(self, name, peripheral, chunk_size=32, max_buffered=360):
        self.name = name.replace(':', '-')
        info = StreamInfo(name, "MotionSenSE", 3, 2, cf_double64, peripheral.address())
        super().__init__(info, chunk_size, max_buffered)

        self.log_dir = os.path.join(yams_dir, "default")
        self.msg = f"📻 {self.tic()} LSL OK. Ready to start..."
        self.msg_fun = f"📻 {self.tic()} LSL OK. Ready to start..."

    def tic(self):
        now = datetime.datetime.now()
        return now.strftime("%H:%M:%S")

    def save_data(self, data):
        self.log_path = os.path.join(self.log_dir, f"{self.name}.txt")
        # Ensure the file exists
        if not os.path.exists(self.log_path):
            with open(self.log_path, 'w') as f: pass

        # Append NumPy array as a line
        with open(self.log_path, 'a') as f:
            np.savetxt(f, [data], fmt='%s')

    def push_sample(self, x):
        formatted = '\t'.join(str(num) for num in x)
        self.msg = f"📻 {self.tic()} last LSL pushed: {formatted}"
        
        fun_msg = "".join(["✅" for i in range(int(time.time())%10)])
        self.msg_fun = f"📻 {self.tic()} {fun_msg}"

        x.append(time.time())
        super().push_sample(x)
        self.save_data(x)

class MsenseController():
    def __init__(self):
        # current YYMMDD
        now = datetime.datetime.now()
        date = now.strftime("%Y-%m-%d")
    
        # init logger
        self.logger = logging.getLogger(__name__)
        os.makedirs(yams_dir, exist_ok=True)
        logging.basicConfig(level=logging.INFO, 
                            format='%(asctime)s [%(levelname)s] %(message)s',
                            handlers=[
                                logging.FileHandler(os.path.join(yams_dir, f"{date}_yams_session.log")),
                                logging.StreamHandler()
                            ])
        self.logger.info(f"Start YAMS session log")
    
        self.auto_reconnect = True
        self.devices = {}
        self.device_name = self.get_dev_dict()
        self.init_adapter()
        self.active_devices = {}
        self.active_outlets = {}

        self.scheduler = BackgroundScheduler()

    def get_dev_dict(self):
        try:
            with open("device_info.json", 'r') as file:
                device_name = json.load(file)
                device_name = {value: key for key, value in device_name.items()}
        except:
            device_name = {}

        print(device_name)
        return device_name

    def init_adapter(self):
        adapters = simplepyble.Adapter.get_adapters()
        assert len(adapters) > 0, "No BT adapter found"
        
        self.adapter = adapters[0]
        # print(f"Selected adapter: {self.adapter.identifier()} [{self.adapter.address()}]")
        self.logger.info(f"Selected adapter: {self.adapter.identifier()} [{self.adapter.address()}]")

    def get_available_devices_checkbox(self, filter_name="MSense"):
        self.scan_devices(filter_name=filter_name)
        return gr.CheckboxGroup(choices=list(self.devices.keys()), 
                                value=list(self.devices.keys()),
                                label="Available devices")

    def scan_devices(self, filter_name="MSense"):
        print("start scanning devices")
        self.logger.info("start device scanning")
        self.adapter.scan_for(5000)
        peripherals = self.adapter.scan_get_results()

        self.devices = {}
        for i, peripheral in enumerate(peripherals):
            if filter_name in peripheral.identifier():
                self.logger.info(f"{i}: {peripheral.identifier()} [{peripheral.address()}]")
                # try to look up device alias
                if peripheral.address() in self.device_name.keys():
                    alias = self.device_name[peripheral.address()]
                    name = f"{peripheral.identifier()} ({alias}) [{peripheral.address()}]"
                else:
                    name = f"{peripheral.identifier()} [{peripheral.address()}]"

                self.devices[name] = peripheral


    def connect_devices(self, names):
        del(self.active_devices)
        self.active_devices = {}
        del(self.active_outlets)
        self.active_outlets = {}

        self.logger.info(f"Start connecting to devices: {names}")
        for n in names:
            gr.Info(f"Connecting to devices: {n}")
            print(f'==== {n}')
            p = self.devices[n]
            print(f"=== {p.identifier()} at {p.address()}")
            p.set_callback_on_connected(lambda: self.logger.info(f"{n} {p.identifier()} is connected"))
            p.set_callback_on_disconnected(lambda: self.logger.info(f"{n} {p.identifier()} is disconnected"))
            p.connect()
            self.active_devices[n] = p
            self.active_outlets[n] = MsenseOutlet(n, p)

    def disconnect_all(self):
        for dev in self.active_devices.values():
            try:
                dev.disconnect()
            except Exception as e:
                print(str(e))
        gr.Warning(f"All devices disconnected")
        self.logger.info("All devices disconnected")

    def tic(self):
        return datetime.datetime.now()
    
    def log(self, msg):
        print(f"{self.tic()}: {msg}")

    def interface(self):
        with gr.Accordion("Initialization"):
            bt_search = gr.Button("Search Bluetooth devices 📱")
            available_devices = gr.CheckboxGroup(label="Available devices", scale=6)
            with gr.Row():
                bt_connect = gr.Button("✅ Connect selected")
                btn_disconnect = gr.Button("❌ Disconnect")

        
        bt_connect.click(self.connect_devices, inputs=available_devices)            
        btn_disconnect.click(self.disconnect_all)

        with gr.Accordion(label="Device control", open=True):
            default_sub = "sub-Test"
            default_ses = "ses-00"

            with gr.Row():
                sub_name = gr.Text(default_sub, label="Subject ID")
                ses_name = gr.Text(default_ses, label="Session ID")
                subject_enc = gr.Number(self.get_participant_encoding(default_sub, default_ses), label='Participant encoding (Read-only)', interactive=False)
                sub_name.change(self.get_participant_encoding, inputs=[sub_name, ses_name], outputs=subject_enc)
                ses_name.change(self.get_participant_encoding, inputs=[sub_name, ses_name], outputs=subject_enc)
        
            with gr.Row():
                btn_start = gr.Button("Start▶️")
                btn_stop = gr.Button("Stop🛑")
                
            btn_start.click(self.start_collection)
            btn_stop.click(self.end_collection)


        self.params = {}
        params = gr.ParamViewer(self.params)
        timer = gr.Timer(value=3)
        timer.tick(fn=self.update_params, outputs=params)

        # erase control
        with gr.Accordion(label="🚨🚨🚨Danger zone🚨🚨🚨", open=False):
            erase_passcode = gr.Number(label="Erase code")
            erase_enable = gr.Checkbox(label="Enable erase feature")
            
            erase_btn = gr.Button("Erase flash data", interactive=False)
            erase_btn.click(self.erase_flash_data, 
                            outputs=[erase_passcode, erase_enable, erase_btn])

            erase_enable.change(self.set_erase_feature, inputs=[erase_enable, erase_passcode], outputs=[erase_btn])

        with gr.Accordion(label="Advanced options", open=False):
            text = gr.Text("MSense", label="Device filter", scale=2)

            auto_reconnect = gr.Checkbox(True, label="Auto reconnect")
            auto_reconnect.change(self.set_auto_reconnect, inputs=auto_reconnect)

            btn_reconnect = gr.Button("Reconnect")
            btn_reconnect.click(self.reconnect)

            btn_service = gr.Button("Get available services")
            btn_service.click(self.get_selected_device_services, inputs=available_devices)
            with gr.Row():
                btn_monitor_start = gr.Button("Start device monitor")
                btn_monitor_stop = gr.Button("Stop device monitor")

                btn_monitor_start.click(self.start_device_monitor)
                btn_monitor_stop.click(self.stop_device_monitor)

        bt_search.click(self.get_available_devices_checkbox, inputs=text, outputs=available_devices)    

    def update_params(self):
        self.params = {}
        for name, device in self.active_devices.items():
            connection_status = "✅ Connected" if device.is_connected() else "🚫 Disconnected"
            self.params[name] = {
                'type': f"{connection_status} | {self.active_outlets[name].msg_fun}",
                'description': self.active_outlets[name].msg
            }

        return self.params

    def set_auto_reconnect(self, status):
        self.auto_reconnect = status

    def start_collection(self):
        timestamp = time.strftime("%y%m%d_%H%M")
        # create log dir
        self.log_dir = os.path.join(yams_dir, 
                                    self.session_info['sub_id'], 
                                    self.session_info['ses_id'], 
                                    f"{self.session_info['participant_enc']}_{timestamp}")
        print(f"create log dir {self.log_dir}")
        os.makedirs(self.log_dir, exist_ok=True)

        gr.Info("▶️ Start data collection...")
        self.logger.info(f"Start data collection with out dir = {self.log_dir}")
        self.logger.info(f"Subject ID = {self.session_info['sub_id']}")
        self.logger.info(f"Session ID = {self.session_info['ses_id']}")
        self.logger.info(f"Participant encoding = {self.session_info['participant_enc']}")

        for name, p in self.active_devices.items():
            print(name, p.is_connected(), p.is_connectable())
            self.collection_ctl(name, True)
            self.active_outlets[name].log_dir = self.log_dir

    def end_collection(self):
        gr.Info("🛑 Stop data collection...")
        self.logger.info("Data collection stopped")
        for name, p in self.active_devices.items():
            print(name, p.is_connected(), p.is_connectable())
            self.collection_ctl(name, False)
    
    def collection_ctl(self, name, start=True):
        peripheral = self.active_devices[name]

        # if starting, do the initialization
        if start:
            # write unix time
            peripheral.write_request("da39c930-1d81-48e2-9c68-d0ae4bbd351f", 
                                     "da39c932-1d81-48e2-9c68-d0ae4bbd351f", 
                                     struct.pack("<Q", int(time.time())))
            # write participant hash
            peripheral.write_request("da39c930-1d81-48e2-9c68-d0ae4bbd351f",
                                     "da39c933-1d81-48e2-9c68-d0ae4bbd351f", 
                                     self.participant_byte)

        service_uuid = "da39c930-1d81-48e2-9c68-d0ae4bbd351f"
        characteristic_uuid = "da39c931-1d81-48e2-9c68-d0ae4bbd351f"
        peripheral.write_request(service_uuid, characteristic_uuid, struct.pack("<I", int(start)))

        self.register_enmo(peripheral, name)

        # 
        if start and self.auto_reconnect:
            self.start_device_monitor()
        elif not start:
            self.stop_device_monitor()
            

    def register_enmo(self, peripheral, name):
        # ENMO 
        service_uuid = "da39c950-1d81-48e2-9c68-d0ae4bbd351f"
        characteristic_uuid = "da39c951-1d81-48e2-9c68-d0ae4bbd351f"
        contents = peripheral.notify(service_uuid, characteristic_uuid, lambda data: self.enmo_handler(data, peripheral, name))

    def enmo_handler(self, data, peripheral, name):
        # print(peripheral.identifier(), data)
        packet_counter = data[4:6]
        ENMO = struct.unpack("<f", data[0:4])
        
        packet_counter = struct.unpack("<H", packet_counter)
        horizontal_array = [ENMO[0], packet_counter[0]]
        print(f"{name}: package counter", horizontal_array)

        self.active_outlets[name].push_sample([ENMO[0], packet_counter[0]])
        

    def get_selected_device_services(self, names):
        for n in names:
            p = self.devices[n]
            print(f'======== Services of device {n}')
            self.get_services(p)

    def get_services(self, peripheral):
        services = peripheral.services()
        service_characteristic_pair = []
        for service in services:
            for characteristic in service.characteristics():
                service_characteristic_pair.append((service.uuid(), characteristic.uuid()))

        for i, (service_uuid, characteristic) in enumerate(service_characteristic_pair):
            print(f"{i}: {service_uuid} {characteristic}")

    def reconnect(self):
        for name, p in self.active_devices.items():
            print(f"{name} {p.identifier()} connection status = {p.is_connected()}")
            if not p.is_connected():
                try:
                    p.connect()
                    self.register_enmo(p, name)
                except Exception as e:
                    print(str(e))

    ######
    def check_and_reconnect_devices(self):
        for name, device in self.active_devices.items():
            try:
                if not device.is_connected():
                    self.logger.warning(f"{device.identifier()} disconnected. Attempting to reconnect..")
                    device.connect()
                    if device.is_connected():
                        self.logger.info(f"Reconnected to {device.identifier()}")
                        self.register_enmo(device, name)
                    else:
                        self.logger.warning(f"Failed to reconnect to {device.identifier()}")
                else:
                    self.logger.debug(f"{device.identifier()} is still connected.")
            except Exception as e:
                self.logger.error(f"Error checking device {device.identifier()}: {e}")

    def start_device_monitor(self, interval_seconds=10):
        if not self.scheduler.running:
            self.scheduler.start()
        if not self.scheduler.get_job("device_monitor"):
            self.scheduler.add_job(self.check_and_reconnect_devices, "interval", seconds=interval_seconds, id="device_monitor")
            self.logger.info("Started device monitor job")

    def stop_device_monitor(self):
        job = self.scheduler.get_job("device_monitor")
        if job:
            job.remove()
            self.logger.info("Stopped device monitor job")

    def get_participant_encoding(self, sub, ses):
        name = f"{sub}-{ses}"
        hash_object = hashlib.sha256(name.encode())
        hex_digest = hash_object.hexdigest()
        integer_representation = int(hex_digest, 16) % 32000
        # print(name, integer_representation)
        self.participant_byte = struct.pack("<I", integer_representation)

        self.session_info = {
            'sub_id': sub,
            'ses_id': ses,
            'participant_enc': integer_representation
        }
        return integer_representation


    def set_erase_feature(self, erase_enable, erase_passcode):
        if erase_enable:
            if erase_passcode == 68:
                gr.Warning("Erase feature is enabled!")
                return gr.Button("Erase flash data", interactive=True)
            else:
                gr.Warning("Incorrect password")
                return gr.Button("Erase flash data", interactive=False)
        else:
            gr.Info("Erase feature disabled")
            return gr.Button("Erase flash data", interactive=False)
        
    def erase_flash_data(self):
        gr.Info("Erasing flash data...")
        self.logger.info("Erasing flash data...")

        service_uuid = "da39c930-1d81-48e2-9c68-d0ae4bbd351f"
        rst_char = "da39c934-1d81-48e2-9c68-d0ae4bbd351f"

        for name, p in self.active_devices.items():
            try:
                p.write_request(service_uuid, rst_char, struct.pack("<I", int(68)))
            except Exception as e:
                self.logger.error(str(e))
                gr.Error(f"⚠️{str(e)}")
        self.logger.info("Erase completed")
        return gr.Number(value=None, label="Erase code"), gr.Checkbox(label="Enable erase feature", value=False), gr.Button("Erase flash data", interactive=False)
