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"
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 dns_response(packet):
    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 != 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)
        if response:
            self.request.sendall(struct.pack("!H", len(response)) + response)

    def _receive_exactly(self, wanted):
        chunks = []
        received = 0
        while received < wanted:
            chunk = self.request.recv(wanted - received)
            if not chunk:
                return None
            chunks.append(chunk)
            received += len(chunk)
        return b"".join(chunks)


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}",
        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()