#!/usr/bin/env python3
"""
Resolve hostname.local via mDNS multicast (pure Python, no dependencies).
Usage: mdns-resolve <hostname.local> [timeout_seconds]
Returns: IP address on stdout, or exit 1 on failure.
"""
import sys
import socket
import struct
import select
import time


def build_mdns_query(hostname: str) -> bytes:
    header = struct.pack('>HHHHHH', 0, 0, 1, 0, 0, 0)
    question = b''
    for part in hostname.rstrip('.').split('.'):
        question += bytes([len(part)]) + part.encode()
    question += b'\x00'
    question += struct.pack('>HH', 1, 1)  # Type A, Class IN
    return header + question


def parse_mdns_response(data: bytes) -> str | None:
    if len(data) < 12:
        return None
    qdcount = struct.unpack('>H', data[4:6])[0]
    ancount = struct.unpack('>H', data[6:8])[0]
    if ancount == 0:
        return None

    offset = 12
    for _ in range(qdcount):
        while offset < len(data) and data[offset] != 0:
            if data[offset] & 0xc0 == 0xc0:
                offset += 2
                break
            offset += data[offset] + 1
        else:
            offset += 1
        offset += 4

    for _ in range(ancount):
        while offset < len(data):
            if data[offset] & 0xc0 == 0xc0:
                offset += 2
                break
            elif data[offset] == 0:
                offset += 1
                break
            else:
                offset += data[offset] + 1
        if offset + 10 > len(data):
            break
        rtype, rclass, ttl, rdlength = struct.unpack('>HHIH', data[offset:offset+10])
        offset += 10
        if rtype == 1 and rdlength == 4 and offset + 4 <= len(data):
            return socket.inet_ntoa(data[offset:offset+4])
        offset += rdlength
    return None


def resolve_mdns(hostname: str, timeout: float = 2.0) -> str | None:
    if not hostname.endswith('.local'):
        hostname += '.local'
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.setblocking(False)
        sock.sendto(build_mdns_query(hostname), ('224.0.0.251', 5353))
        end_time = time.time() + timeout
        while time.time() < end_time:
            ready, _, _ = select.select([sock], [], [], 0.1)
            if ready:
                try:
                    data, _ = sock.recvfrom(4096)
                    ip = parse_mdns_response(data)
                    if ip:
                        sock.close()
                        return ip
                except Exception:
                    pass
        sock.close()
    except Exception:
        pass
    return None


def main():
    if len(sys.argv) < 2:
        print("Usage: mdns-resolve <hostname.local> [timeout]", file=sys.stderr)
        sys.exit(1)
    hostname = sys.argv[1]
    timeout = float(sys.argv[2]) if len(sys.argv) > 2 else 2.0
    ip = resolve_mdns(hostname, timeout)
    if ip:
        print(ip)
    else:
        sys.exit(1)


if __name__ == "__main__":
    main()