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

from __future__ import annotations

import argparse
import concurrent.futures
import datetime as dt
import json
import os
import pathlib
import socket
import time
import urllib.parse
import urllib.request
from dataclasses import dataclass
from typing import Any

try:
    from tools import tdengine_smoke
except ModuleNotFoundError:
    import tdengine_smoke


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


@dataclass(frozen=True)
class SendWorkloadResult:
    sent_registers: int
    sent_locations: int
    elapsed_seconds: float


def send_one_phone(args: argparse.Namespace,
                   phone: str,
                   phone_index: int,
                   event_time: dt.datetime) -> tuple[bytes, int]:
    if args.connection_mode == "session":
        return 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,
        )

    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)
    return ack, sent


def send_phone_workload(args: argparse.Namespace,
                        phones: list[str],
                        event_time: dt.datetime) -> SendWorkloadResult:
    sent_locations = 0
    start = time.monotonic()
    workers = max(1, int(getattr(args, "workers", 1)))

    def run_one(index_and_phone: tuple[int, str]) -> tuple[bytes, int]:
        phone_index, phone = index_and_phone
        return send_one_phone(args, phone, phone_index, event_time)

    indexed_phones = list(enumerate(phones))
    if workers == 1 or len(indexed_phones) <= 1:
        results = [run_one(item) for item in indexed_phones]
    else:
        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
            results = list(executor.map(run_one, indexed_phones))

    for phone, (ack, sent) in zip(phones, results):
        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)
    return SendWorkloadResult(
        sent_registers=len(phones),
        sent_locations=sent_locations,
        elapsed_seconds=elapsed,
    )


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 raw_frame_count_sql(phone: str, items: list[dict[str, Any]]) -> str:
    sqls = raw_frame_count_sqls(phone, items)
    if not sqls:
        raise ValueError("no raw uri found in jt808 location items")
    return sqls[0]


def raw_frame_count_sqls(phone: str, items: list[dict[str, Any]]) -> list[str]:
    seen: set[str] = set()
    sqls: list[str] = []
    for item in items:
        raw_uri = item.get("rawUri") or ""
        if not raw_uri or raw_uri in seen:
            continue
        seen.add(raw_uri)
        sqls.append(tdengine_smoke.raw_frame_count_sql(
            protocol="JT808",
            vehicle_key=item.get("vehicleKey") or ("jt808:" + phone),
            phone=phone,
            raw_uri=raw_uri,
        ))
    return sqls


def verify_tdengine_raw_frames(
    rest_url: str,
    username: str,
    password: str,
    sqls: list[str],
    timeout_seconds: float,
    request_timeout: float,
) -> int:
    if not sqls:
        return 0
    verified = 0
    for sql in sqls:
        count = wait_for_tdengine_raw_frames(
            rest_url, username, password, sql, 1, timeout_seconds, request_timeout)
        if count < 1:
            raise RuntimeError(f"tdengine raw_frames row not visible for query: {sql}")
        verified += 1
    return verified


def wait_for_tdengine_raw_frames(
    rest_url: str,
    username: str,
    password: str,
    sql: str,
    minimum: int,
    timeout_seconds: float,
    request_timeout: float,
) -> int:
    deadline = time.monotonic() + timeout_seconds
    last_count = 0
    while time.monotonic() < deadline:
        last_count = tdengine_smoke.query_count(rest_url, username, password, sql, request_timeout)
        if last_count >= minimum:
            return last_count
        time.sleep(1)
    return last_count


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 expected_history_count(args: argparse.Namespace) -> int:
    minimum = min(args.frames, args.expect_history_count)
    if args.verify_pagination and args.frames > 1:
        minimum = max(minimum, 2)
    return minimum


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)]
    workload = send_phone_workload(args, phones, event_time)
    first_phone = phones[0]
    page = wait_for_locations(
        args.history_base_url,
        first_phone,
        date_from,
        date_to,
        expected_history_count(args),
        args.history_timeout,
        args.http_timeout,
    )
    items = page.get("items") or []
    if len(items) < expected_history_count(args):
        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

    tdengine_raw_frames = None
    if args.tdengine_rest_url:
        tdengine_raw_frames = verify_tdengine_raw_frames(
            args.tdengine_rest_url,
            args.tdengine_username,
            args.tdengine_password,
            raw_frame_count_sqls(first_phone, items),
            args.history_timeout,
            args.http_timeout,
        )

    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,
        "workers": max(1, int(args.workers)),
        "phones": phones,
        "sentRegisters": workload.sent_registers,
        "sentLocations": workload.sent_locations,
        "sendElapsedSeconds": round(workload.elapsed_seconds, 3),
        "sendLocationsPerSecond": round(workload.sent_locations / workload.elapsed_seconds, 2),
        "historyRowsForFirstPhone": len(items),
        "tdengineRawFrames": tdengine_raw_frames,
        "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("--workers", type=int, default=1, help="parallel phone sessions; 1 keeps legacy serial behavior")
    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("--tdengine-rest-url", default=os.environ.get("TDENGINE_REST_URL", ""))
    p.add_argument("--tdengine-username", default=os.environ.get("TDENGINE_USERNAME", "root"))
    p.add_argument("--tdengine-password", default=os.environ.get("TDENGINE_PASSWORD", "taosdata"))
    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()