from __future__ import annotations

import argparse
import socket
import time
from dataclasses import dataclass
from typing import Optional
from urllib.parse import urlparse

from pymuclient.common import DiscoveryResult, resolve_local_ip, run_cli


HSFZ_PORT_TCP = 6801
HSFZ_PORT_UDP = 6811
HSFZ_TYPE_DIAGNOSTIC_REQ_RES = 0x0001
HSFZ_TYPE_ACKNOWLEDGE_TRANSFER = 0x0002
HSFZ_REQUEST_VEHICLE_ANNOUNCEMENT = bytes([0x00, 0x00, 0x00, 0x00, 0x00, 0x11])
HSFZ_BROADCAST_ARBITRATION = 0xDF
HSFZ_TESTER_ADDRESS = 0xF4
HSFZ_BROADCAST_UDP = "255.255.255.255"
HSFZ_BROADCAST_LINK_LOCAL_UDP = "169.254.255.255"


@dataclass
class HSFZMessage:
    msg_type: int
    source: int
    destination: int
    payload: bytes

    @property
    def length(self) -> int:
        return 2 + len(self.payload)

    @property
    def bytes(self) -> bytes:
        length_bytes = self.length.to_bytes(4, "big")
        type_bytes = self.msg_type.to_bytes(2, "big")
        return length_bytes + type_bytes + bytes([self.source, self.destination]) + self.payload

    @staticmethod
    def from_socket(client: socket.socket) -> Optional["HSFZMessage"]:
        header = _read_exact(client, 6)
        if not header:
            return None
        length = int.from_bytes(header[:4], "big")
        msg_type = int.from_bytes(header[4:6], "big")
        if length == 0:
            return None
        address_bytes = _read_exact(client, 2)
        if not address_bytes:
            return None
        source, destination = address_bytes[0], address_bytes[1]
        payload_length = max(length - 2, 0)
        payload = _read_exact(client, payload_length) if payload_length else b""
        if payload_length and not payload:
            return None
        return HSFZMessage(msg_type=msg_type, source=source, destination=destination, payload=payload)


class HSFZAdapter:
    name = "HSFZ"
    default_url = "enet://0.0.0.0:6801"
    server_help = "ZGW URL (enet://<host>:<port>/<iface>), default broadcasts to discover"
    header_width = 2
    header_mask = 0xFF
    default_header = "df"
    discovery_failure_message = "Could not find a ZGW."

    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument(
            "-f",
            "--force",
            action="store_true",
            help="Force entering REPL even if no ECU responds to negotiation (kept for parity).",
        )

    def handle_arguments(self, args: argparse.Namespace) -> None:
        return

    def parse_url(self, value: str) -> tuple[str, int, str]:
        if "://" not in value:
            value = f"enet://{value}"
        parsed = urlparse(value)
        if not parsed.hostname:
            raise ValueError("Invalid server URL, must be enet://<host>:[<port>]")
        host = parsed.hostname
        port = parsed.port or HSFZ_PORT_TCP
        interface = parsed.path.lstrip("/") or "en0"
        return host, port, interface

    def discover(self) -> Optional[DiscoveryResult]:
        return _discover_zgw()

    def format_discovery_message(self, result: DiscoveryResult) -> str:
        return f"Found ZGW @ {result.interface_ip} for {result.vin} via {result.server_ip}"

    def format_connect_message(self, interface: str, host: str, port: int) -> str:
        return f"Connecting to ZGW @ {interface} via {host}:{port}"

    def connect(self, client: socket.socket, host: str, port: int, interface: str) -> None:
        return

    def send_request(self, client: socket.socket, target: int, payload: bytes) -> list[tuple[int, bytes]]:
        return _send_request(client, target, payload)


def _read_exact(client: socket.socket, length: int) -> bytes:
    if length <= 0:
        return b""
    buffer = bytearray()
    while len(buffer) < length:
        try:
            chunk = client.recv(length - len(buffer))
        except socket.timeout:
            return b""
        if not chunk:
            return b""
        buffer.extend(chunk)
    return bytes(buffer)


def _extract_vin(payload: bytes) -> str:
    try:
        text = payload.decode("ascii", errors="ignore")
    except Exception:
        return "unknown"
    marker = "BMWVIN"
    idx = text.find(marker)
    if idx == -1:
        return "unknown"
    vin = text[idx + len(marker) : idx + len(marker) + 17]
    vin = "".join(ch for ch in vin if ch.isalnum())
    return vin or "unknown"


def _discover_zgw(timeout: float = 1.0, retries: int = 3) -> Optional[DiscoveryResult]:
    udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    udp.settimeout(timeout)
    try:
        for _ in range(retries):
            for broadcast in (HSFZ_BROADCAST_UDP, HSFZ_BROADCAST_LINK_LOCAL_UDP):
                try:
                    udp.sendto(HSFZ_REQUEST_VEHICLE_ANNOUNCEMENT, (broadcast, HSFZ_PORT_UDP))
                except OSError:
                    continue
            try:
                data, addr = udp.recvfrom(256)
            except socket.timeout:
                continue
            if len(data) < 6 or data[4:6] != b"\x00\x11":
                continue
            length = int.from_bytes(data[0:4], "big")
            if len(data) < 6 + length:
                continue
            payload = data[6 : 6 + length]
            vin = _extract_vin(payload)
            server_ip = addr[0]
            interface_ip = resolve_local_ip(server_ip)
            return DiscoveryResult(interface_ip=interface_ip, server_ip=server_ip, vin=vin)
    finally:
        udp.close()
    return None


def _send_request(client: socket.socket, destination: int, payload: bytes) -> list[tuple[int, bytes]]:
    source = HSFZ_TESTER_ADDRESS
    message = HSFZMessage(
        msg_type=HSFZ_TYPE_DIAGNOSTIC_REQ_RES,
        source=source,
        destination=destination,
        payload=payload,
    )
    client.sendall(message.bytes)
    responses: list[tuple[int, bytes]] = []
    is_broadcast = destination == HSFZ_BROADCAST_ARBITRATION
    deadline = time.time() + 1.0
    while True:
        remaining = deadline - time.time()
        if remaining <= 0:
            break
        client.settimeout(remaining)
        response = HSFZMessage.from_socket(client)
        if response is None:
            break
        if response.msg_type == HSFZ_TYPE_ACKNOWLEDGE_TRANSFER:
            continue
        if response.msg_type != HSFZ_TYPE_DIAGNOSTIC_REQ_RES:
            continue
        uds_bytes = response.payload
        if not uds_bytes:
            continue
        is_negative = len(uds_bytes) >= 3 and uds_bytes[0] == 0x7F
        is_pending = is_negative and uds_bytes[2] == 0x78
        if is_pending and not is_broadcast:
            deadline = time.time() + 5.0
            continue
        responses.append((response.source, uds_bytes))
        if not is_broadcast and not is_pending:
            break
    return responses


def main(argv: list[str] | None = None) -> None:
    adapter = HSFZAdapter()
    run_cli(adapter, argv, prog="hsfz-client", description="Terminal client for HSFZ servers")


if __name__ == "__main__":
    main()
