#!/usr/bin/env python3
"""
UAV Speaker Server v1.2.0
Ubuntu 服务器：接收 Android 上传的 WAV 音频，通过串口转发给 ESP32-S3
Features:
  - Auto-detect ESP32 serial port by heartbeat (UAVSPK_HB)
  - Auto-reconnect on disconnect / hot-plug
  - Resample any WAV to 44.1kHz mono 16bit via ffmpeg
  - systemd service compatible
"""

import os
import sys
import struct
import math
import time
import glob
import serial
import serial.tools.list_ports
import threading
import subprocess
from flask import Flask, request, jsonify

VERSION = "1.4.1"
app = Flask(__name__)
UPLOAD_FOLDER = "/tmp/audio_uploads"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# ============== 串口配置 ==============
SERIAL_BAUD = 1500000
SERIAL_PORT = None          # Auto-detected at runtime
HEARTBEAT_INTERVAL = 5      # Seconds between auto-detect retries
HEARTBEAT_PREFIX = b"UAVSPK_HB_"  # e.g. UAVSPK_HB_E072A1F5FB6C
# 兼容旧版固件心跳前缀（过渡使用，建议尽快升级到 UAVSPK_HB_）
HEARTBEAT_PREFIX_LEGACY = b"ESP32_HB_"
# =====================================

# ============== 音频配置 ==============
# 可通过环境变量 UAVSPK_MAX_AUDIO_SECONDS 配置，默认 30 秒
MAX_AUDIO_SECONDS = int(os.environ.get("UAVSPK_MAX_AUDIO_SECONDS", "30"))
# =====================================

# 44.1kHz Mono 16bit = 88200 bytes/s
# Send every 20ms: 44100 * 0.02 * 2 = 1764 bytes
SEND_INTERVAL_MS = 20
SEND_CHUNK_BYTES = int(44100 * SEND_INTERVAL_MS / 1000 * 2)

ser = None
serial_lock = threading.Lock()
serial_detect_thread = None

# ============== 串口心跳超时检测 ==============
# 设备每 3 秒发送一次心跳，超过 10 秒未收到视为断开
last_heartbeat_time = 0.0
HEARTBEAT_TIMEOUT_SEC = 10.0
# =============================================

# ============== 播放控制 ==============
play_stop_event = threading.Event()
current_play_thread = None
playback_status = {
    "playing": False,
    "file": None,
    "started_at": None,
    "stopped_at": None,
    "sent_bytes": 0,
    "total_bytes": 0,
}
status_lock = threading.Lock()
# =====================================


def list_candidate_ports():
    """列出可能的串口设备（排除系统串口和已知的非对讲机设备）"""
    candidates = []
    for p in serial.tools.list_ports.comports():
        dev = p.device
        # 优先匹配常见的 USB-UART 芯片 VID/PID
        vid = getattr(p, 'vid', None)
        pid = getattr(p, 'pid', None)
        # CH343 (WCH): VID=0x1a86, PID=0x55d3/0x55d4
        # CP2102 (Silicon Labs): VID=0x10c4, PID=0xea60
        # FT232 (FTDI): VID=0x0403, PID=0x6001
        is_likely_esp32 = False
        if vid == 0x1a86 and pid in (0x55d3, 0x55d4):
            is_likely_esp32 = True
        if vid == 0x10c4 and pid == 0xea60:
            is_likely_esp32 = True
        if vid == 0x0403 and pid == 0x6001:
            is_likely_esp32 = True
        if '/ttyACM' in dev or '/ttyUSB' in dev:
            is_likely_esp32 = True

        if is_likely_esp32:
            candidates.append((dev, vid, pid, p.description))
    return candidates


def is_port_in_use(port: str) -> bool:
    """检查串口是否被其他进程占用（通过 /proc/*/fd/ 检查，不实际打开设备）"""
    try:
        real_port = os.path.realpath(port)
        my_pid = str(os.getpid())
        for fd_dir in glob.glob('/proc/[0-9]*/fd'):
            try:
                pid = fd_dir.split('/')[2]
                if pid == my_pid:
                    continue  # 跳过自己
                for fd in os.listdir(fd_dir):
                    fd_path = os.path.join(fd_dir, fd)
                    try:
                        if os.readlink(fd_path) == real_port:
                            try:
                                with open(f'/proc/{pid}/comm', 'r') as f:
                                    comm = f.read().strip()
                                print(f"[AutoDetect] {port}: already used by process {pid} ({comm}), skipping")
                            except Exception:
                                print(f"[AutoDetect] {port}: already used by process {pid}, skipping")
                            return True
                    except (OSError, ValueError):
                        continue
            except (OSError, PermissionError):
                continue
    except Exception:
        pass
    return False


def try_open_port(port: str, baud: int, timeout: float = 5.0):
    """
    尝试打开串口并监听心跳。
    返回 (serial_obj, success) 或 (None, False)
    如果串口被其他应用占用，直接跳过不尝试打开，避免影响其他程序。
    """
    # 先检查是否被其他进程占用，避免 probing 导致其他程序 crash
    if is_port_in_use(port):
        return None, False

    try:
        s = serial.Serial(
            port, baud,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=0.1,
            write_timeout=1
        )
        s.dtr = False
        s.rts = False
    except serial.SerialException as e:
        # 串口被占用或不存在
        print(f"[AutoDetect] {port}: cannot open ({e})")
        return None, False
    except Exception as e:
        print(f"[AutoDetect] {port}: unexpected error ({e})")
        return None, False

    # 清空缓冲区，等待心跳
    s.reset_input_buffer()
    deadline = time.time() + timeout
    buffer = b""

    while time.time() < deadline:
        try:
            waiting = s.in_waiting
            if waiting > 0:
                buffer += s.read(waiting)
                if HEARTBEAT_PREFIX in buffer or HEARTBEAT_PREFIX_LEGACY in buffer:
                    print(f"[AutoDetect] {port}: heartbeat found!")
                    return s, True
                # 限制缓冲区大小，避免内存泄漏
                if len(buffer) > 4096:
                    buffer = buffer[-2048:]
            time.sleep(0.05)
        except serial.SerialException:
            s.close()
            return None, False
        except Exception as e:
            print(f"[AutoDetect] {port}: read error ({e})")
            s.close()
            return None, False

    s.close()
    print(f"[AutoDetect] {port}: no heartbeat within {timeout}s")
    return None, False


def auto_detect_serial():
    """自动扫描并连接对讲机设备串口"""
    global ser, SERIAL_PORT

    candidates = list_candidate_ports()
    if not candidates:
        print("[AutoDetect] No candidate serial ports found")
        return None

    print(f"[AutoDetect] Scanning {len(candidates)} candidate port(s):")
    for dev, vid, pid, desc in candidates:
        print(f"  {dev}  VID={vid:#06x} PID={pid:#06x} ({desc})")

    for dev, vid, pid, desc in candidates:
        print(f"[AutoDetect] Probing {dev}...")
        s, ok = try_open_port(dev, SERIAL_BAUD)
        if ok and s is not None:
            SERIAL_PORT = dev
            with serial_lock:
                ser = s
            print(f"[AutoDetect] Connected to device on {dev}")
            return s

    print("[AutoDetect] No device found on any port")
    return None


def serial_monitor_thread():
    """
    后台线程：持续监控串口连接状态。
    - 如果串口断开，自动重新扫描
    - 定期轮询所有候选端口
    """
    global ser, SERIAL_PORT
    print("[SerialMonitor] Monitor thread started")
    last_detect_time = 0

    while True:
        try:
            with serial_lock:
                is_open = (ser is not None and ser.is_open)

            if not is_open:
                now = time.time()
                if now - last_detect_time >= HEARTBEAT_INTERVAL:
                    last_detect_time = now
                    print("[SerialMonitor] Serial disconnected, auto-detecting...")
                    auto_detect_serial()
                else:
                    time.sleep(1)
            else:
                # 串口正常，检查心跳是否超时
                if not is_serial_connected():
                    print("[SerialMonitor] Heartbeat timeout, closing serial...")
                    close_serial()
                    time.sleep(1)
                else:
                    time.sleep(1)
        except Exception as e:
            print(f"[SerialMonitor] error: {e}")
            time.sleep(2)


def get_serial():
    """获取当前串口对象（如果未连接则返回 None）"""
    with serial_lock:
        if ser is not None and ser.is_open:
            return ser
    return None


def is_serial_connected():
    """判断串口是否真正连接（is_open + 心跳未超时）"""
    with serial_lock:
        if ser is None or not ser.is_open:
            return False
    # 即使 is_open 为 True，也要检查心跳超时
    global last_heartbeat_time
    if last_heartbeat_time > 0 and (time.time() - last_heartbeat_time) > HEARTBEAT_TIMEOUT_SEC:
        return False
    return True


def close_serial():
    """安全关闭串口并清理状态"""
    global ser, SERIAL_PORT, last_heartbeat_time
    with serial_lock:
        try:
            if ser:
                ser.close()
        except Exception:
            pass
        ser = None
        SERIAL_PORT = None
        last_heartbeat_time = 0.0


def process_pcm(pcm_data: bytes) -> bytes:
    """噪声门 + 线性增益"""
    samples = len(pcm_data) // 2
    if samples == 0:
        return pcm_data
    fmt = "<" + "h" * samples
    samples_list = struct.unpack(fmt, pcm_data)

    NOISE_GATE = 150
    gated = [s if abs(s) >= NOISE_GATE else 0 for s in samples_list]

    max_abs = max(abs(s) for s in gated)
    if max_abs == 0:
        return pcm_data
    TARGET_PEAK = 29490
    gain = TARGET_PEAK / max_abs

    amplified = [int(max(-32768, min(32767, s * gain))) for s in gated]
    print(f"[Audio] gate={NOISE_GATE}, gain={gain:.2f}x, "
          f"peak_before={max_abs}, peak_after={max(abs(s) for s in amplified)}")
    return struct.pack(fmt, *amplified)


def resample_to_44k_mono(filepath: str, max_seconds: float = None) -> bytes:
    """使用 ffmpeg 将任意格式音频重采样为 44.1kHz 单声道 16bit raw PCM，限制最长时长"""
    limit = max_seconds if max_seconds is not None else MAX_AUDIO_SECONDS
    print(f"[FFmpeg] Resampling {filepath} to 44100Hz mono 16bit (max {limit}s)...")
    cmd = [
        "ffmpeg", "-y", "-i", filepath,
        "-t", str(limit),
        "-ar", "44100", "-ac", "1", "-f", "s16le", "-"
    ]
    result = subprocess.run(cmd, capture_output=True)
    if result.returncode != 0:
        err = result.stderr.decode("utf-8", errors="replace").strip()
        print(f"[FFmpeg] Error: {err}")
        return None
    pcm = result.stdout
    print(f"[FFmpeg] Resampled to {len(pcm)} bytes ({len(pcm)//2} samples)")
    return pcm


def _send_audio_worker(filepath: str, stop_event: threading.Event):
    """后台线程：发送 PCM 数据到串口，支持中途停止"""
    global playback_status
    s = get_serial()
    if s is None:
        print("[Serial] Not available, skipping send")
        with status_lock:
            playback_status["playing"] = False
        return

    pcm_data = resample_to_44k_mono(filepath)
    if pcm_data is None:
        with status_lock:
            playback_status["playing"] = False
        return

    pcm_data = process_pcm(pcm_data)

    total = len(pcm_data)
    sent = 0

    preview = pcm_data[:32]
    print(f"[Serial] PCM preview: {preview.hex()}")
    nonzero = sum(1 for b in pcm_data if b != 0)
    print(f"[Serial] PCM non-zero: {nonzero}/{len(pcm_data)} ({100*nonzero/len(pcm_data):.1f}%)")

    print(f"[Serial] Sending {total} bytes to device @ {SERIAL_BAUD}...")

    with serial_lock:
        s.reset_input_buffer()
        s.reset_output_buffer()

    start_time = time.time()
    next_send_time = start_time

    with status_lock:
        playback_status["playing"] = True
        playback_status["sent_bytes"] = 0
        playback_status["total_bytes"] = total
        playback_status["started_at"] = time.strftime("%Y-%m-%d %H:%M:%S")
        playback_status["stopped_at"] = None

    while sent < total and not stop_event.is_set():
        end = min(sent + SEND_CHUNK_BYTES, total)
        chunk = pcm_data[sent:end]

        now = time.time()
        wait = next_send_time - now
        if wait > 0:
            time.sleep(wait)

        with serial_lock:
            s.write(chunk)
        sent += len(chunk)
        next_send_time += SEND_INTERVAL_MS / 1000.0

        if sent % (SEND_CHUNK_BYTES * 10) == 0 or sent == total:
            elapsed = time.time() - start_time
            rate = sent / elapsed if elapsed > 0 else 0
            print(f"[Serial] Progress: {sent}/{total} bytes ({rate:.1f} bytes/s)")
            with status_lock:
                playback_status["sent_bytes"] = sent

    elapsed = time.time() - start_time
    stopped = stop_event.is_set()
    with status_lock:
        playback_status["playing"] = False
        playback_status["sent_bytes"] = sent
        playback_status["stopped_at"] = time.strftime("%Y-%m-%d %H:%M:%S")
    if stopped:
        print(f"[Serial] Send STOPPED by user ({sent}/{total} bytes in {elapsed:.2f}s)")
    else:
        print(f"[Serial] Send complete ({sent} bytes in {elapsed:.2f}s, "
              f"avg {sent/elapsed:.1f} bytes/s)")


def start_playback(filepath: str):
    """启动音频播放（先停止当前播放）"""
    global current_play_thread, play_stop_event
    stop_playback()
    play_stop_event = threading.Event()
    t = threading.Thread(target=_send_audio_worker, args=(filepath, play_stop_event), daemon=True)
    t.start()
    current_play_thread = t
    return True


def stop_playback():
    """停止当前播放"""
    global current_play_thread, play_stop_event
    play_stop_event.set()
    if current_play_thread and current_play_thread.is_alive():
        current_play_thread.join(timeout=2.0)
    current_play_thread = None
    return True


# ------------------- 串口回读线程：接收设备调试信息 -------------------
def serial_reader_thread():
    """后台线程：持续读取串口回传数据并打印（过滤心跳）"""
    global last_heartbeat_time
    print("[Serial RX] Reader thread started")
    while True:
        try:
            s = get_serial()
            if s and s.is_open:
                with serial_lock:
                    waiting = s.in_waiting
                    if waiting > 0:
                        data = s.read(waiting)
                    else:
                        data = b""

                if data:
                    # 过滤心跳行，同时记录心跳时间
                    lines = data.split(b'\n')
                    filtered = []
                    for line in lines:
                        line = line.strip()
                        if line:
                            if line.startswith(HEARTBEAT_PREFIX) or line.startswith(HEARTBEAT_PREFIX_LEGACY):
                                last_heartbeat_time = time.time()
                            else:
                                filtered.append(line)
                    if filtered:
                        for line in filtered:
                            try:
                                text = line.decode("utf-8", errors="replace")
                                print(f"[Serial RX] {text}")
                            except Exception:
                                print(f"[Serial RX] (hex) {line.hex()}")
            time.sleep(0.05)
        except serial.SerialException as e:
            print(f"[Serial RX] SerialException: {e}, closing...")
            close_serial()
            time.sleep(1)
        except Exception as e:
            print(f"[Serial RX] error: {e}")
            time.sleep(1)


@app.route("/version", methods=["GET"])
def version():
    serial_status = "connected" if is_serial_connected() else "disconnected"
    return jsonify({
        "version": VERSION,
        "serial_port": SERIAL_PORT,
        "serial_baud": SERIAL_BAUD,
        "serial_status": serial_status,
        "sample_rate": 44100,
        "channels": 1,
        "bits_per_sample": 16,
        "max_audio_seconds": MAX_AUDIO_SECONDS,
        "server_time": time.strftime("%Y-%m-%d %H:%M:%S")
    })


@app.route("/", methods=["GET"])
def index():
    return f"""
    <h2>UAV Speaker Server v{VERSION}</h2>
    <p>POST /upload (multipart/form-data, field='audio')</p>
    <p>GET /version (JSON API)</p>
    <p>Audio format: any WAV -> server resamples to 44.1kHz mono 16bit</p>
    <p>Serial: auto-detect (heartbeat) @ {SERIAL_BAUD} baud</p>
    """


import zlib

@app.route("/upload", methods=["POST"])
def upload_audio():
    if "audio" not in request.files:
        return jsonify({"error": "No audio file"}), 400

    file = request.files["audio"]
    if file.filename == "":
        return jsonify({"error": "Empty filename"}), 400

    compressed = request.args.get('compressed') == '1'
    filepath = os.path.join(UPLOAD_FOLDER, "latest_recording.wav")

    if compressed:
        # 解压 zlib 压缩的数据
        try:
            raw_bytes = file.read()
            decompressed = zlib.decompress(raw_bytes)
            print(f"[HTTP] Received compressed audio: {len(raw_bytes)} bytes "
                  f"-> decompressed {len(decompressed)} bytes")
            with open(filepath, "wb") as f:
                f.write(decompressed)
            print(f"[HTTP] Saved decompressed WAV: {filepath} ({os.path.getsize(filepath)} bytes)")
        except zlib.error as e:
            print(f"[HTTP] Decompression failed: {e}")
            return jsonify({"error": f"Decompression failed: {e}"}), 400
        except Exception as e:
            print(f"[HTTP] Decompression error: {e}")
            return jsonify({"error": f"Decompression error: {e}"}), 400
    else:
        file.save(filepath)
        print(f"[HTTP] Received audio: {filepath} ({os.path.getsize(filepath)} bytes)")

    start_playback(filepath)
    return jsonify({"success": True, "playback_started": True}), 200


@app.route("/stop", methods=["POST"])
def stop_audio():
    """停止当前播放"""
    stop_playback()
    return jsonify({"success": True, "message": "Playback stopped"}), 200


@app.route("/status", methods=["GET"])
def get_status():
    """获取服务器状态（含播放状态）"""
    serial_status_val = "connected" if is_serial_connected() else "disconnected"
    with status_lock:
        status = dict(playback_status)
    return jsonify({
        "version": VERSION,
        "serial_port": SERIAL_PORT,
        "serial_baud": SERIAL_BAUD,
        "serial_status": serial_status_val,
        "sample_rate": 44100,
        "max_audio_seconds": MAX_AUDIO_SECONDS,
        "server_time": time.strftime("%Y-%m-%d %H:%M:%S"),
        "playback": status,
    })


@app.route("/ports", methods=["GET"])
def list_ports():
    candidates = list_candidate_ports()
    return jsonify({
        "ports": [p[0] for p in candidates],
        "configured": SERIAL_PORT,
        "candidates": [
            {"device": p[0], "vid": f"{p[1]:#06x}", "pid": f"{p[2]:#06x}", "desc": p[3]}
            for p in candidates
        ]
    })


if __name__ == "__main__":
    # 启动串口监控线程（自动检测 + 热插拔）
    serial_detect_thread = threading.Thread(
        target=serial_monitor_thread, daemon=True
    )
    serial_detect_thread.start()

    # 启动串口回读线程
    threading.Thread(target=serial_reader_thread, daemon=True).start()

    # 启动 Flask 服务
    app.run(host="0.0.0.0", port=5000, debug=False)
