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


DOIP_VERSION = 0x02
DOIP_INVERSE_VERSION = 0xFF - DOIP_VERSION
DOIP_PORT_TCP = 13400
DOIP_PORT_UDP = 13400
DOIP_PAYLOAD_VEHICLE_IDENT_REQUEST = 0x0001
DOIP_PAYLOAD_VEHICLE_IDENT_REQUEST_EID = 0x0002
DOIP_PAYLOAD_VEHICLE_IDENT_REQUEST_VIN = 0x0003
DOIP_PAYLOAD_VEHICLE_IDENT_RESPONSE = 0x0004
DOIP_PAYLOAD_ROUTING_ACTIVATION_REQUEST = 0x0005
DOIP_PAYLOAD_ROUTING_ACTIVATION_RESPONSE = 0x0006
DOIP_PAYLOAD_DIAGNOSTIC_MESSAGE = 0x8001
DOIP_PAYLOAD_DIAGNOSTIC_ACK = 0x8002
DOIP_PAYLOAD_DIAGNOSTIC_NACK = 0x8003

DOIP_TESTER_ADDRESS = 0x0E80
DOIP_ROUTING_SUCCESS = 0x10
DOIP_ROUTING_INVALID_SOURCE = 0x00
DOIP_BROADCAST_ADDRESS = 0xE400
DOIP_BROADCAST_UDP = "255.255.255.255"
DOIP_BROADCAST_LINK_LOCAL_UDP = "169.254.255.255"


def _parse_tester_address(value: str) -> int:
    try:
        address = int(value, 0)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("Tester address must be a number (e.g. 0x0E80)") from exc
    if not (0 <= address <= 0xFFFF):
        raise argparse.ArgumentTypeError("Tester address must be between 0x0000 and 0xFFFF")
    return address


@dataclass
class DoIPMessage:
    version: int
    inverse_version: int
    payload_type: int
    payload: bytes

    @property
    def bytes(self) -> bytes:
        header = _build_header(self.payload_type, len(self.payload))
        return header + self.payload

    @staticmethod
    def from_socket(client: socket.socket) -> Optional["DoIPMessage"]:
        header = _read_exact(client, 8)
        if header is None:
            return None
        version = header[0]
        inverse_version = header[1]
        payload_type = (header[2] << 8) | header[3]
        payload_length = int.from_bytes(header[4:8], "big")
        payload = _read_exact(client, payload_length) if payload_length else b""
        if payload_length and payload is None:
            return None
        return DoIPMessage(
            version=version,
            inverse_version=inverse_version,
            payload_type=payload_type,
            payload=payload,
        )


class DoIPAdapter:
    name = "DoIP"
    default_url = "doip://0.0.0.0:13400"
    server_help = "DoIP URL (doip://<host>:<port>/<iface>), default broadcasts to discover"
    header_width = 4
    header_mask = 0xFFFF
    default_header = "df"
    discovery_failure_message = "Could not find a DoIP server."

    def __init__(self) -> None:
        self.tester_address = DOIP_TESTER_ADDRESS

    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument(
            "--tester-address",
            default=DOIP_TESTER_ADDRESS,
            type=_parse_tester_address,
            help="Tester logical address (hex or decimal, default: 0x0E80)",
        )

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

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

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

    def format_discovery_message(self, result: DiscoveryResult) -> str:
        return f"Found DoIP server @ {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 DoIP server @ {interface} via {host}:{port}"

    def connect(self, client: socket.socket, host: str, port: int, interface: str) -> None:
        try:
            response_code = _send_routing_activation(client, self.tester_address)
        except Exception as exc:
            raise RuntimeError(f"Routing activation failed: {exc}")
        print(f"Routing activation accepted (code 0x{response_code:02X})")

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


def _read_exact(client: socket.socket, length: int) -> Optional[bytes]:
    if length <= 0:
        return b""
    buffer = bytearray()
    while len(buffer) < length:
        try:
            chunk = client.recv(length - len(buffer))
        except socket.timeout:
            return None
        if not chunk:
            raise ConnectionError("DoIP peer closed the connection")
        buffer.extend(chunk)
    return bytes(buffer)


def _build_header(payload_type: int, payload_length: int) -> bytes:
    return bytes(
        [
            DOIP_VERSION,
            DOIP_INVERSE_VERSION,
            (payload_type >> 8) & 0xFF,
            payload_type & 0xFF,
        ]
    ) + int(payload_length).to_bytes(4, "big")


def _build_routing_activation_request(source_address: int) -> bytes:
    payload = bytearray()
    payload.extend(int(source_address).to_bytes(2, "big"))
    payload.append(0x00)
    payload.extend([0x00] * 8)
    header = _build_header(DOIP_PAYLOAD_ROUTING_ACTIVATION_REQUEST, len(payload))
    return header + bytes(payload)


def _build_diagnostic_message(source: int, target: int, payload: bytes) -> bytes:
    diag_payload = bytearray()
    diag_payload.extend(int(source).to_bytes(2, "big"))
    diag_payload.extend(int(target).to_bytes(2, "big"))
    diag_payload.extend(payload)
    header = _build_header(DOIP_PAYLOAD_DIAGNOSTIC_MESSAGE, len(diag_payload))
    return header + bytes(diag_payload)


def _extract_vin(payload: bytes) -> str:
    if len(payload) < 17:
        return "unknown"
    vin_bytes = payload[:17]
    vin = "".join(chr(b) for b in vin_bytes if 32 <= b < 127 and b != 0x00)
    vin = "".join(ch for ch in vin if ch.isalnum())
    return vin or "unknown"


def _discover_server(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)
    request = _build_header(DOIP_PAYLOAD_VEHICLE_IDENT_REQUEST, 0)
    try:
        for _ in range(retries):
            for broadcast in (DOIP_BROADCAST_UDP, DOIP_BROADCAST_LINK_LOCAL_UDP):
                try:
                    udp.sendto(request, (broadcast, DOIP_PORT_UDP))
                except OSError:
                    continue
            try:
                data, addr = udp.recvfrom(256)
            except socket.timeout:
                continue
            if len(data) < 8:
                continue
            payload_type = (data[2] << 8) | data[3]
            if payload_type != DOIP_PAYLOAD_VEHICLE_IDENT_RESPONSE:
                continue
            payload_length = int.from_bytes(data[4:8], "big")
            if len(data) < 8 + payload_length:
                continue
            payload = data[8 : 8 + payload_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_routing_activation(
    client: socket.socket,
    tester_address: int,
    retries: int = 2,
    timeout: float = 5.0,
) -> int:
    request = _build_routing_activation_request(tester_address)
    for attempt in range(retries):
        client.sendall(request)
        deadline = time.time() + timeout
        while True:
            remaining = deadline - time.time()
            if remaining <= 0:
                break
            client.settimeout(max(0.1, remaining))
            response = DoIPMessage.from_socket(client)
            if response is None:
                continue
            if response.payload_type == DOIP_PAYLOAD_ROUTING_ACTIVATION_RESPONSE:
                if len(response.payload) < 5:
                    raise RuntimeError("Routing activation response too short")
                response_code = response.payload[4]
                if response_code == DOIP_ROUTING_INVALID_SOURCE:
                    raise RuntimeError("Invalid source address")
                if response_code != DOIP_ROUTING_SUCCESS:
                    raise RuntimeError(f"Routing activation denied (response code 0x{response_code:02X})")
                return response_code
        if attempt < retries - 1:
            continue
    raise RuntimeError("Did not receive routing activation response")


def _send_request(
    client: socket.socket,
    tester_address: int,
    target: int,
    payload: bytes,
) -> list[tuple[int, bytes]]:
    message = _build_diagnostic_message(tester_address, target, payload)
    client.sendall(message)
    responses: list[tuple[int, bytes]] = []
    is_broadcast = target == DOIP_BROADCAST_ADDRESS
    deadline = time.time() + 1.0
    while True:
        remaining = deadline - time.time()
        if remaining <= 0:
            break
        client.settimeout(remaining)
        response = DoIPMessage.from_socket(client)
        if response is None:
            break
        if response.payload_type == DOIP_PAYLOAD_DIAGNOSTIC_ACK:
            continue
        if response.payload_type == DOIP_PAYLOAD_DIAGNOSTIC_NACK:
            print("Error: DoIP diagnostic message NACK")
            break
        if response.payload_type != DOIP_PAYLOAD_DIAGNOSTIC_MESSAGE:
            continue
        if len(response.payload) < 4:
            continue
        source = int.from_bytes(response.payload[0:2], "big")
        uds_bytes = response.payload[4:]
        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((source, uds_bytes))
        if not is_broadcast and not is_pending:
            break
    return responses


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


if __name__ == "__main__":
    main()
