import cv2
import threading
import sys, os, time
import argparse
import numpy as np
import serial
import json
import glob
import shutil
from collections import deque
from datetime import datetime
from PyQt5.QtWidgets import (QApplication, QLabel, QWidget, QVBoxLayout,
                              QHBoxLayout, QMainWindow, QStackedWidget,
                              QFrame, QListWidget, QListWidgetItem,
                              QSizePolicy, QProgressBar, QGridLayout)
from PyQt5.QtGui import QImage, QPixmap, QFont
from PyQt5.QtCore import Qt, QTimer, pyqtSignal, QObject, QSize
from PIL import Image, ImageDraw, ImageFont

try:
    from surround_view import FisheyeCameraModel, BirdView
    import surround_view.param_settings as settings
except ImportError as e:
    print("[ERROR] Cannot import surround_view modules: {}".format(e))
    sys.exit(1)

# ==========================================================
# 遥控键码
# ==========================================================
KEY_UP    = 0x01
KEY_DOWN  = 0x09
KEY_LEFT  = 0x04
KEY_RIGHT = 0x06
KEY_BACK  = 0x02
KEY_OK    = 0x05

# ==========================================================
# 参数
# ==========================================================
VIDEO_DIR     = "/videos"
SEGMENT_SECS  = 180
RECORD_FPS    = 10
CELL_W, CELL_H = 640, 360
GRID_W, GRID_H = 1280, 720
DISK_WARN_PCT = 90
SEND_INTERVAL_MS = 100

MODE_360  = 0
MODE_QUAD = 1

DIR_CN     = {"front":"前方","left":"左方","back":"后方","right":"右方"}
WARN_COLOR = (255, 80,  80)
SAFE_COLOR = (100, 255, 100)

# ==========================================================
# 配置
# ==========================================================
CONFIG_PATH = os.path.join(os.getcwd(), "surround_config.json")
DEFAULT_CONFIG = {
    "carousel_interval": 4,
    "brightness":        50,
    "radar_unit":        0,
    "show_radar":        1,
    "radar_alarm_dist":  50,
    "display_mode":      MODE_360,
}

def load_config():
    if os.path.exists(CONFIG_PATH):
        try:
            with open(CONFIG_PATH, "r") as f:
                cfg = json.load(f)
            for k, v in DEFAULT_CONFIG.items():
                cfg.setdefault(k, v)
            return cfg
        except: pass
    return dict(DEFAULT_CONFIG)

def save_config(cfg):
    try:
        with open(CONFIG_PATH, "w") as f:
            json.dump(cfg, f, indent=2, ensure_ascii=False)
    except: pass

# ==========================================================
# 工具函数
# ==========================================================
def get_disk_info_str(path):
    try:
        usage = shutil.disk_usage(path)
        pct = usage.used * 100 // usage.total
        return "存储: {:.1f}G/{:.1f}G ({}%)".format(usage.used/1024**3, usage.total/1024**3, pct), pct
    except: return "磁盘未知", 0

class FontManager:
    _fonts = {}
    @classmethod
    def get(cls, size=20):
        if size not in cls._fonts:
            p = "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc"
            if not os.path.exists(p): p = None
            cls._fonts[size] = ImageFont.truetype(p, size) if p else ImageFont.load_default()
        return cls._fonts[size]

def add_texts_to_pil(pil_image, text_list):
    draw = ImageDraw.Draw(pil_image)
    for item in text_list:
        x, y, text, color, size = item
        font = FontManager.get(size)
        draw.text((x+1, y+1), text, font=font, fill=(0,0,0))
        draw.text((x, y), text, font=font, fill=color)

# ==========================================================
# 串口通信 (TTYS0)
# ==========================================================
class SerialWorker(threading.Thread):
    def __init__(self, port="/dev/ttyS0", baudrate=115200):
        super().__init__(daemon=True)
        self.port, self.baudrate = port, baudrate
        self.running = True
        self._rx_lock = threading.Lock()
        self._radar_dist, self._stopped, self._active_dirs = 0, False, []
        self._tx_lock = threading.Lock()
        self._tx_data = [0,0,0,0, 50] # F,L,B,R, Thr
        self._rx_enabled = True

    def set_rx_enabled(self, e): self._rx_enabled = e
    def update_send(self, f, l, b, r, thr):
        with self._tx_lock: self._tx_data = [f,l,b,r, thr]

    def run(self):
        try:
            ser = serial.Serial(self.port, self.baudrate, timeout=0.05)
        except: return

        def _send_loop():
            while self.running:
                with self._tx_lock:
                    d = self._tx_data
                    chk = sum(d[:4] + [0, d[4]]) & 0xFF
                    frame = bytes([0x55, d[0], d[1], d[2], d[3], 0x00, d[4], chk, 0xAA])
                ser.write(frame)
                time.sleep(SEND_INTERVAL_MS/1000.0)

        threading.Thread(target=_send_loop, daemon=True).start()
        buf = bytearray()
        while self.running:
            if not self._rx_enabled: 
                time.sleep(0.1); continue
            b = ser.read(1)
            if not b: continue
            buf.append(b[0])
            if len(buf) == 9:
                if buf[0] == 0x55 and buf[8] == 0xAA:
                    with self._rx_lock:
                        self._active_dirs = [n for i,n in enumerate(["front","left","back","right"]) if buf[i+1]]
                        self._radar_dist, self._stopped = buf[5], buf[6] > 0
                buf.clear()
            elif len(buf) > 9: buf.clear()

# ==========================================================
# 渲染线程 (消除黑边关键)
# ==========================================================
class RenderWorker(threading.Thread):
    def __init__(self, bird_worker, raw_worker, serial_worker, carousel, cfg, sw, sh):
        super().__init__(daemon=True)
        self.bird_w, self.raw_w, self.ser_w = bird_worker, raw_worker, serial_worker
        self.carousel, self.cfg = carousel, cfg
        self.sw, self.sh = sw, sh
        self.running = True
        self._out_q = deque(maxlen=2)

    def get_qimage(self):
        try: return self._out_q[-1]
        except: return None

    def run(self):
        while self.running:
            bird = self.bird_w.get_bird()
            if bird is None: time.sleep(0.02); continue
            
            active, dist, stopped = self.ser_w.get_state()
            cur_dir, rem = self.carousel.update(active)
            
            # 分配宽度
            bw = self.sw // 3
            fw = self.sw - bw
            
            # 强制拉伸(消除黑边)
            bird_part = cv2.resize(bird, (bw, self.sh))
            raw_f = self.raw_w.get_frame(cur_dir if cur_dir else "front")
            if raw_f is None: raw_f = np.zeros((self.sh, fw, 3), np.uint8)
            front_part = cv2.resize(raw_f, (fw, self.sh))

            # 亮度
            alpha = self.cfg.get("brightness", 50) / 50.0
            if alpha != 1.0: front_part = cv2.convertScaleAbs(front_part, alpha=alpha)

            canvas = np.hstack((bird_part, front_part))
            pil_img = Image.fromarray(cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB))
            
            # 文字叠加 (字号加大)
            texts = [(10, 10, "鸟瞰视图", (0,255,100), 25)]
            if cur_dir:
                texts.append((bw+20, 20, "{} 有人".format(DIR_CN[cur_dir]), WARN_COLOR, 40))
            
            # 雷达
            if self.cfg.get("show_radar"):
                txt = "雷达: {}cm".format(dist)
                texts.append((bw+20, self.sh-80, txt, (255,255,0), 30))

            add_texts_to_pil(pil_img, texts)
            arr = np.array(pil_img)
            self._out_q.append(QImage(arr.data, arr.shape[1], arr.shape[0], QImage.Format_RGB888))
            
            # 反向串口发送
            presence = [1 if d in active else 0 for d in ["front","left","back","right"]]
            self.ser_w.update_send(*presence, self.cfg.get("radar_alarm_dist", 50))
            time.sleep(0.03)

# ==========================================================
# 四分屏渲染
# ==========================================================
class QuadWorker(threading.Thread):
    def __init__(self, raw_worker, serial_worker, cfg, sw, sh):
        super().__init__(daemon=True)
        self.raw_w, self.ser_w, self.cfg = raw_worker, serial_worker, cfg
        self.sw, self.sh = sw, sh
        self.running = True
        self._out_q = deque(maxlen=2)

    def run(self):
        cw, ch = self.sw // 2, self.sh // 2
        while self.running:
            frames = []
            for n in ["front","left","back","right"]:
                f = self.raw_w.get_frame(n)
                frames.append(cv2.resize(f if f is not None else np.zeros((ch,cw,3), np.uint8), (cw,ch)))
            
            top = np.hstack((frames[0], frames[1]))
            bot = np.hstack((frames[2], frames[3]))
            canvas = np.vstack((top, bot))
            
            rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
            self._out_q.append(QImage(rgb.data, self.sw, self.sh, QImage.Format_RGB888))
            self.ser_w.update_send(0,0,0,0, self.cfg.get("radar_alarm_dist", 50))
            time.sleep(0.03)

# ==========================================================
# 回放面板 (充满右侧)
# ==========================================================
class PlaybackPanel(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self._video_list = []
        self._play_cap = None
        self._timer = QTimer()
        self._timer.timeout.connect(self._step)
        
        layout = QHBoxLayout(self)
        # 左侧列表放大
        self.list_w = QListWidget()
        self.list_w.setFixedWidth(300)
        self.list_w.setStyleSheet("font-size: 22px; background:#0d1117; color:white;")
        layout.addWidget(self.list_w)
        
        # 右侧画面
        self.view = QLabel("选择视频按 OK")
        self.view.setStyleSheet("background:black; border:2px solid #333;")
        self.view.setAlignment(Qt.AlignCenter)
        self.view.setScaledContents(True) # ★ 关键：拉伸充满
        layout.addWidget(self.view, 1)

    def load(self):
        self.list_w.clear()
        self._video_list = sorted(glob.glob(VIDEO_DIR + "/rec_*.mp4"), reverse=True)
        for f in self._video_list:
            item = QListWidgetItem(os.path.basename(f))
            item.setSizeHint(QSize(0, 60))
            self.list_w.addItem(item)
        self.list_w.setCurrentRow(0)

    def on_key(self, code):
        if code == KEY_UP: self.list_w.setCurrentRow(max(0, self.list_w.currentRow()-1))
        elif code == KEY_DOWN: self.list_w.setCurrentRow(min(len(self._video_list)-1, self.list_w.currentRow()+1))
        elif code == KEY_OK: self._start_play()

    def _start_play(self):
        if self._play_cap: self._play_cap.release()
        self._play_cap = cv2.VideoCapture(self._video_list[self.list_w.currentRow()])
        self._timer.start(33)

    def _step(self):
        ret, frame = self._play_cap.read()
        if ret:
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            qimg = QImage(rgb.data, rgb.shape[1], rgb.shape[0], QImage.Format_RGB888)
            self.view.setPixmap(QPixmap.fromImage(qimg))
        else: self._timer.stop()

    def cleanup(self):
        self._timer.stop()
        if self._play_cap: self._play_cap.release()

# ==========================================================
# 设置面板 (字体放大)
# ==========================================================
class SettingsPanel(QWidget):
    def __init__(self, cfg, parent=None):
        super().__init__(parent)
        self.cfg = cfg
        self.cur = 0
        self.sub_mode = False # False=Menu, True=Playback
        
        self.menu = [
            {"n":"显示模式", "k":"display_mode", "opts":["360环视","四分屏"]},
            {"n":"轮播间隔", "k":"carousel_interval", "min":1, "max":10},
            {"n":"画面亮度", "k":"brightness", "min":0, "max":100},
            {"n":"雷达阈值", "k":"radar_alarm_dist", "min":10, "max":200},
            {"n":"视频回放", "k":"page"}
        ]
        
        self.stack = QStackedWidget()
        # Menu Page
        self.menu_page = QWidget()
        mv = QVBoxLayout(self.menu_page)
        self.list_v = QListWidget()
        # ★ 列表样式：大字体，大行高
        self.list_v.setStyleSheet("""
            QListWidget { font-size: 26px; background:#1a1a2e; color:white; border:none; }
            QListWidget::item { height: 70px; border-bottom:1px solid #333; padding-left:20px; }
            QListWidget::item:selected { background:#e94560; }
        """)
        mv.addWidget(self.list_v)
        self.val_l = QLabel()
        self.val_l.setAlignment(Qt.AlignCenter)
        self.val_l.setStyleSheet("font-size: 50px; color:#e94560; padding:20px;")
        mv.addWidget(self.val_l)
        
        self.stack.addWidget(self.menu_page)
        self.play_panel = PlaybackPanel()
        self.stack.addWidget(self.play_panel)
        
        layout = QVBoxLayout(self)
        layout.setContentsMargins(0,0,0,0)
        layout.addWidget(self.stack)

    def load(self):
        self.sub_mode = False
        self.stack.setCurrentIndex(0)
        self._refresh()

    def _refresh(self):
        self.list_v.clear()
        for i, m in enumerate(self.menu):
            self.list_v.addItem(m["n"])
        self.list_v.setCurrentRow(self.cur)
        # 显示当前值
        item = self.menu[self.cur]
        if "opts" in item:
            self.val_l.setText(item["opts"][self.cfg[item["k"]]])
        elif "min" in item:
            self.val_l.setText(str(self.cfg[item["k"]]))
        else:
            self.val_l.setText("按 OK 进入")

    def on_key(self, code):
        if self.sub_mode:
            if code == KEY_BACK:
                self.sub_mode = False
                self.play_panel.cleanup()
                self.stack.setCurrentIndex(0)
            else: self.play_panel.on_key(code)
            return

        if code == KEY_UP: self.cur = (self.cur-1)%len(self.menu); self._refresh()
        elif code == KEY_DOWN: self.cur = (self.cur+1)%len(self.menu); self._refresh()
        elif code == KEY_LEFT or code == KEY_RIGHT:
            item = self.menu[self.cur]
            step = 1 if code == KEY_RIGHT else -1
            if "opts" in item:
                self.cfg[item["k"]] = (self.cfg[item["k"]] + step) % len(item["opts"])
            elif "min" in item:
                self.cfg[item["k"]] = max(item["min"], min(item["max"], self.cfg[item["k"]] + step * 5))
            self._refresh()
        elif code == KEY_OK:
            if self.menu[self.cur]["k"] == "page":
                self.sub_mode = True
                self.play_panel.load()
                self.stack.setCurrentIndex(1)
            elif "opts" in self.menu[self.cur]:
                self.cfg[self.menu[self.cur]["k"]] = (self.cfg[self.menu[self.cur]["k"]] + 1) % len(self.menu[self.cur]["opts"])
                self._refresh()

# ==========================================================
# 主窗口
# ==========================================================
class MainWindow(QMainWindow):
    def __init__(self, multi_cam, ser_w, remote_w, cfg):
        super().__init__()
        self.cfg = cfg
        self.ser_w = ser_w
        self.showFullScreen()
        sw, sh = self.width(), self.height()
        
        self.stack = QStackedWidget()
        self.setCentralWidget(self.stack)
        
        # View Page
        self.view_l = QLabel()
        self.view_l.setScaledContents(True) # ★ 强制铺满
        self.stack.addWidget(self.view_l)
        
        self.set_panel = SettingsPanel(cfg)
        self.stack.addWidget(self.set_panel)
        
        # Workers
        self.raw_w = RawCaptureWorker(multi_cam.caps, multi_cam.names)
        self.raw_w.start()
        self.bird_w = BirdViewWorker(self.raw_w, multi_cam.camera_models, multi_cam.birdview)
        self.bird_w.start()
        
        self.render_360 = RenderWorker(self.bird_w, self.raw_w, ser_w, CarouselScheduler(), cfg, sw, sh)
        self.render_360.start()
        self.render_quad = QuadWorker(self.raw_w, ser_w, cfg, sw, sh)
        self.render_quad.start()
        
        RecordWorker(self.raw_w).start()

        self.timer = QTimer()
        self.timer.timeout.connect(self._flush)
        self.timer.start(20)
        
        remote_w.signals.key_pressed.connect(self._on_key)

    def _flush(self):
        if self.stack.currentIndex() != 0: return
        qimg = self.render_quad._out_q[-1] if self.cfg["display_mode"] == 1 else self.render_360.get_qimage()
        if qimg: self.view_l.setPixmap(QPixmap.fromImage(qimg))

    def _on_key(self, code):
        if self.stack.currentIndex() == 0:
            if code == KEY_OK:
                self.set_panel.load()
                self.stack.setCurrentIndex(1)
        else:
            if code == KEY_BACK and not self.set_panel.sub_mode:
                save_config(self.cfg)
                self.ser_w.set_rx_enabled(self.cfg["display_mode"] == 0)
                self.stack.setCurrentIndex(0)
            else:
                self.set_panel.on_key(code)

# ==========================================================
# 入口
# ==========================================================
def main():
    cfg = load_config()
    app = QApplication(sys.argv)
    
    ser_w = SerialWorker()
    ser_w.start()
    rem_w = RemoteWorker()
    rem_w.start()
    
    cam = MultiCameraBirdView()
    if not cam.running: return
    
    win = MainWindow(cam, ser_w, rem_w, cfg)
    win.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()
