#!/usr/bin/env python3
"""V1/V2 身份预检与 EPFP/4001 固件升级工具。"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import select
import socket
import stat
import struct
import sys
import time
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple


INFO_PORT = 8942
UPGRADE_PORT = 8943
INFO_COMMAND = 30101
V1_INFO_PORT = 4941
V1_INFO_COMMAND = 0x5A0E2104
UPGRADE_COMMAND = 4001
MAX_JSON_SIZE = 1024 * 1024
MAX_STATUS_SIZE = 1024 * 1024
MAX_ROM_SIZE = 0xFFFFFFFF
STATUS_SIZE = 21


def configure_redirected_output() -> None:
    """输出被管道捕获时统一使用 UTF-8，避免中文日志乱码。"""

    for stream in (sys.stdout, sys.stderr):
        if not stream.isatty() and hasattr(stream, "reconfigure"):
            stream.reconfigure(encoding="utf-8", errors="backslashreplace")


configure_redirected_output()


class UpgradeError(RuntimeError):
    """表示预检、协议或升级过程失败。"""


class EventLogger:
    """同时向终端和 JSONL 文件写入带时间戳的事件。"""

    def __init__(self, log_dir: Path) -> None:
        log_dir.mkdir(parents=True, exist_ok=True)
        stamp = time.strftime("%Y%m%d-%H%M%S")
        self.path = log_dir / ("device-upgrade-" + stamp + ".jsonl")
        self._stream = self.path.open("a", encoding="utf-8", buffering=1)

    def close(self) -> None:
        self._stream.close()

    def emit(self, event: str, **fields: Any) -> None:
        record: Dict[str, Any] = {
            "time": time.strftime("%Y-%m-%d %H:%M:%S"),
            "event": event,
        }
        record.update(fields)
        line = json.dumps(record, ensure_ascii=False, separators=(",", ":"))
        print(line, flush=True)
        self._stream.write(line + "\n")


def recv_exact(sock: socket.socket, size: int) -> bytes:
    """读取指定字节数；对端提前关闭时立即失败。"""

    data = bytearray()
    while len(data) < size:
        chunk = sock.recv(size - len(data))
        if not chunk:
            raise UpgradeError("设备在响应完整之前关闭了连接")
        data.extend(chunk)
    return bytes(data)


def query_device_info(host: str, timeout: float, logger: EventLogger) -> Dict[str, Any]:
    """通过 EPTP/30101 查询并严格校验设备基础信息。"""

    session = int(time.time() * 1000) & 0xFFFFFFFF
    request_object = {"cmd": INFO_COMMAND, "session": session}
    request = json.dumps(request_object, separators=(",", ":")).encode("utf-8")
    logger.emit("info_request", host=host, port=INFO_PORT, request=request_object)

    with socket.create_connection((host, INFO_PORT), timeout=timeout) as sock:
        sock.settimeout(timeout)
        sock.sendall(b"EPTP" + struct.pack("!I", len(request)) + request)
        header = recv_exact(sock, 8)
        if header[:4] != b"EPTP":
            raise UpgradeError("基础信息响应同步头不是 EPTP：%r" % (header[:4],))
        payload_size = struct.unpack("!I", header[4:])[0]
        if payload_size == 0 or payload_size > MAX_JSON_SIZE:
            raise UpgradeError("基础信息响应长度不合理：%d" % payload_size)
        payload = recv_exact(sock, payload_size)

    try:
        response = json.loads(payload.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise UpgradeError("基础信息响应不是有效 UTF-8 JSON：%s" % exc) from exc
    if not isinstance(response, dict):
        raise UpgradeError("基础信息响应 JSON 顶层不是对象")
    if response.get("cmd") != INFO_COMMAND:
        raise UpgradeError("基础信息响应 cmd 不匹配：%r" % response.get("cmd"))
    if response.get("session") != session:
        raise UpgradeError(
            "基础信息响应 session 不匹配：期望 %d，实际 %r"
            % (session, response.get("session"))
        )
    if response.get("code") != 0:
        raise UpgradeError("基础信息响应 code 非 0：%r" % response.get("code"))
    logger.emit("device_info", response=response)
    return response


def query_v1_device_info(host: str, timeout: float, logger: EventLogger) -> Dict[str, Any]:
    """通过 V1 工厂 TCP/2104 获取设备身份与固件字段。"""
    logger.emit("v1_info_request", host=host, port=V1_INFO_PORT,
                command="0x%08X" % V1_INFO_COMMAND)
    with socket.create_connection((host, V1_INFO_PORT), timeout=timeout) as sock:
        sock.settimeout(timeout)
        sock.sendall(struct.pack("!II", V1_INFO_COMMAND, 0))
        header = recv_exact(sock, 8)
        command, size = struct.unpack("!II", header)
        if command != V1_INFO_COMMAND or size > MAX_JSON_SIZE:
            raise UpgradeError("V1 固件信息响应 CMD 或长度异常")
        payload = recv_exact(sock, size)
    try:
        response = json.loads(payload.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise UpgradeError("V1 固件信息响应不是有效 JSON：%s" % exc) from exc
    if not isinstance(response, dict) or response.get("code") != 0:
        raise UpgradeError("V1 固件信息查询失败或响应格式异常")
    if not response.get("sn") or not response.get("dev_type"):
        raise UpgradeError("V1 固件信息缺少 SN 或 dev_type，拒绝升级")
    logger.emit("device_info", protocol="V1", response=response)
    return response


def response_data(response: Dict[str, Any]) -> Dict[str, Any]:
    """取得响应 data 对象，格式错误时拒绝继续。"""

    data = response.get("data")
    if not isinstance(data, dict):
        raise UpgradeError("基础信息响应缺少 data 对象")
    return data


def collect_models(data: Dict[str, Any]) -> List[str]:
    """收集 inner、outer 和 hardware 中可用的型号。"""

    models: List[str] = []
    for section_name in ("inner", "outer", "hardware"):
        section = data.get(section_name)
        if isinstance(section, dict):
            value = section.get("type")
            if isinstance(value, str) and value and value not in models:
                models.append(value)
    return models


def collect_versions(data: Dict[str, Any]) -> Dict[str, Any]:
    """收集各固件分区和硬件版本，便于升级前后比对。"""

    versions: Dict[str, Any] = {}
    for section_name in ("inner", "outer", "hardware"):
        section = data.get(section_name)
        if isinstance(section, dict):
            versions[section_name] = section.get("version")
    return versions


def validate_identity(
    response: Dict[str, Any], expected_sn: str, expected_model: Optional[str]
) -> Tuple[Dict[str, Any], List[str]]:
    """对完整 SN 做强校验，并按需校验型号。"""

    data = response_data(response)
    actual_sn = data.get("sn")
    if actual_sn != expected_sn:
        raise UpgradeError("设备 SN 不匹配：期望 %s，实际 %r" % (expected_sn, actual_sn))
    models = collect_models(data)
    if expected_model and expected_model not in models:
        raise UpgradeError(
            "设备型号不匹配：期望 %s，设备报告 %s"
            % (expected_model, models if models else "无可用型号")
        )
    return data, models


def infer_rom_version(path: Path) -> Optional[str]:
    """仅从路径中识别明确的四段数字版本号，不猜测其他命名。"""

    candidates: List[str] = []
    pattern = re.compile(r"(?<!\d)(\d+\.\d+\.\d+\.\d+)(?!\d)")
    for part in path.parts:
        candidates.extend(pattern.findall(part))
    return candidates[-1] if candidates else None


def validate_rom(path: Path) -> Dict[str, Any]:
    """确认 ROM 为普通非空文件，并计算 SHA-256。"""

    absolute = path.expanduser().resolve(strict=True)
    metadata = os.stat(str(absolute))
    if not stat.S_ISREG(metadata.st_mode):
        raise UpgradeError("ROM 不是普通文件：%s" % absolute)
    if metadata.st_size <= 0:
        raise UpgradeError("ROM 文件为空：%s" % absolute)
    if metadata.st_size > MAX_ROM_SIZE:
        raise UpgradeError("ROM 大于协议 uint32 长度上限：%d" % metadata.st_size)

    digest = hashlib.sha256()
    with absolute.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return {
        "path": str(absolute),
        "size": metadata.st_size,
        "modified": time.strftime(
            "%Y-%m-%d %H:%M:%S", time.localtime(metadata.st_mtime)
        ),
        "sha256": digest.hexdigest(),
        "target_version_from_path": infer_rom_version(absolute),
    }


def parse_status(payload: bytes) -> Dict[str, int]:
    """解析 21 字节升级状态业务数据。"""

    if len(payload) < STATUS_SIZE:
        raise UpgradeError("升级状态不足 21 字节：%d" % len(payload))
    status, session, download_total, downloaded, flash_total, flashed = struct.unpack(
        "!BIIIII", payload[:STATUS_SIZE]
    )
    if status not in (0, 1, 2, 3):
        raise UpgradeError("未知升级状态值：%d" % status)
    return {
        "status": status,
        "session": session,
        "download_total": download_total,
        "downloaded": downloaded,
        "flash_total": flash_total,
        "flashed": flashed,
    }


def extract_status_payloads(buffer: bytearray) -> Iterable[Tuple[str, bytes]]:
    """从 TCP 字节流中提取完整 EPFP 帧或兼容设备实测的裸状态。"""

    while buffer:
        # 同步头被拆包时等待后续数据，避免把半个 EPFP 当裸状态。
        if len(buffer) < 4 and b"EPFP".startswith(bytes(buffer)):
            return

        if buffer[:4] == b"EPFP":
            if len(buffer) < 12:
                return
            command, payload_size = struct.unpack("!II", buffer[4:12])
            if command != UPGRADE_COMMAND:
                raise UpgradeError("升级状态帧 cmd 不匹配：%d" % command)
            if payload_size < STATUS_SIZE or payload_size > MAX_STATUS_SIZE:
                raise UpgradeError("升级状态帧长度不合理：%d" % payload_size)
            if len(buffer) < 12 + payload_size:
                return
            payload = bytes(buffer[12 : 12 + payload_size])
            del buffer[: 12 + payload_size]
            yield "framed", payload
            continue

        if buffer[0] in (0, 1, 2, 3):
            # 合法裸状态内部也可能包含字节串 EPFP，不能在其业务数据中误重同步。
            if len(buffer) < STATUS_SIZE:
                return
            payload = bytes(buffer[:STATUS_SIZE])
            del buffer[:STATUS_SIZE]
            yield "raw", payload
            continue

        sync_position = buffer.find(b"EPFP")
        if sync_position > 0:
            # 首字节不是合法 status，且后方存在同步头时，丢弃无效前缀。
            del buffer[:sync_position]
            continue
        if len(buffer) < STATUS_SIZE:
            return
        # 保留这 21 字节交给状态解析器报告明确的非法 status。
        payload = bytes(buffer[:STATUS_SIZE])
        del buffer[:STATUS_SIZE]
        yield "raw", payload


def transfer_rom(
    host: str,
    rom_path: Path,
    rom_size: int,
    connect_timeout: float,
    stall_timeout: float,
    total_timeout: float,
    chunk_size: int,
    logger: EventLogger,
) -> Dict[str, Any]:
    """发送 ROM，并在同一循环中处理回包、停滞和终态。"""

    result: Dict[str, Any] = {"status": None, "bytes_sent": 0}
    receive_buffer = bytearray()
    pending = bytearray(b"EPFP" + struct.pack("!II", UPGRADE_COMMAND, rom_size))
    started = time.monotonic()
    last_progress = started
    previous_progress: Optional[Tuple[int, int]] = None
    upload_complete = False

    logger.emit(
        "upgrade_connect", host=host, port=UPGRADE_PORT, rom_size=rom_size
    )
    with socket.create_connection((host, UPGRADE_PORT), timeout=connect_timeout) as sock:
        sock.setblocking(False)
        with rom_path.open("rb") as rom:
            while True:
                now = time.monotonic()
                if now - started >= total_timeout:
                    raise UpgradeError("升级超过总超时 %.1f 秒" % total_timeout)
                if now - last_progress >= stall_timeout:
                    raise UpgradeError("设备下载或烧录进度连续 %.1f 秒无变化" % stall_timeout)

                readable, writable, exceptional = select.select(
                    [sock], [] if upload_complete else [sock], [sock], 0.5
                )
                if exceptional:
                    raise UpgradeError("升级 TCP 连接出现异常状态")

                if readable:
                    chunk = sock.recv(4096)
                    if not chunk:
                        raise UpgradeError("设备在 status=2 之前关闭升级连接")
                    logger.emit("upgrade_rx", size=len(chunk), hex=chunk.hex())
                    receive_buffer.extend(chunk)
                    for frame_type, payload in extract_status_payloads(receive_buffer):
                        status = parse_status(payload)
                        logger.emit("device_status", frame=frame_type, **status)
                        progress = (status["downloaded"], status["flashed"])
                        if previous_progress is None or progress != previous_progress:
                            previous_progress = progress
                            last_progress = time.monotonic()
                        result.update(status)
                        if status["status"] == 3:
                            raise UpgradeError("设备返回升级失败状态 status=3")
                        if status["status"] == 2:
                            incomplete_counters = {}
                            if (
                                status["download_total"] > 0
                                and status["downloaded"] < status["download_total"]
                            ):
                                incomplete_counters["download"] = {
                                    "total": status["download_total"],
                                    "current": status["downloaded"],
                                }
                            if (
                                status["flash_total"] > 0
                                and status["flashed"] < status["flash_total"]
                            ):
                                incomplete_counters["flash"] = {
                                    "total": status["flash_total"],
                                    "current": status["flashed"],
                                }
                            if incomplete_counters:
                                logger.emit(
                                    "protocol_warning",
                                    message="status=2，但设备进度计数尚未达到其自身总量",
                                    counters=incomplete_counters,
                                )
                            logger.emit("protocol_success", status=2)
                            return result

                if writable:
                    if not pending:
                        next_chunk = rom.read(chunk_size)
                        if next_chunk:
                            pending.extend(next_chunk)
                        else:
                            upload_complete = True
                            logger.emit(
                                "client_upload_complete",
                                bytes_sent=result["bytes_sent"],
                                rom_size=rom_size,
                            )
                    if pending:
                        sent = sock.send(pending)
                        if sent <= 0:
                            raise UpgradeError("升级 TCP 连接无法继续发送")
                        # 前 12 字节是 EPFP 头，不计入 ROM 上传进度。
                        previous_wire_bytes = result.get("wire_bytes_sent", 0)
                        wire_bytes = previous_wire_bytes + sent
                        result["wire_bytes_sent"] = wire_bytes
                        result["bytes_sent"] = max(0, min(rom_size, wire_bytes - 12))
                        del pending[:sent]


def build_parser() -> argparse.ArgumentParser:
    """创建命令行参数解析器。"""

    parser = argparse.ArgumentParser(
        description="设备身份预检与 EPFP/4001 ROM 升级；默认仅预检。"
    )
    parser.add_argument("--host", required=True, help="目标设备 IPv4/主机名")
    parser.add_argument("--protocol", choices=("v1", "v2"), default="v2",
                        help="身份查询协议；固件传输均使用 EPFP/4001")
    parser.add_argument("--expected-sn", required=True, help="必须精确匹配的完整设备 SN")
    parser.add_argument("--expected-model", help="可选，必须出现在设备型号字段中的型号")
    parser.add_argument("--rom", required=True, type=Path, help="ROM 文件绝对或相对路径")
    parser.add_argument("--log-dir", required=True, type=Path, help="本次任务日志目录")
    parser.add_argument(
        "--execute",
        action="store_true",
        help="实际传输 ROM；未指定时只查询身份、校验 ROM 和输出预检信息",
    )
    parser.add_argument("--connect-timeout", type=float, default=10.0)
    parser.add_argument("--stall-timeout", type=float, default=60.0)
    parser.add_argument("--total-timeout", type=float, default=600.0)
    parser.add_argument("--chunk-size", type=int, default=256 * 1024)
    return parser


def validate_numeric_options(args: argparse.Namespace) -> None:
    """拒绝会让超时或分块保护失效的参数。"""

    if args.connect_timeout <= 0:
        raise UpgradeError("--connect-timeout 必须大于 0")
    if args.stall_timeout <= 0:
        raise UpgradeError("--stall-timeout 必须大于 0")
    if args.total_timeout <= 0:
        raise UpgradeError("--total-timeout 必须大于 0")
    if args.chunk_size <= 0 or args.chunk_size > 16 * 1024 * 1024:
        raise UpgradeError("--chunk-size 必须在 1 到 16777216 之间")


def main(argv: Optional[List[str]] = None) -> int:
    """执行预检；仅在明确给出 --execute 时进入升级传输。"""

    args = build_parser().parse_args(argv)
    logger: Optional[EventLogger] = None
    try:
        validate_numeric_options(args)
        logger = EventLogger(args.log_dir.expanduser().resolve())
        logger.emit("mode", execute=bool(args.execute), log_file=str(logger.path))
        rom = validate_rom(args.rom)
        logger.emit("rom", **rom)
        if args.protocol == "v1":
            response = query_v1_device_info(args.host, args.connect_timeout, logger)
            if response.get("sn") != args.expected_sn:
                raise UpgradeError("设备 SN 不匹配：期望 %s，实际 %r" % (args.expected_sn, response.get("sn")))
            if args.expected_model and args.expected_model not in str(response.get("dev_type", "")):
                raise UpgradeError("设备型号不匹配：期望 %s，实际 %r" % (args.expected_model, response.get("dev_type")))
            data = {"sn": response.get("sn"), "inner": {"type": response.get("dev_type"),
                    "version": response.get("firmware")}, "outer": {"version": response.get("dev_ver")},
                    "hardware": {"version": response.get("hardware")}}
            models = [str(response["dev_type"])]
        else:
            response = query_device_info(args.host, args.connect_timeout, logger)
            data, models = validate_identity(response, args.expected_sn, args.expected_model)
        summary = {
            "host": args.host,
            "identity_protocol": args.protocol,
            "upgrade_protocol_compatibility": "V1 身份查询 + EPFP/4001；仍须依据目标固件资料确认兼容性",
            "sn": data.get("sn"),
            "mac": data.get("mac"),
            "models": models,
            "versions": collect_versions(data),
            "expected_model": args.expected_model,
            "rom": rom,
            "rom_model_verification": (
                "由调用方提供并已核对设备型号"
                if args.expected_model
                else "无法从协议或 ROM 自动验证目标型号"
            ),
        }
        logger.emit("preflight", **summary)

        if not args.execute:
            logger.emit("preflight_only", message="未指定 --execute，未连接 8943，未发送 ROM")
            return 0

        result = transfer_rom(
            host=args.host,
            rom_path=Path(rom["path"]),
            rom_size=int(rom["size"]),
            connect_timeout=args.connect_timeout,
            stall_timeout=args.stall_timeout,
            total_timeout=args.total_timeout,
            chunk_size=args.chunk_size,
            logger=logger,
        )
        if result.get("status") != 2:
            raise UpgradeError("升级没有取得 status=2 成功终态")
        logger.emit(
            "upgrade_result",
            protocol_success=True,
            postcheck_required=True,
            result=result,
        )
        return 0
    except (OSError, ValueError, UpgradeError) as exc:
        if logger is not None:
            logger.emit("error", type=type(exc).__name__, message=str(exc))
        else:
            print("错误：%s" % exc, file=sys.stderr, flush=True)
        return 1
    finally:
        if logger is not None:
            logger.close()


if __name__ == "__main__":
    raise SystemExit(main())