#!/usr/bin/env python3
import hashlib
import ipaddress
import json
import os
import socket
import threading
from typing import Dict, Tuple

from dnslib import A, DNSHeader, DNSRecord, QTYPE, RCODE, RR

LISTEN_HOST = os.getenv("I2PDNS_LISTEN_HOST", "0.0.0.0")
LISTEN_PORT = int(os.getenv("I2PDNS_LISTEN_PORT", "53"))
DEFAULT_TTL = int(os.getenv("I2PDNS_TTL", "60"))
POOL_CIDR = os.getenv("I2PDNS_POOL_CIDR", "172.31.0.0/16")
MAP_FILE = os.getenv("I2PDNS_MAP_FILE", "/var/lib/i2pdns/map.json")

POOL = ipaddress.ip_network(POOL_CIDR, strict=True)
if POOL.version != 4 or POOL.num_addresses < 4:
    raise RuntimeError(f"I2PDNS_POOL_CIDR must be an IPv4 network with >=4 addresses: {POOL_CIDR}")

_lock = threading.Lock()


def _load_map() -> Dict[str, Dict[str, str]]:
    try:
        with open(MAP_FILE, "r", encoding="utf-8") as fh:
            data = json.load(fh)
    except FileNotFoundError:
        data = {}
    except Exception:
        data = {}

    host_to_ip = data.get("host_to_ip") if isinstance(data, dict) else None
    ip_to_host = data.get("ip_to_host") if isinstance(data, dict) else None
    if not isinstance(host_to_ip, dict):
        host_to_ip = {}
    if not isinstance(ip_to_host, dict):
        ip_to_host = {}
    return {"host_to_ip": host_to_ip, "ip_to_host": ip_to_host}


def _save_map(data: Dict[str, Dict[str, str]]) -> None:
    os.makedirs(os.path.dirname(MAP_FILE), exist_ok=True)
    tmp = MAP_FILE + ".tmp"
    with open(tmp, "w", encoding="utf-8") as fh:
        json.dump(data, fh, sort_keys=True)
    os.replace(tmp, MAP_FILE)


def _candidate_index(host: str, host_space: int) -> int:
    digest = hashlib.blake2b(host.encode("utf-8"), digest_size=8).digest()
    return (int.from_bytes(digest, "big") % host_space) + 1


def _index_to_ip(idx: int) -> str:
    return str(ipaddress.IPv4Address(int(POOL.network_address) + idx))


def _allocate_ip_for_host(host: str) -> str:
    # Exclude network and broadcast addresses in the pool.
    host_space = int(POOL.num_addresses) - 2
    if host_space <= 0:
        raise RuntimeError("I2P DNS pool has no usable host addresses")

    with _lock:
        data = _load_map()
        host_to_ip = data["host_to_ip"]
        ip_to_host = data["ip_to_host"]

        existing = host_to_ip.get(host)
        if existing:
            return existing

        start = _candidate_index(host, host_space)
        idx = start
        while True:
            ip = _index_to_ip(idx)
            current = ip_to_host.get(ip)
            if current is None or current == host:
                host_to_ip[host] = ip
                ip_to_host[ip] = host
                _save_map(data)
                return ip

            idx += 1
            if idx > host_space:
                idx = 1
            if idx == start:
                raise RuntimeError("I2P DNS pool exhausted")


def _is_i2p_name(name: str) -> bool:
    return name.endswith(".i2p")


def _answer(req: DNSRecord) -> DNSRecord:
    reply = DNSRecord(DNSHeader(id=req.header.id, qr=1, aa=1, ra=1), q=req.q)

    qname = str(req.q.qname).rstrip(".").lower()
    qtype = QTYPE[req.q.qtype]

    if not _is_i2p_name(qname):
        reply.header.rcode = RCODE.NXDOMAIN
        return reply

    if qtype in ("A", "ANY"):
        ip = _allocate_ip_for_host(qname)
        reply.add_answer(RR(str(req.q.qname), QTYPE.A, rdata=A(ip), ttl=DEFAULT_TTL))
        return reply

    if qtype == "AAAA":
        # We only synthesize IPv4 addresses for SOCKS routing compatibility.
        return reply

    # Unsupported RR types get NOERROR with empty answer section.
    return reply


def _serve_udp(sock: socket.socket) -> None:
    while True:
        data, addr = sock.recvfrom(4096)
        try:
            req = DNSRecord.parse(data)
            resp = _answer(req).pack()
            sock.sendto(resp, addr)
        except Exception:
            continue


def _handle_tcp_conn(conn: socket.socket) -> None:
    try:
        hdr = conn.recv(2)
        if len(hdr) != 2:
            return
        size = int.from_bytes(hdr, "big")
        payload = conn.recv(size)
        if len(payload) != size:
            return
        req = DNSRecord.parse(payload)
        resp = _answer(req).pack()
        conn.sendall(len(resp).to_bytes(2, "big") + resp)
    except Exception:
        pass
    finally:
        conn.close()


def _serve_tcp(sock: socket.socket) -> None:
    while True:
        conn, _ = sock.accept()
        threading.Thread(target=_handle_tcp_conn, args=(conn,), daemon=True).start()


def main() -> None:
    print(f"[i2pdns] listening on {LISTEN_HOST}:{LISTEN_PORT}", flush=True)
    print(f"[i2pdns] pool={POOL} map_file={MAP_FILE}", flush=True)

    udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp.bind((LISTEN_HOST, LISTEN_PORT))

    tcp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    tcp.bind((LISTEN_HOST, LISTEN_PORT))
    tcp.listen(128)

    threading.Thread(target=_serve_udp, args=(udp,), daemon=True).start()
    _serve_tcp(tcp)


if __name__ == "__main__":
    main()