#!/usr/bin/env python3
"""Small stdlib RENKEI Processor TCP/JSON client for driver authors."""

from __future__ import annotations

import argparse
import json
import math
import os
import select
import socket
import ssl
import sys
import time
from decimal import Decimal, InvalidOperation
from typing import Any

PROTOCOL_VERSION = "1.0"
MAX_FRAME_BYTES = 1024 * 1024
MAX_REQUEST_ID_CODE_POINTS = 64


def parse_target_percent(value: str) -> float:
    try:
        parsed = Decimal(value)
    except InvalidOperation as exc:
        raise argparse.ArgumentTypeError("must be a number from 0 to 100") from exc
    if not parsed.is_finite() or parsed < 0 or parsed > 100:
        raise argparse.ArgumentTypeError("must be a number from 0 to 100")
    if parsed != parsed.quantize(Decimal("0.01")):
        raise argparse.ArgumentTypeError("must have no more than two decimal places")
    return float(parsed)


def parse_ping_interval(value: str) -> float:
    try:
        parsed = float(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("must be a number greater than 0 and less than 90") from exc
    if not math.isfinite(parsed) or parsed <= 0 or parsed >= 90:
        raise argparse.ArgumentTypeError("must be a number greater than 0 and less than 90")
    return parsed


def parse_request_id(value: str) -> str:
    if not 1 <= len(value) <= MAX_REQUEST_ID_CODE_POINTS:
        raise argparse.ArgumentTypeError("must contain 1 to 64 Unicode code points")
    return value


def parse_motor_target(value: str) -> dict[str, Any]:
    motor_id, separator, target = value.partition("=")
    if not separator or not motor_id or not target:
        raise argparse.ArgumentTypeError("must use MOTOR_ID=TARGET_PERCENT")
    return {
        "motor_id": motor_id,
        "target_percent": parse_target_percent(target),
    }


class RenkeiClient:
    def __init__(
        self,
        host: str,
        port: int,
        token: str,
        use_tls: bool = True,
        insecure: bool = False,
        cafile: str | None = None,
        timeout: float = 10.0,
    ) -> None:
        self.host = host
        self.port = port
        self.token = token
        self.use_tls = use_tls
        self.insecure = insecure
        self.cafile = cafile
        self.timeout = timeout
        self.sock: socket.socket | ssl.SSLSocket | None = None
        self.receive_buffer = bytearray()

    def __enter__(self) -> "RenkeiClient":
        self.connect()
        return self

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

    def connect(self) -> None:
        raw = socket.create_connection((self.host, self.port), timeout=self.timeout)
        if self.use_tls:
            if self.insecure:
                context = ssl._create_unverified_context()
            else:
                context = ssl.create_default_context(cafile=self.cafile)
            self.sock = context.wrap_socket(raw, server_hostname=self.host)
        else:
            self.sock = raw
        self.receive_buffer.clear()

    def close(self) -> None:
        if self.sock is not None:
            self.sock.close()
            self.sock = None
        self.receive_buffer.clear()

    def send(self, message: dict[str, Any]) -> None:
        if self.sock is None:
            raise RuntimeError("client is not connected")
        payload = json.dumps(message, separators=(",", ":")).encode("utf-8")
        if len(payload) > MAX_FRAME_BYTES:
            raise ValueError("outbound frame exceeds 1 MiB")
        self.sock.sendall(payload + b"\n")

    def receive(self) -> dict[str, Any]:
        if self.sock is None:
            raise RuntimeError("client is not connected")
        while b"\n" not in self.receive_buffer:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise EOFError("connection closed")
            self.receive_buffer.extend(chunk)
            if len(self.receive_buffer) > MAX_FRAME_BYTES and b"\n" not in self.receive_buffer:
                raise ValueError("received frame exceeds 1 MiB")

        line, _, remainder = self.receive_buffer.partition(b"\n")
        self.receive_buffer = bytearray(remainder)
        if len(line) > MAX_FRAME_BYTES:
            raise ValueError("received frame exceeds 1 MiB")
        return json.loads(line.decode("utf-8"))

    def message_available(self, timeout: float) -> bool:
        if self.sock is None:
            raise RuntimeError("client is not connected")
        if b"\n" in self.receive_buffer:
            return True
        if isinstance(self.sock, ssl.SSLSocket) and self.sock.pending() > 0:
            return True
        readable, _, _ = select.select([self.sock], [], [], timeout)
        return bool(readable)

    def request(self, message: dict[str, Any], expected_type: str) -> dict[str, Any]:
        expected_request_id = message.get("request_id")
        self.send(message)
        while True:
            response = self.receive()
            if expected_request_id is not None:
                if response.get("request_id") == expected_request_id:
                    return response
            elif "error_code" in response or response.get("type") == expected_type:
                return response
            print(json.dumps(response, sort_keys=True), file=sys.stderr)

    def hello(self) -> dict[str, Any]:
        return self.request(
            {
                "type": "hello",
                "protocol_version": PROTOCOL_VERSION,
                "auth": {
                    "token": self.token,
                },
            },
            "hello",
        )

    def ping(self) -> dict[str, Any]:
        return self.request(
            {
                "type": "ping",
            },
            "pong",
        )

    def query_state(self, scope: str = "all") -> dict[str, Any]:
        return self.request(
            {
                "type": "query_state",
                "scope": scope,
            },
            "query_state",
        )

    def query_baseline(self) -> list[dict[str, Any]]:
        response = self.query_state("all")
        if response.get("error_code") != "response_too_large":
            return [response]

        responses: list[dict[str, Any]] = []
        for scope in ("motors", "groups"):
            scoped = self.query_state(scope)
            responses.append(scoped)
            if "error_code" in scoped:
                break
        return responses

    def subscribe(self) -> dict[str, Any]:
        return self.request(
            {
                "type": "subscribe",
            },
            "subscribe",
        )

    def command(self, payload: dict[str, Any]) -> dict[str, Any]:
        payload.setdefault("type", "command")
        return self.request(payload, "command_response")

    def group_command(self, payload: dict[str, Any]) -> dict[str, Any]:
        payload.setdefault("type", "group_command")
        return self.request(payload, "command_response")

    def multi_motor_move(self, payload: dict[str, Any]) -> dict[str, Any]:
        payload.setdefault("type", "multi_motor_move")
        return self.request(payload, "command_response")


def print_json(payload: dict[str, Any]) -> None:
    print(json.dumps(payload, indent=2, sort_keys=True))


def establish_baseline(client: RenkeiClient) -> None:
    responses = client.query_baseline()
    for response in responses:
        print_json(response)
    if any("error_code" in response for response in responses):
        raise RuntimeError("could not establish a complete query_state baseline")


def add_connection_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--host", required=True, help="Processor IP address or hostname.")
    parser.add_argument("--port", type=int, help="Processor TCP port. Defaults to 17101 for TLS or 17102 for plaintext.")
    parser.add_argument("--token", default=os.environ.get("RENKEI_TOKEN"), help="Integration token. Defaults to RENKEI_TOKEN.")
    parser.add_argument("--plaintext", action="store_true", help="Use plaintext TCP instead of TLS.")
    parser.add_argument("--insecure", action="store_true", help="Disable TLS certificate verification for lab systems.")
    parser.add_argument("--cafile", help="Custom CA bundle for TLS certificate verification.")
    parser.add_argument("--timeout", type=float, default=10.0, help="Socket timeout in seconds.")


def connect_from_args(args: argparse.Namespace) -> RenkeiClient:
    if not args.token:
        raise SystemExit("--token or RENKEI_TOKEN is required")
    port = args.port if args.port is not None else (17102 if args.plaintext else 17101)
    return RenkeiClient(
        host=args.host,
        port=port,
        token=args.token,
        use_tls=not args.plaintext,
        insecure=args.insecure,
        cafile=args.cafile,
        timeout=args.timeout,
    )


def run_ping(args: argparse.Namespace) -> None:
    with connect_from_args(args) as client:
        print_json(client.hello())
        print_json(client.ping())


def run_query(args: argparse.Namespace) -> None:
    with connect_from_args(args) as client:
        print_json(client.hello())
        responses = client.query_baseline() if args.scope == "all" else [client.query_state(args.scope)]
        for response in responses:
            print_json(response)


def run_listen(args: argparse.Namespace) -> None:
    with connect_from_args(args) as client:
        print_json(client.hello())
        establish_baseline(client)
        print_json(client.subscribe())
        next_ping = time.monotonic() + args.ping_interval
        while True:
            wait_seconds = max(0.0, next_ping - time.monotonic())
            if client.message_available(wait_seconds):
                print_json(client.receive())
            if time.monotonic() >= next_ping:
                client.send(
                    {
                        "type": "ping",
                    }
                )
                next_ping = time.monotonic() + args.ping_interval


def run_command(args: argparse.Namespace) -> None:
    payload: dict[str, Any] = {
        "motor_id": args.motor_id,
        "op": args.op,
    }
    if args.op == "move":
        if args.target_percent is None:
            raise SystemExit("--target-percent is required for move")
        payload["target_percent"] = args.target_percent
    elif args.target_percent is not None:
        raise SystemExit("--target-percent is valid only for move")
    if args.request_id is not None:
        payload["request_id"] = args.request_id

    with connect_from_args(args) as client:
        print_json(client.hello())
        establish_baseline(client)
        print_json(client.subscribe())
        print_json(client.command(payload))


def run_group_command(args: argparse.Namespace) -> None:
    payload: dict[str, Any] = {
        "group_id": args.group_id,
        "op": args.op,
    }
    if args.target_percent is not None:
        payload["target_percent"] = args.target_percent
    if args.op == "move" and args.target_percent is None:
        raise SystemExit("--target-percent is required for group move")
    if args.op == "stop" and args.target_percent is not None:
        raise SystemExit("--target-percent is valid only for group move")
    if args.request_id is not None:
        payload["request_id"] = args.request_id

    with connect_from_args(args) as client:
        print_json(client.hello())
        establish_baseline(client)
        print_json(client.subscribe())
        print_json(client.group_command(payload))


def run_multi_motor_move(args: argparse.Namespace) -> None:
    motor_ids = [target["motor_id"] for target in args.targets]
    if len(set(motor_ids)) != len(motor_ids):
        raise SystemExit("--target motor IDs must be unique")

    payload: dict[str, Any] = {"targets": args.targets}
    if args.request_id is not None:
        payload["request_id"] = args.request_id

    with connect_from_args(args) as client:
        print_json(client.hello())
        establish_baseline(client)
        print_json(client.subscribe())
        print_json(client.multi_motor_move(payload))


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="RENKEI Processor integration API reference client.")
    add_connection_args(parser)
    subparsers = parser.add_subparsers(dest="command_name", required=True)

    ping_parser = subparsers.add_parser("ping", help="Authenticate and send ping.")
    ping_parser.set_defaults(func=run_ping)

    query_parser = subparsers.add_parser("query", help="Authenticate and query baseline state.")
    query_parser.add_argument(
        "--scope",
        choices=["all", "motors", "groups"],
        default="all",
        help="Baseline scope (default: all; automatically falls back to motors then groups if oversized).",
    )
    query_parser.set_defaults(func=run_query)

    listen_parser = subparsers.add_parser("listen", help="Authenticate, subscribe, and print events.")
    listen_parser.add_argument(
        "--ping-interval",
        type=parse_ping_interval,
        default=30.0,
        help="Keepalive interval in seconds. Must remain below the processor's 90-second inbound deadline (default: 30).",
    )
    listen_parser.set_defaults(func=run_listen)

    command_parser = subparsers.add_parser("command", help="Send a single-motor command after the full lifecycle.")
    command_parser.add_argument("--motor-id", required=True)
    command_parser.add_argument(
        "--op",
        choices=["move", "open", "close", "stop", "cycle_motion"],
        required=True,
    )
    command_parser.add_argument("--target-percent", type=parse_target_percent)
    command_parser.add_argument(
        "--request-id",
        type=parse_request_id,
        help="Optional 1-to-64-code-point retry key and response correlation ID.",
    )
    command_parser.set_defaults(func=run_command)

    group_parser = subparsers.add_parser("group-command", help="Send a group command after the full lifecycle.")
    group_parser.add_argument("--group-id", required=True)
    group_parser.add_argument("--op", choices=["move", "stop"], required=True)
    group_parser.add_argument("--target-percent", type=parse_target_percent)
    group_parser.add_argument(
        "--request-id",
        type=parse_request_id,
        help="Optional 1-to-64-code-point retry key and response correlation ID.",
    )
    group_parser.set_defaults(func=run_group_command)

    multi_parser = subparsers.add_parser(
        "multi-motor-move",
        help="Send synchronized independent percentage targets after the full lifecycle.",
    )
    multi_parser.add_argument(
        "--target",
        dest="targets",
        action="append",
        type=parse_motor_target,
        required=True,
        metavar="MOTOR_ID=PERCENT",
        help="Motor target; repeat for each unique motor.",
    )
    multi_parser.add_argument(
        "--request-id",
        type=parse_request_id,
        help="Optional 1-to-64-code-point retry key and response correlation ID.",
    )
    multi_parser.set_defaults(func=run_multi_motor_move)

    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    try:
        args.func(args)
    except KeyboardInterrupt:
        return 130
    return 0


if __name__ == "__main__":
    sys.exit(main())