"""Tiny local DNS forwarder for cloudflared when FlClash Fake-IP DNS is enabled.

It resolves A/AAAA records through Cloudflare DoH over the configured HTTP proxy,
then returns a normal DNS response on loopback. It deliberately listens only on
127.0.0.1 and does not persist credentials or query data.
"""

from __future__ import annotations

import ipaddress
import json
import os
import socket
import struct
import sys
from urllib.parse import urlencode
from urllib.request import ProxyHandler, Request, build_opener


LISTEN_HOST = "127.0.0.1"
LISTEN_PORT = int(os.environ.get("DOH_DNS_PORT", "53535"))
DOH_URL = "https://cloudflare-dns.com/dns-query"
LOG_PATH = os.environ.get("DOH_DNS_LOG", "D:/work/code/ai-ecology/c2a/data/cloudflared-doh-dns.log")


def log(message: str) -> None:
    try:
        with open(LOG_PATH, "a", encoding="utf-8") as stream:
            stream.write(message + "\n")
    except OSError:
        pass


def read_question(packet: bytes) -> tuple[str, int, int, int]:
    offset = 12
    labels: list[str] = []
    while offset < len(packet):
        length = packet[offset]
        offset += 1
        if length == 0:
            break
        if length & 0xC0:
            raise ValueError("compressed question name is unsupported")
        labels.append(packet[offset : offset + length].decode("ascii"))
        offset += length
    qname = ".".join(labels) + "."
    qtype, qclass = struct.unpack("!HH", packet[offset : offset + 4])
    return qname, qtype, qclass, offset + 4


def doh_lookup(name: str, qtype: int) -> list[tuple[int, str, int]]:
    query = urlencode({"name": name, "type": "A" if qtype == 1 else "AAAA"})
    request = Request(
        f"{DOH_URL}?{query}",
        headers={"Accept": "application/dns-json", "User-Agent": "c2a-cloudflared-dns/1"},
    )
    proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy") or "http://127.0.0.1:7890"
    opener = build_opener(ProxyHandler({"https": proxy, "http": proxy}) if proxy else ProxyHandler({}))
    with opener.open(request, timeout=15) as response:
        payload = json.load(response)
    records: list[tuple[int, str, int]] = []
    for answer in payload.get("Answer") or []:
        if int(answer.get("type") or 0) == qtype:
            records.append((qtype, str(answer.get("data") or ""), int(answer.get("TTL") or 60)))
    return records


def build_response(query: bytes) -> bytes:
    txid = query[:2]
    qname, qtype, qclass, question_end = read_question(query)
    question = query[12:question_end]
    records: list[tuple[int, str, int]] = []
    if qclass == 1 and qtype in (1, 28):
        records = doh_lookup(qname, qtype)
    log(f"query name={qname} type={qtype} records={[item[1] for item in records]}")
    answers = bytearray()
    for record_type, value, ttl in records:
        try:
            packed = ipaddress.ip_address(value).packed
        except ValueError:
            continue
        answers.extend(struct.pack("!HHHIH", 0xC00C, record_type, 1, max(1, ttl), len(packed)))
        answers.extend(packed)
    flags = 0x8180 if records else 0x8183
    header = txid + struct.pack("!HHHHH", flags, 1, len(records), 0, 0)
    return header + question + bytes(answers)


def main() -> int:
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind((LISTEN_HOST, LISTEN_PORT))
    print(f"DNS forwarder listening on {LISTEN_HOST}:{LISTEN_PORT}", flush=True)
    while True:
        packet, address = sock.recvfrom(4096)
        try:
            response = build_response(packet)
        except Exception:
            log("query failed")
            response = packet[:2] + struct.pack("!HHHHH", 0x8182, 1, 0, 0, 0) + packet[12:]
        sock.sendto(response, address)


if __name__ == "__main__":
    raise SystemExit(main())