import threading
import time
import serial
import numpy as np
import keyboard  # 需要 pip install keyboard
from enum import Enum

# === 官方驱动引用 (请确保这3个文件在同级目录) ===
from class_ch341 import *
from class_sensorcmd import *
from class_finger import *

# ==========================================
# PART 1: 传感器驱动 (完整版 TactileSensorDAQ)
# ==========================================
DEF_MAX_FINGER_NUM = 3
PCA_ADDR = 0x70
SAMPLE_RATE_MS = 10


class EnumCh341ConnectStatus(Enum):
    CH341_CONNECT_INIT = 0
    CH341_CONNECT_OPEN = 1
    CH341_CONNECT_SET_SPEED = 2
    CH341_CONNECT_SAMPLE_START = 3
    CH341_CONNECT_CHECK = 4


class TactileSensorDAQ:
    def __init__(self):
        # 1. 硬件初始化
        self.ch341 = ClassCh341()
        self.fingers = list()
        # 初始化3个传感器对象
        for i in range(DEF_MAX_FINGER_NUM):
            self.fingers.append(ClassFinger(4 + i, self.ch341))

        # 2. 状态机变量
        self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_INIT
        self.ch341CheckTimer = 0
        self.pcaAddr = PCA_ADDR
        self.syncTimer = 0

        # 3. 数据容器
        self.raw_data = np.zeros(12, dtype=np.float32)
        self.offset = np.zeros(12, dtype=np.float32)
        self.clean_data = np.zeros(12, dtype=np.float32)

        # 4. 线程控制
        self.running = False
        self.lock = threading.Lock()
        self.thread = None

    def _set_sensor_enable(self, idx):
        _pack = list()
        _pack.append(idx)
        self.ch341.write(self.pcaAddr, _pack)

    def _update_state_machine(self):
        if self.connectStatus == EnumCh341ConnectStatus.CH341_CONNECT_INIT:
            if self.ch341.init(): self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_OPEN
        elif self.connectStatus == EnumCh341ConnectStatus.CH341_CONNECT_OPEN:
            if self.ch341.open():
                self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_SET_SPEED
            else:
                self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_INIT
        elif self.connectStatus == EnumCh341ConnectStatus.CH341_CONNECT_SET_SPEED:
            if self.ch341.set_speed(self.ch341.IIC_SPEED_400):
                self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_SAMPLE_START
            else:
                self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_SAMPLE_START
        elif self.connectStatus == EnumCh341ConnectStatus.CH341_CONNECT_SAMPLE_START:
            self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_CHECK

    def _read_hardware(self):
        connectedSensorChan = 0
        temp_data_buffer = []

        for fingerIndex in range(len(self.fingers)):
            self._set_sensor_enable(1 << (self.fingers[fingerIndex].pcaIdx))
            connectedSensorChan |= 1 << (self.fingers[fingerIndex].pcaIdx)
            current_finger = self.fingers[fingerIndex]

            if not current_finger.connect:
                if current_finger.checkSensor():
                    print(f"[Sensor] Finger {fingerIndex} Connected!")
            else:
                current_finger.capRead()
                for unit_i in range(current_finger.projectPara.ydds_num):
                    fn = current_finger.readData.nf[unit_i]
                    ft = current_finger.readData.tf[unit_i]
                    temp_data_buffer.append(fn)
                    temp_data_buffer.append(ft)

        if len(temp_data_buffer) == 12:
            with self.lock:
                self.raw_data = np.array(temp_data_buffer, dtype=np.float32)
                self.clean_data = self.raw_data - self.offset

        if (time.time() - self.syncTimer) > 1.0:
            self.syncTimer = time.time()
            self._set_sensor_enable(connectedSensorChan)
            for f in self.fingers:
                if f.connect:
                    f.snsCmd.setSensorSync(0)
                    break

    def _thread_worker(self):
        while self.running:
            if self.connectStatus != EnumCh341ConnectStatus.CH341_CONNECT_CHECK:
                self._update_state_machine()
                time.sleep(0.1)
                continue

            start_time = time.time()
            try:
                self._read_hardware()
            except Exception as e:
                print(f"Read Error: {e}")

            self.ch341CheckTimer += (time.time() - start_time) * 1000
            if self.ch341CheckTimer >= 1000:
                self.ch341CheckTimer = 0
                if not self.ch341.connectCheck():
                    print("[Sensor] CH341 Disconnected!")
                    self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_INIT

            elapsed = (time.time() - start_time) * 1000
            sleep_time = (SAMPLE_RATE_MS - elapsed) / 1000.0
            if sleep_time > 0: time.sleep(sleep_time)

    def start(self):
        if self.running: return
        self.running = True
        self.thread = threading.Thread(target=self._thread_worker)
        self.thread.daemon = True
        self.thread.start()
        print("[System] Sensor Thread Started.")

    def stop(self):
        self.running = False
        if self.thread: self.thread.join()
        self.ch341.disconnect()
        print("[System] Sensor Thread Stopped.")

    def tare(self):
        print("[System] Taring sensors... please wait.")
        time.sleep(1)
        with self.lock:
            self.offset = np.copy(self.raw_data)
        print("[System] Tare complete.")

    def get_max_force(self):
        with self.lock:
            data = np.copy(self.clean_data)
        max_force = 0.0
        for i in range(0, 12, 2):
            if data[i] > max_force: max_force = data[i]
        return max_force


# ==========================================
# PART 2: 执行器驱动 (RobotDriver)
# ==========================================
class RobotDriver:
    def __init__(self, port='COM9', baud=115200):
        try:
            self.ser = serial.Serial(port, baud, timeout=1)
            print(f"[System] Robot Connected to {port}")
            time.sleep(2)
        except Exception as e:
            print(f"[Error] Robot Connection failed: {e}")
            self.ser = None

    def _send(self, cmd):
        if self.ser:
            # 加上 \r\n 以防万一
            full_cmd = f"{cmd}\r\n"
            self.ser.write(full_cmd.encode())
            time.sleep(0.05)

    def motor_open(self):
        self._send("M:OPEN")

    def motor_close(self):
        self._send("M:CLOSE")

    def motor_stop(self):
        self._send("M:STOP")

    def set_servo(self, servo_id, angle):
        self._send(f"S{servo_id}:{angle}")

    def set_config(self, mode):
        print(f"[Robot] Switching to Config Mode {mode}...")
        if mode == 0:
            self.set_servo(1, 90)
            self.set_servo(2, 90)
        elif mode == 1:
            self.set_servo(1, 30)
            self.set_servo(2, 150)
        elif mode == 2:
            self.set_servo(1, 120)
            self.set_servo(2, 60)
        time.sleep(1)

    def close(self):
        self.motor_stop()
        if self.ser: self.ser.close()


# ==========================================
# PART 3: 全局急停逻辑 (SafetyGuard)
# ==========================================
class SafetyGuard:
    def __init__(self, robot):
        self.robot = robot
        self.is_paused = False
        keyboard.add_hotkey('space', self.toggle_safety)

    def toggle_safety(self):
        self.is_paused = not self.is_paused
        if self.is_paused:
            self.robot.motor_stop()
            print("\n\n" + "!" * 40)
            print("!!! 紧急停止触发 (EMERGENCY STOP) !!!")
            print("!!! 电机已锁死，程序挂起         !!!")
            print("!!! 再按一次 [空格] 恢复运行      !!!")
            print("!" * 40 + "\n")
        else:
            print("\n" + "=" * 40)
            print(">>> 解除急停，继续任务 (Resuming)...")
            print("=" * 40 + "\n")

    def check_pause(self):
        if not self.is_paused: return None
        while self.is_paused: time.sleep(0.1)
        time.sleep(0.5)
        return True


# ==========================================
# PART 4: 主逻辑 (Main Controller)
# ==========================================

def run_auto_grasp_task(robot, sensor, safety):
    print("\n>>> 任务开始 <<<")

    # 1. 构型输入
    try:
        mode_str = input("请输入目标构型 (0:初始, 1:错位, 2:对握): ").strip()
        mode = int(mode_str)
        if mode not in [0, 1, 2]: raise ValueError
    except:
        print("[Error] 无效输入")
        return

    # 2. 变构型
    safety.check_pause()
    robot.set_config(mode)

    # === 限制任务开始时的张开等待，避免过度张开 ===
    OPEN_WAIT_TIME = 1.0  # 秒

    safety.check_pause()
    print(f"[Task] 初始化：直线电机张开... (等待 {OPEN_WAIT_TIME} 秒)")
    robot.motor_open()

    # 3. 延时 (带倒计时的充分等待)
    steps = int(OPEN_WAIT_TIME * 10)  # 转换为 0.1s 的步数
    print(f"[Task] 正在张开并等待放置番茄 (按空格可急停)...")

    for i in range(steps):
        if safety.check_pause():
            print("[Task] 暂停恢复，重置倒计时...")

        time.sleep(0.1)

        # 每秒打印一次倒计时
        seconds_left = OPEN_WAIT_TIME - (i / 10)
        if i % 10 == 0:
            print(f"{int(seconds_left)}...", end=' ', flush=True)

    print("Go!")
    # ===============================================

    # 4. 开始闭合
    safety.check_pause()
    print("[Task] 直线电机开始闭合...")
    robot.motor_close()

    # 5. 力反馈循环
    # ================= 修改位置：力控阈值 =================
    FORCE_THRESHOLD = 0.3
    # ===================================================
    start_time = time.time()

    try:
        while True:
            # 急停恢复检查
            recover_flag = safety.check_pause()
            if recover_flag:
                print("[Task] 恢复运动：重新下发闭合指令...")
                robot.motor_close()
                start_time = time.time()

            # 获取传感器数据
            current_max_force = sensor.get_max_force()

            # 打印状态
            if (time.time() * 1000) % 200 < 20:
                print(f"\r[Grasping] Force: {current_max_force:.1f} / {FORCE_THRESHOLD}", end="")

            # 触发判断
            if current_max_force > FORCE_THRESHOLD:
                print(f"\n[Task] 触觉触发！停止。")
                break

            # 超时保护 (15秒，防止闭合行程也很长)
            if time.time() - start_time > 15.0:
                print("\n[Task] 抓取超时 (未检测到受力)。")
                break

            time.sleep(0.01)

    except KeyboardInterrupt:
        print("\n[Task] 人工强制中断！")

    # 6. 停止
    robot.motor_stop()
    print("[Task] 任务结束。\n")


if __name__ == "__main__":
    sensor = TactileSensorDAQ()
    sensor.start()

    # ================= 修改位置：串口号 =================
    # 请务必确认这里是 STM32 的串口号，而不是传感器的
    robot = RobotDriver(port='COM9')
    # ==================================================


    safety = SafetyGuard(robot)


    try:
        print("等待传感器稳定...")

        time.sleep(2)
        sensor.tare()

        while True:
            print("========================")
            print(" [Enter] 运行抓取任务")
            print(" [Space] 随时急停/恢复")
            print(" [t]     重新去皮")
            print(" [q]     退出")

            cmd = input("Command > ").strip().lower()

            if cmd == '':
                run_auto_grasp_task(robot, sensor, safety)

            elif cmd == 't':
                sensor.tare()
            elif cmd == 'q':
                break

    except KeyboardInterrupt:
        pass
    finally:
        robot.close()
        sensor.stop()
        try:
            keyboard.unhook_all()
        except:
            pass
        print("System All Shutdown.")