import http.server
import ipaddress
import os
import signal
import socket
import socketserver
import ssl
import struct
import threading


DNS_PORT = int(os.environ.get("DNS_PORT", "8053"))
HTTP_PORT = int(os.environ.get("HTTP_PORT", "8081"))
HTTPS_PORT = int(os.environ.get("HTTPS_PORT", "8443"))
GATEWAY_IP = ipaddress.IPv4Address(os.environ.get("GATEWAY_IP", "192.168.20.30"))
REDIRECT_URL = os.environ.get("REDIRECT_URL", f"http://{GATEWAY_IP}:8000/")
TARGET_DOMAIN = "manuals.playstation.net"
# Exact names forwarded to the upstream resolver (homebrew store: catalog + downloads).
FORWARDED_DOMAINS = frozenset(
    name.strip().lower()
    for name in os.environ.get(
        "FORWARDED_DOMAINS",
        "homebrew.page,itsplk.github.io,store.spectrumlibrary.online,api.github.com,github.com,raw.githubusercontent.com,release-assets.githubusercontent.com,prosperopatches.com,cdn.prosperopatches.com,image.api.playstation.com,gmedia.playstation.com",
    ).split(",")
    if name.strip()
)
UPSTREAM_DNS = str(ipaddress.ip_address(os.environ.get("UPSTREAM_DNS", "1.1.1.1")))
UPSTREAM_TIMEOUT = 3.0
CERTIFICATE = "/srv/certificates/manuals.crt"
PRIVATE_KEY = "/srv/certificates/manuals.key"


def parse_question(packet):
    if len(packet) < 17:
        raise ValueError("short DNS query")

    labels = []
    offset = 12
    while True:
        if offset >= len(packet):
            raise ValueError("unterminated DNS name")
        length = packet[offset]
        offset += 1
        if length == 0:
            break
        if length & 0xC0 or length > 63 or offset + length > len(packet):
            raise ValueError("unsupported DNS name")
        labels.append(packet[offset : offset + length].decode("ascii"))
        offset += length

    if offset + 4 > len(packet):
        raise ValueError("short DNS question")
    qtype, qclass = struct.unpack_from("!HH", packet, offset)
    return ".".join(labels).lower(), qtype, qclass, offset + 4


def error_response(packet, rcode):
    transaction_id, request_flags = struct.unpack_from("!HH", packet, 0)
    response_flags = 0x8000 | (request_flags & 0x0100) | rcode
    try:
        _, _, _, question_end = parse_question(packet)
        question = packet[12:question_end]
        question_count = 1
    except (UnicodeDecodeError, ValueError):
        question = b""
        question_count = 0
    return (
        struct.pack("!HHHHHH", transaction_id, response_flags, question_count, 0, 0, 0)
        + question
    )


def receive_exactly(connection, wanted):
    chunks = []
    received = 0
    while received < wanted:
        chunk = connection.recv(wanted - received)
        if not chunk:
            raise OSError("upstream DNS connection closed early")
        chunks.append(chunk)
        received += len(chunk)
    return b"".join(chunks)


def forward_query(packet, use_tcp=False):
    family = socket.AF_INET6 if ":" in UPSTREAM_DNS else socket.AF_INET
    socket_type = socket.SOCK_STREAM if use_tcp else socket.SOCK_DGRAM
    with socket.socket(family, socket_type) as upstream:
        upstream.settimeout(UPSTREAM_TIMEOUT)
        upstream.connect((UPSTREAM_DNS, 53))
        if use_tcp:
            upstream.sendall(struct.pack("!H", len(packet)) + packet)
            response_length = struct.unpack("!H", receive_exactly(upstream, 2))[0]
            response = receive_exactly(upstream, response_length)
        else:
            upstream.send(packet)
            response = upstream.recv(65535)

    if len(response) < 12 or response[:2] != packet[:2]:
        raise ValueError("invalid upstream DNS response")
    return response


def dns_response(packet, use_tcp=False):
    if len(packet) < 12:
        return b""

    transaction_id, request_flags, question_count = struct.unpack_from("!HHH", packet, 0)
    response_flags = 0x8400 | (request_flags & 0x0100)  # response, authoritative, preserve RD
    try:
        name, qtype, qclass, question_end = parse_question(packet)
        if question_count != 1:
            raise ValueError("expected one DNS question")
    except (UnicodeDecodeError, ValueError):
        return struct.pack("!HHHHHH", transaction_id, response_flags | 1, 0, 0, 0, 0)

    question = packet[12:question_end]
    if name in FORWARDED_DOMAINS:
        try:
            return forward_query(packet, use_tcp=use_tcp)
        except (OSError, ValueError) as error:
            print(f"DNS upstream failure for {name}: {error}", flush=True)
            return error_response(packet, 2)  # SERVFAIL

    if name != TARGET_DOMAIN:
        return struct.pack("!HHHHHH", transaction_id, response_flags | 3, 1, 0, 0, 0) + question

    if qclass == 1 and qtype in (1, 255):
        answer = (
            b"\xc0\x0c"
            + struct.pack("!HHIH", 1, 1, 60, 4)
            + GATEWAY_IP.packed
        )
        return struct.pack("!HHHHHH", transaction_id, response_flags, 1, 1, 0, 0) + question + answer

    # NOERROR with no data prevents IPv6 or unsupported record fallback.
    return struct.pack("!HHHHHH", transaction_id, response_flags, 1, 0, 0, 0) + question


class ThreadingUDPServer(socketserver.ThreadingUDPServer):
    allow_reuse_address = True
    daemon_threads = True


class ThreadingTCPServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True


class UDPDNSHandler(socketserver.BaseRequestHandler):
    def handle(self):
        packet, connection = self.request
        try:
            name, qtype, _, _ = parse_question(packet)
            print(f"DNS UDP {self.client_address[0]} {name} type={qtype}", flush=True)
        except (UnicodeDecodeError, ValueError):
            print(f"DNS UDP {self.client_address[0]} malformed", flush=True)
        response = dns_response(packet)
        if response:
            connection.sendto(response, self.client_address)


class TCPDNSHandler(socketserver.BaseRequestHandler):
    def handle(self):
        length_bytes = self._receive_exactly(2)
        if length_bytes is None:
            return
        wanted = struct.unpack("!H", length_bytes)[0]
        packet = self._receive_exactly(wanted)
        if packet is None:
            return
        try:
            name, qtype, _, _ = parse_question(packet)
            print(f"DNS TCP {self.client_address[0]} {name} type={qtype}", flush=True)
        except (UnicodeDecodeError, ValueError):
            print(f"DNS TCP {self.client_address[0]} malformed", flush=True)
        response = dns_response(packet, use_tcp=True)
        if response:
            self.request.sendall(struct.pack("!H", len(response)) + response)

    def _receive_exactly(self, wanted):
        try:
            return receive_exactly(self.request, wanted)
        except OSError:
            return None


class RedirectHandler(http.server.BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def do_GET(self):
        if self.path == "/_health":
            body = b"ok\n"
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return

        self.send_response(302)
        self.send_header("Location", REDIRECT_URL)
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", "0")
        self.end_headers()

    def do_HEAD(self):
        self.send_response(302)
        self.send_header("Location", REDIRECT_URL)
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", "0")
        self.end_headers()

    def log_message(self, message, *args):
        user_agent = self.headers.get("User-Agent", "-")
        print(
            f"HTTP {self.client_address[0]} {message % args} user-agent={user_agent!r}",
            flush=True,
        )


class HTTPServer(http.server.ThreadingHTTPServer):
    allow_reuse_address = True
    daemon_threads = True


def serve(server, name):
    thread = threading.Thread(target=server.serve_forever, name=name, daemon=True)
    thread.start()
    return thread


if __name__ == "__main__":
    dns_udp = ThreadingUDPServer(("0.0.0.0", DNS_PORT), UDPDNSHandler)
    dns_tcp = ThreadingTCPServer(("0.0.0.0", DNS_PORT), TCPDNSHandler)
    http = HTTPServer(("0.0.0.0", HTTP_PORT), RedirectHandler)
    https = HTTPServer(("0.0.0.0", HTTPS_PORT), RedirectHandler)

    tls = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    tls.load_cert_chain(CERTIFICATE, PRIVATE_KEY)
    https.socket = tls.wrap_socket(https.socket, server_side=True)

    servers = [dns_udp, dns_tcp, http, https]
    serve(dns_udp, "dns-udp")
    serve(dns_tcp, "dns-tcp")
    serve(http, "http-redirect")
    serve(https, "https-redirect")

    print(
        f"PS5 guide gateway: {TARGET_DOMAIN} -> {GATEWAY_IP}; redirect -> {REDIRECT_URL}; "
        f"forward {','.join(sorted(FORWARDED_DOMAINS))} -> {UPSTREAM_DNS}",
        flush=True,
    )

    stopped = threading.Event()

    def stop(*_):
        stopped.set()

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    stopped.wait()
    for server in servers:
        server.shutdown()
        server.server_close()