#!/usr/bin/env python3
"""UART host runner for the TQ printer product-test protocol."""

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional

try:
    import serial
    from serial.tools import list_ports
except ImportError as exc:  # pragma: no cover - depends on operator env
    raise SystemExit(
        "缺少 pyserial，请先执行：python3 -m pip install pyserial"
    ) from exc


TOKEN_READY = "TQTEST_READY"
TOKEN_PROGRESS = "TQTEST_PROGRESS"
TOKEN_RESULT = "TQTEST_RESULT"
TOKEN_ERROR = "TQTEST_ERROR"
TOKEN_CMD = "TQTEST_CMD"


@dataclass
class ProtocolFrame:
    token: str
    payload: Dict[str, Any]
    raw: str


class ProtocolError(RuntimeError):
    pass


class ProductTestClient:
    def __init__(
        self,
        port: str,
        baud: int,
        timeout: float,
        echo_logs: bool = False,
        verbose: bool = False,
        jsonl_path: Optional[Path] = None,
    ) -> None:
        self.port = port
        self.baud = baud
        self.timeout = timeout
        self.echo_logs = echo_logs
        self.verbose = verbose
        self.jsonl_path = jsonl_path
        self.seq = 1
        self.serial = serial.Serial(
            port=port,
            baudrate=baud,
            timeout=0.2,
            write_timeout=2.0,
        )
        self.jsonl = jsonl_path.open("a", encoding="utf-8") if jsonl_path else None
        self._last_probe = 0.0

    def close(self) -> None:
        if self.jsonl:
            self.jsonl.close()
            self.jsonl = None
        if self.serial and self.serial.is_open:
            self.serial.close()

    def __enter__(self) -> "ProductTestClient":
        return self

    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
        self.close()

    def _log_jsonl(self, direction: str, item: Dict[str, Any]) -> None:
        if not self.jsonl:
            return
        record = {
            "ts": time.time(),
            "direction": direction,
            **item,
        }
        self.jsonl.write(json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n")
        self.jsonl.flush()

    def _read_line(self, deadline: float) -> Optional[str]:
        while time.monotonic() < deadline:
            raw = self.serial.readline()
            if not raw:
                continue
            text = raw.decode("utf-8", errors="replace").rstrip("\r\n")
            if not text:
                continue
            self._log_jsonl("rx", {"raw": text})
            if self.echo_logs and not text.startswith("TQTEST_"):
                print(f"[log] {text}")
            return text
        return None

    def hard_reset(self, delay_s: float) -> None:
        self.serial.reset_input_buffer()
        self.serial.dtr = False
        self.serial.rts = True
        time.sleep(0.1)
        self.serial.rts = False
        time.sleep(delay_s)
        self.serial.reset_input_buffer()
        self._last_probe = 0.0

    @staticmethod
    def parse_frame(line: str) -> Optional[ProtocolFrame]:
        if not line.startswith("TQTEST_"):
            return None
        parts = line.split(" ", 1)
        token = parts[0]
        if token not in {TOKEN_READY, TOKEN_PROGRESS, TOKEN_RESULT, TOKEN_ERROR}:
            return ProtocolFrame(
                token=token,
                payload={"protocol_error": "unknown_token"},
                raw=line,
            )
        payload_text = parts[1] if len(parts) == 2 else "{}"
        try:
            payload = json.loads(payload_text)
        except json.JSONDecodeError:
            payload = {"parse_error": True, "raw_payload": payload_text}
        if not isinstance(payload, dict):
            payload = {"value": payload}
        return ProtocolFrame(token=token, payload=payload, raw=line)

    @staticmethod
    def _frame_parse_error(frame: ProtocolFrame, context: str) -> ProtocolError:
        raw = frame.raw
        if len(raw) > 240:
            raw = raw[:237] + "..."
        if frame.payload.get("protocol_error") == "unknown_token":
            return ProtocolError(
                f"{context} 收到未知协议帧: token={frame.token} raw={raw!r}"
            )
        return ProtocolError(f"{context} 收到协议帧 JSON 损坏: token={frame.token} raw={raw!r}")

    def wait_ready(self, timeout_s: float) -> Dict[str, Any]:
        deadline = time.monotonic() + timeout_s
        while time.monotonic() < deadline:
            now = time.monotonic()
            if now - self._last_probe >= 1.0:
                self._last_probe = now
                self.send_command("PING", {})
            line = self._read_line(deadline)
            if line is None:
                break
            frame = self.parse_frame(line)
            if not frame:
                continue
            if frame.payload.get("parse_error") or frame.payload.get("protocol_error"):
                raise self._frame_parse_error(frame, "等待 READY")
            if frame.token == TOKEN_READY:
                return frame.payload
            if frame.token == TOKEN_RESULT and frame.payload.get("name") == "PING":
                return {
                    "version": frame.payload.get("data", {}).get("version"),
                    "device": "TQ_printer",
                    "probe": "PING",
                }
        raise ProtocolError(f"等待 {TOKEN_READY} 超时")

    def send_command(self, command: str, payload: Optional[Dict[str, Any]] = None) -> int:
        seq = self.seq
        self.seq += 1
        body = {
            "seq": seq,
            "cmd": command,
            "payload": payload or {},
        }
        line = TOKEN_CMD + " " + json.dumps(body, ensure_ascii=False, separators=(",", ":"))
        self._log_jsonl("tx", {"raw": line, "seq": seq, "cmd": command})
        self.serial.write((line + "\n").encode("utf-8"))
        self.serial.flush()
        return seq

    def command(
        self,
        command: str,
        payload: Optional[Dict[str, Any]] = None,
        timeout_s: Optional[float] = None,
    ) -> Dict[str, Any]:
        seq = self.send_command(command, payload)
        deadline = time.monotonic() + (timeout_s or self.timeout)
        while time.monotonic() < deadline:
            line = self._read_line(deadline)
            if line is None:
                break
            frame = self.parse_frame(line)
            if not frame:
                continue
            if frame.payload.get("parse_error") or frame.payload.get("protocol_error"):
                raise self._frame_parse_error(frame, command)
            frame_seq = int(frame.payload.get("seq", -1))
            if frame_seq != seq:
                if frame.token == TOKEN_READY:
                    raise ProtocolError(f"{command} 执行期间设备重启，等待结果失败 seq={seq}")
                continue
            if frame.token == TOKEN_PROGRESS:
                if self.verbose:
                    stage = frame.payload.get("stage", "-")
                    print(f"  - {command}: {stage}")
                continue
            if frame.token == TOKEN_RESULT:
                return frame.payload
            if frame.token == TOKEN_ERROR:
                raise ProtocolError(f"{command} 返回错误帧: {frame.payload}")
        raise ProtocolError(f"{command} 等待结果超时 seq={seq}")


def list_serial_ports() -> None:
    ports = list(list_ports.comports())
    if not ports:
        print("未发现串口")
        return
    for item in ports:
        desc = item.description or "-"
        hwid = item.hwid or "-"
        print(f"{item.device}\t{desc}\t{hwid}")


def default_steps(args: argparse.Namespace) -> List[Dict[str, Any]]:
    steps: List[Dict[str, Any]] = [
        {"cmd": "PING", "timeout": 10},
        {"cmd": "STATUS", "timeout": 10},
        {"cmd": "WAIT_WIFI", "payload": {"timeout_ms": args.wifi_timeout_ms}, "timeout": args.wifi_timeout_ms / 1000 + 5},
        {"cmd": "SCREEN", "payload": {"text": "小Q打印机\n屏幕测试\n显示正常"}, "timeout": 15},
        {"cmd": "SCREEN_FLOW", "payload": {"step_ms": 400}, "timeout": 20},
        {"cmd": "PRINTER_STATUS", "timeout": 10},
        {"cmd": "AUDIO_PROBE", "payload": {"samples": 320, "timeout_ms": 3000}, "timeout": 15},
        {"cmd": "VOICE_SESSION", "payload": {"timeout_ms": args.voice_timeout_ms}, "timeout": args.voice_timeout_ms / 1000 + 10},
        {"cmd": "NEGATIVE", "payload": {"case": "invalid_print"}, "timeout": 10},
    ]

    if args.print:
        steps.append(
            {
                "cmd": "PRINT_PATTERN",
                "payload": {
                    "height": args.print_height,
                    "wait_done": True,
                    "ignore_precheck": args.ignore_precheck,
                    "timeout_ms": args.print_timeout_ms,
                },
                "timeout": args.print_timeout_ms / 1000 + 10,
            }
        )

    if args.image:
        image_cmd = "IMAGE_PRINT" if args.image_print else "IMAGE_GENERATE"
        steps.append(
            {
                "cmd": image_cmd,
                "payload": {
                    "prompt": args.prompt,
                    "size": args.image_size,
                    "prompt_extend": True,
                    "ignore_precheck": args.ignore_precheck,
                    "generation_timeout_ms": args.image_timeout_ms,
                    "download_timeout_ms": args.image_download_timeout_ms,
                    "print_timeout_ms": args.print_timeout_ms,
                },
                "timeout": (args.image_timeout_ms + args.image_download_timeout_ms + args.print_timeout_ms) / 1000 + 20,
            }
        )

    if args.stability_loops > 0:
        steps.append(
            {
                "cmd": "STABILITY",
                "payload": {
                    "loops": args.stability_loops,
                    "include_print": args.stability_print,
                    "include_audio": True,
                },
                "timeout": max(30, args.stability_loops * (8 if args.stability_print else 3)),
            }
        )

    steps.append({"cmd": "VOICE_STOP", "payload": {"timeout_ms": 3000}, "timeout": 10})
    steps.append({"cmd": "STATUS", "timeout": 10})
    return steps


def load_plan(path: Path) -> List[Dict[str, Any]]:
    data = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, list):
        raise SystemExit("测试计划必须是 JSON 数组")
    steps: List[Dict[str, Any]] = []
    for idx, item in enumerate(data, start=1):
        if not isinstance(item, dict) or "cmd" not in item:
            raise SystemExit(f"测试计划第 {idx} 项缺少 cmd")
        steps.append(item)
    return steps


def _printer_status_from_result(result: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    data = result.get("data")
    if not isinstance(data, dict):
        return None
    printer = data.get("printer")
    if isinstance(printer, dict):
        return printer
    if "temperature_c" in data or "battery_percent" in data or "has_paper" in data:
        return data
    return None


def _printer_status_summary(result: Dict[str, Any]) -> str:
    printer = _printer_status_from_result(result)
    if not printer:
        return ""

    temp = printer.get("temperature_c")
    temp_min = printer.get("temperature_min_c")
    temp_max = printer.get("temperature_max_c")
    battery = printer.get("battery_percent")
    has_paper = printer.get("has_paper")

    parts: List[str] = []
    if has_paper is not None:
        parts.append("paper=ok" if has_paper else "paper=missing")
    if battery is not None:
        parts.append(f"batt={battery}%")
    if temp is not None:
        parts.append(f"temp={float(temp):.1f}C")
    return " printer={" + " ".join(parts) + "}" if parts else ""


def summarize_result(result: Dict[str, Any]) -> str:
    status = result.get("status", "-")
    rc = result.get("rc", "-")
    message = result.get("message", "")
    suffix = f" msg={message}" if message else ""
    if status == "ok" and rc == 0 and not suffix:
        return f"ok{_printer_status_summary(result)}"
    return f"status={status} rc={rc}{suffix}{_printer_status_summary(result)}"


def run_steps(
    client: ProductTestClient,
    steps: Iterable[Dict[str, Any]],
    continue_on_fail: bool,
) -> int:
    failures = 0
    for index, step in enumerate(steps, start=1):
        cmd = str(step["cmd"])
        payload = step.get("payload") or {}
        timeout = float(step.get("timeout", client.timeout))
        try:
            result = client.command(cmd, payload=payload, timeout_s=timeout)
        except ProtocolError as exc:
            failures += 1
            print(f"[{index}] {cmd}: FAIL {exc}")
            if not continue_on_fail:
                break
            continue

        status = result.get("status")
        ok = status == "ok"
        print(f"[{index}] {cmd}: {'OK' if ok else 'FAIL'} {summarize_result(result)}")
        if not ok:
            failures += 1
            if not continue_on_fail:
                break
    return failures


def parse_args(argv: List[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run TQ printer product tests over UART.")
    parser.add_argument("-p", "--port", help="串口设备，例如 /dev/cu.usbserial-xxx")
    parser.add_argument("-b", "--baud", type=int, default=115200, help="UART 波特率")
    parser.add_argument("--list-ports", action="store_true", help="列出可用串口后退出")
    parser.add_argument("--ready-timeout", type=float, default=30.0, help="等待 TQTEST_READY 的秒数")
    parser.add_argument("--timeout", type=float, default=60.0, help="单条命令默认超时秒数")
    parser.add_argument("--reset", action="store_true", help="等待协议就绪前通过串口控制线硬复位设备")
    parser.add_argument("--reset-delay", type=float, default=2.0, help="硬复位后等待设备启动的秒数")
    parser.add_argument("--jsonl", type=Path, help="保存收发帧 JSONL")
    parser.add_argument("--echo-logs", action="store_true", help="打印非协议串口日志")
    parser.add_argument("--verbose", action="store_true", help="显示每条命令的中间阶段")
    parser.add_argument("--continue-on-fail", action="store_true", help="失败后继续执行后续步骤")
    parser.add_argument("--plan", type=Path, help="从 JSON 测试计划读取步骤")

    parser.add_argument("--print", action="store_true", help="执行实际打印测试")
    parser.add_argument("--print-height", type=int, default=96, help="打印测试图案高度")
    parser.add_argument("--print-timeout-ms", type=int, default=90000, help="打印完成等待超时")
    parser.add_argument("--ignore-precheck", action="store_true", help="打印时跳过纸张/温度/电量预检")

    parser.add_argument("--image", action="store_true", help="执行云端图像生成测试")
    parser.add_argument("--image-print", action="store_true", help="图像生成后继续打印")
    parser.add_argument("--prompt", default="可爱的猫咪线稿，白色背景，边缘清晰，适合热敏打印", help="图像生成提示词")
    parser.add_argument("--image-size", default="512*768", help="图像生成尺寸")
    parser.add_argument("--image-timeout-ms", type=int, default=90000, help="图像生成超时")
    parser.add_argument("--image-download-timeout-ms", type=int, default=30000, help="图像下载超时")

    parser.add_argument("--wifi-timeout-ms", type=int, default=20000, help="等待 Wi-Fi 就绪超时")
    parser.add_argument("--voice-timeout-ms", type=int, default=12000, help="语音会话启动超时")
    parser.add_argument("--stability-loops", type=int, default=0, help="稳定性循环次数，0 表示跳过")
    parser.add_argument("--stability-print", action="store_true", help="稳定性循环包含走纸")
    return parser.parse_args(argv)


def main(argv: List[str]) -> int:
    args = parse_args(argv)
    if args.list_ports:
        list_serial_ports()
        return 0
    if not args.port:
        raise SystemExit("必须指定 --port，或使用 --list-ports 查看串口")

    steps = load_plan(args.plan) if args.plan else default_steps(args)
    with ProductTestClient(
        port=args.port,
        baud=args.baud,
        timeout=args.timeout,
        echo_logs=args.echo_logs,
        verbose=args.verbose,
        jsonl_path=args.jsonl,
    ) as client:
        if args.reset:
            if args.verbose:
                print(f"复位设备: {args.port}")
            client.hard_reset(args.reset_delay)
        ready = client.wait_ready(args.ready_timeout)
        if args.verbose:
            print("READY " + json.dumps(ready, ensure_ascii=False, separators=(",", ":")))
        else:
            print("设备就绪")
        failures = run_steps(client, steps, args.continue_on_fail)

    if failures:
        print(f"测试完成，失败项: {failures}")
        return 1
    print("测试完成，全部通过")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))