"""UDP BitTorrent tracker probe (BEP 15).

Protocol packing inspired by newTrackon (MIT).
"""

from __future__ import annotations

import asyncio
import os
import random
import socket
import struct
import time

from app.probes.types import ProbeOutcome
from app.services.ssrf import assert_connected_ip_safe, assert_safe_destination


def _transaction_id() -> int:
    return random.randint(0, 0x7FFFFFFF)


def create_connect_request() -> tuple[bytes, int]:
    connection_id = 0x41727101980
    action = 0
    tid = _transaction_id()
    buf = struct.pack("!qii", connection_id, action, tid)
    return buf, tid


def parse_connect_response(buf: bytes, sent_tid: int) -> int:
    if len(buf) < 16:
        raise RuntimeError(f"connect response too short: {len(buf)}")
    action, tid = struct.unpack_from("!ii", buf, 0)
    if tid != sent_tid:
        raise RuntimeError("transaction id mismatch on connect")
    if action == 3:
        raise RuntimeError("tracker returned UDP error on connect")
    if action != 0:
        raise RuntimeError(f"unexpected connect action: {action}")
    return struct.unpack_from("!q", buf, 8)[0]


def create_announce_request(
    connection_id: int,
    info_hash: bytes,
    peer_id: bytes,
    *,
    num_want: int = 0,
    port: int = 0,
) -> tuple[bytes, int]:
    action = 1
    tid = _transaction_id()
    buf = struct.pack("!qii", connection_id, action, tid)
    buf += struct.pack("!20s20s", info_hash, peer_id)
    buf += struct.pack("!qqq", 0, 0, 0)  # downloaded, left, uploaded
    buf += struct.pack("!i", 2)  # event = started
    buf += struct.pack("!i", 0)  # IP
    buf += struct.pack("!i", _transaction_id())  # key
    buf += struct.pack("!i", num_want)
    buf += struct.pack("!H", port)
    return buf, tid


def parse_announce_response(buf: bytes, sent_tid: int) -> tuple[int, int, int]:
    """Return (interval, leechers, seeders). Peer bytes are ignored/discarded."""
    if len(buf) < 20:
        raise RuntimeError(f"announce response too short: {len(buf)}")
    action, tid = struct.unpack_from("!ii", buf, 0)
    if tid != sent_tid:
        raise RuntimeError("transaction id mismatch on announce")
    if action == 3:
        raise RuntimeError("tracker returned UDP error on announce")
    if action != 1:
        raise RuntimeError(f"unexpected announce action: {action}")
    interval, leechers, seeders = struct.unpack_from("!iii", buf, 8)
    # Remaining bytes would be peers — intentionally not parsed into storage.
    return interval, leechers, seeders


async def probe_udp(
    host: str,
    port: int,
    *,
    timeout: float = 10.0,
    info_hash: bytes | None = None,
    peer_id: bytes | None = None,
) -> ProbeOutcome:
    info_hash = info_hash or os.urandom(20)
    peer_id = peer_id or (b"-RASTRO-" + os.urandom(12))
    if len(peer_id) != 20:
        peer_id = (peer_id + os.urandom(20))[:20]

    loop = asyncio.get_running_loop()
    t0 = time.perf_counter()
    try:
        infos = await loop.getaddrinfo(host, port, type=socket.SOCK_DGRAM)
    except OSError as exc:
        return ProbeOutcome(
            status="down",
            response_valid=False,
            error_kind="dns",
            error_detail=str(exc)[:512],
            latency_ms=(time.perf_counter() - t0) * 1000,
        )

    ips = [str(item[4][0]) for item in infos]
    try:
        assert_safe_destination(host, ips)
    except ValueError as exc:
        return ProbeOutcome(
            status="down",
            response_valid=False,
            error_kind="blocked",
            error_detail=str(exc)[:512],
            latency_ms=(time.perf_counter() - t0) * 1000,
        )

    last_error = "udp announce failed"
    for af, socktype, proto, _, sockaddr in infos:
        ip = str(sockaddr[0])
        try:
            assert_connected_ip_safe(ip)
        except ValueError as exc:
            last_error = str(exc)
            continue

        sock = socket.socket(af, socktype, proto)
        sock.setblocking(False)
        try:
            await asyncio.wait_for(loop.sock_connect(sock, sockaddr), timeout=timeout)
            connect_req, connect_tid = create_connect_request()
            await asyncio.wait_for(loop.sock_sendall(sock, connect_req), timeout=timeout)
            connect_buf = await asyncio.wait_for(loop.sock_recv(sock, 2048), timeout=timeout)
            connection_id = parse_connect_response(connect_buf, connect_tid)

            announce_req, announce_tid = create_announce_request(
                connection_id, info_hash, peer_id, num_want=0
            )
            await asyncio.wait_for(loop.sock_sendall(sock, announce_req), timeout=timeout)
            announce_buf = await asyncio.wait_for(loop.sock_recv(sock, 2048), timeout=timeout)
            interval, leechers, seeders = parse_announce_response(announce_buf, announce_tid)
            latency = (time.perf_counter() - t0) * 1000
            family = "4" if af == socket.AF_INET else "6"
            return ProbeOutcome(
                status="up",
                response_valid=True,
                latency_ms=latency,
                ip_address=ip,
                ip_family=family,
                tracker_interval_seconds=interval,
                seeders=seeders,
                leechers=leechers,
            )
        except TimeoutError:
            last_error = "udp timeout"
        except OSError as exc:
            last_error = f"udp error: {exc}"
        except RuntimeError as exc:
            last_error = str(exc)
        finally:
            sock.close()

    return ProbeOutcome(
        status="down",
        response_valid=False,
        error_kind="udp",
        error_detail=last_error[:512],
        latency_ms=(time.perf_counter() - t0) * 1000,
    )