"""设备 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


UPGRADE_COMMAND = 4001


MAX_JSON_SIZE = 1024 * 1024


MAX_STATUS_SIZE = 1024 * 1024


MAX_ROM_SIZE = 0xFFFFFFFF


STATUS_SIZE = 21


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


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 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]