#!/usr/bin/env python3
"""JT808 TCP ingest to TDengine history smoke and light-load verifier."""

from __future__ import annotations

import argparse
import datetime as dt
import json
import pathlib
import socket
import time
import urllib.parse
import urllib.request
from typing import Any


def bcd(value: str) -> bytes:
    digits = "".join(ch for ch in value if ch.isdigit())
    if len(digits) % 2:
        digits = "0" + digits
    return bytes((int(digits[i]) << 4) | int(digits[i + 1]) for i in range(0, len(digits), 2))


def write_u16(value: int) -> bytes:
    return bytes([(value >> 8) & 0xFF, value & 0xFF])


def write_u32(value: int) -> bytes:
    return bytes([
        (value >> 24) & 0xFF,
        (value >> 16) & 0xFF,
        (value >> 8) & 0xFF,
        value & 0xFF,
    ])


def ascii_fixed(value: str, length: int) -> bytes:
    raw = value.encode("ascii", errors="ignore")[:length]
    return raw + b" " * (length - len(raw))


def bcc(data: bytes) -> int:
    result = 0
    for byte in data:
        result ^= byte
    return result


def escape_payload(data: bytes) -> bytes:
    out = bytearray()
    for byte in data:
        if byte == 0x7E:
            out += b"\x7D\x02"
        elif byte == 0x7D:
            out += b"\x7D\x01"
        else:
            out.append(byte)
    return bytes(out)


def unescape_payload(data: bytes) -> bytes:
    out = bytearray()
    index = 0
    while index < len(data):
        byte = data[index]
        if byte == 0x7D and index + 1 < len(data):
            marker = data[index + 1]
            if marker == 0x02:
                out.append(0x7E)
                index += 2
                continue
            if marker == 0x01:
                out.append(0x7D)
                index += 2
                continue
        out.append(byte)
        index += 1
    return bytes(out)


def build_frame(message_id: int, phone: str, serial: int, body: bytes) -> bytes:
    payload = (
        write_u16(message_id)
        + write_u16(len(body) & 0x03FF)
        + bcd(phone)[-6:]
        + write_u16(serial & 0xFFFF)
        + body
    )
    payload += bytes([bcc(payload)])
    return b"\x7E" + escape_payload(payload) + b"\x7E"


def build_register_body(device_id: str, plate: str) -> bytes:
    return (
        write_u16(44)
        + write_u16(4401)
        + ascii_fixed("MAKER", 5)
        + ascii_fixed("TYPE-A", 20)
        + ascii_fixed(device_id, 7)
        + bytes([1])
        + plate.encode("ascii", errors="ignore")
    )


def build_location_body(event_time: dt.datetime, sequence: int = 0) -> bytes:
    timestamp = event_time.strftime("%y%m%d%H%M%S")
    return (
        write_u32(0)
        + write_u32(0x00030000)
        + write_u32(39916527 + sequence)
        + write_u32(116397128 + sequence)
        + write_u16(50)
        + write_u16(523)
        + write_u16(90)
        + bcd(timestamp)
    )


def parse_event_time(value: str) -> dt.datetime:
    parsed = dt.datetime.fromisoformat(value.replace("Z", "+00:00"))
    if parsed.tzinfo is not None:
        parsed = parsed.astimezone(dt.timezone.utc).replace(tzinfo=None)
    return parsed


def send_frame(host: str, port: int, frame: bytes, timeout: float, read_response: bool, linger_ms: int = 0) -> bytes:
    with socket.create_connection((host, port), timeout=timeout) as sock:
        sock.settimeout(timeout)
        sock.sendall(frame)
        if not read_response:
            if linger_ms > 0:
                time.sleep(linger_ms / 1000)
            return b""
        try:
            return sock.recv(1024)
        except socket.timeout:
            return b""


def send_register(host: str, port: int, phone: str, serial: int, timeout: float) -> bytes:
    device_id = "D" + phone[-6:]
    plate = "B" + phone[-5:]
    return send_frame(
        host,
        port,
        build_frame(0x0100, phone, serial, build_register_body(device_id, plate)),
        timeout,
        True,
    )


def send_location(host: str,
                  port: int,
                  phone: str,
                  serial: int,
                  event_time: dt.datetime,
                  timeout: float,
                  linger_ms: int) -> bytes:
    return send_frame(
        host,
        port,
        build_frame(0x0200, phone, serial, build_location_body(event_time, serial)),
        timeout,
        False,
        linger_ms,
    )


def recv_optional(sock: socket.socket, timeout: float) -> bytes:
    previous_timeout = sock.gettimeout()
    sock.settimeout(timeout)
    try:
        return sock.recv(1024)
    except socket.timeout:
        return b""
    finally:
        sock.settimeout(previous_timeout)


def send_phone_sequence(
    host: str,
    port: int,
    phone: str,
    phone_index: int,
    frames: int,
    event_time: dt.datetime,
    timeout: float,
    rate: float,
    linger_ms: int,
    post_register_ms: int,
) -> tuple[bytes, int]:
    with socket.create_connection((host, port), timeout=timeout) as sock:
        sock.settimeout(timeout)
        device_id = "D" + phone[-6:]
        plate = "B" + phone[-5:]
        sock.sendall(build_frame(0x0100, phone, 100 + phone_index, build_register_body(device_id, plate)))
        ack = recv_optional(sock, timeout)
        if post_register_ms > 0:
            time.sleep(post_register_ms / 1000)
        sent = 0
        for frame_index in range(frames):
            serial = 1000 + phone_index * frames + frame_index
            sock.sendall(build_frame(
                0x0200,
                phone,
                serial,
                build_location_body(event_time + dt.timedelta(seconds=frame_index), serial),
            ))
            sent += 1
            if rate > 0:
                time.sleep(1 / rate)
        if linger_ms > 0:
            time.sleep(linger_ms / 1000)
        return ack, sent


def query_json(base_url: str, path: str, params: dict[str, Any], timeout: float) -> dict[str, Any]:
    url = base_url.rstrip("/") + path + "?" + urllib.parse.urlencode(params)
    request = urllib.request.Request(url, headers={"Accept": "application/json"})
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def location_params(phone: str, date_from: str, date_to: str, limit: int) -> dict[str, Any]:
    return {
        "phone": phone,
        "dateFrom": date_from,
        "dateTo": date_to,
        "limit": limit,
        "order": "ASC",
    }


def next_page_params(params: dict[str, Any], page: dict[str, Any]) -> dict[str, Any]:
    cursor = page.get("nextCursor") or {}
    next_params = dict(params)
    if cursor:
        next_params["cursorTs"] = cursor["ts"]
        next_params["cursorId"] = cursor["id"]
    return next_params


def archive_path(root: pathlib.Path, uri: str) -> pathlib.Path:
    prefix = "archive://"
    if not uri.startswith(prefix):
        raise ValueError(f"unsupported archive uri: {uri}")
    return root / uri[len(prefix):]


def wait_for_locations(
    history_base_url: str,
    phone: str,
    date_from: str,
    date_to: str,
    minimum: int,
    timeout_seconds: float,
    request_timeout: float,
) -> dict[str, Any]:
    deadline = time.monotonic() + timeout_seconds
    params = location_params(phone, date_from, date_to, max(minimum, 10))
    last_page: dict[str, Any] = {}
    while time.monotonic() < deadline:
        last_page = query_json(history_base_url, "/api/event-history/jt808/locations", params, request_timeout)
        if len(last_page.get("items") or []) >= minimum:
            return last_page
        time.sleep(1)
    return last_page


def verify_pagination(history_base_url: str, phone: str, date_from: str, date_to: str, request_timeout: float) -> bool:
    first_params = location_params(phone, date_from, date_to, 1)
    first = query_json(history_base_url, "/api/event-history/jt808/locations", first_params, request_timeout)
    if not first.get("items") or not first.get("nextCursor"):
        return False
    second = query_json(
        history_base_url,
        "/api/event-history/jt808/locations",
        next_page_params(first_params, first),
        request_timeout,
    )
    return bool(second.get("items"))


def generated_phone(start_phone: str, offset: int) -> str:
    return str(int(start_phone) + offset)


def run(args: argparse.Namespace) -> dict[str, Any]:
    event_time = parse_event_time(args.event_time)
    date_from = (event_time - dt.timedelta(minutes=5)).isoformat()
    date_to = (event_time + dt.timedelta(minutes=args.frames + 5)).isoformat()
    phones = [generated_phone(args.start_phone, index) for index in range(args.phones)]
    sent_locations = 0
    sent_registers = 0
    start = time.monotonic()

    for phone_index, phone in enumerate(phones):
        if args.connection_mode == "session":
            ack, sent = send_phone_sequence(
                args.tcp_host,
                args.tcp_port,
                phone,
                phone_index,
                args.frames,
                event_time,
                args.tcp_timeout,
                args.rate,
                args.linger_ms,
                args.post_register_ms,
            )
        else:
            ack = send_register(args.tcp_host, args.tcp_port, phone, 100 + phone_index, args.tcp_timeout)
            if args.post_register_ms > 0:
                time.sleep(args.post_register_ms / 1000)
            sent = 0
            for frame_index in range(args.frames):
                serial = 1000 + phone_index * args.frames + frame_index
                send_location(
                    args.tcp_host,
                    args.tcp_port,
                    phone,
                    serial,
                    event_time + dt.timedelta(seconds=frame_index),
                    args.tcp_timeout,
                    args.linger_ms,
                )
                sent += 1
                if args.rate > 0:
                    time.sleep(1 / args.rate)
        sent_registers += 1
        if args.require_register_ack and not ack:
            raise RuntimeError(f"register ack timeout for {phone}")
        sent_locations += sent

    elapsed = max(time.monotonic() - start, 0.001)
    first_phone = phones[0]
    page = wait_for_locations(
        args.history_base_url,
        first_phone,
        date_from,
        date_to,
        min(args.frames, args.expect_history_count),
        args.history_timeout,
        args.http_timeout,
    )
    items = page.get("items") or []
    if len(items) < min(args.frames, args.expect_history_count):
        raise RuntimeError(f"history rows not visible for {first_phone}: got {len(items)}")

    archive_checked = 0
    if args.archive_root:
        root = pathlib.Path(args.archive_root)
        for item in items[: args.archive_check_limit]:
            path = archive_path(root, item["rawUri"])
            if not path.exists():
                raise RuntimeError(f"raw archive missing: {path}")
            archive_checked += 1

    pagination_ok = None
    if args.verify_pagination and args.frames > 1:
        pagination_ok = verify_pagination(args.history_base_url, first_phone, date_from, date_to, args.http_timeout)
        if not pagination_ok:
            raise RuntimeError("cursor pagination verification failed")

    return {
        "tcpHost": args.tcp_host,
        "tcpPort": args.tcp_port,
        "connectionMode": args.connection_mode,
        "phones": phones,
        "sentRegisters": sent_registers,
        "sentLocations": sent_locations,
        "sendElapsedSeconds": round(elapsed, 3),
        "sendLocationsPerSecond": round(sent_locations / elapsed, 2),
        "historyRowsForFirstPhone": len(items),
        "archiveChecked": archive_checked,
        "paginationOk": pagination_ok,
        "firstFrameId": items[0].get("frameId") if items else None,
        "firstRawUri": items[0].get("rawUri") if items else None,
    }


def parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--tcp-host", default="127.0.0.1")
    p.add_argument("--tcp-port", type=int, default=808)
    p.add_argument("--history-base-url", default="http://127.0.0.1:20200")
    p.add_argument("--connection-mode", choices=["per-frame", "session"], default="per-frame")
    p.add_argument("--start-phone", default="13079962000")
    p.add_argument("--phones", type=int, default=1)
    p.add_argument("--frames", type=int, default=3, help="location frames per phone")
    p.add_argument("--rate", type=float, default=0, help="location send rate per second; 0 means unlimited")
    p.add_argument("--linger-ms", type=int, default=2500, help="time to keep each terminal connection open after writes")
    p.add_argument("--post-register-ms", type=int, default=1000, help="pause after register ack before location frames")
    p.add_argument("--event-time", default="2024-01-02T03:04:05")
    p.add_argument("--expect-history-count", type=int, default=1)
    p.add_argument("--verify-pagination", action="store_true")
    p.add_argument("--archive-root", default="")
    p.add_argument("--archive-check-limit", type=int, default=3)
    p.add_argument("--require-register-ack", action="store_true", default=True)
    p.add_argument("--tcp-timeout", type=float, default=2)
    p.add_argument("--http-timeout", type=float, default=5)
    p.add_argument("--history-timeout", type=float, default=30)
    return p


def main() -> None:
    args = parser().parse_args()
    summary = run(args)
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()