import socket
import threading
import json
import base64
import os
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from typing import Callable, Any


class Request:
    def __init__(self, path: str, client_id: str):
        self.path = path
        self.client_id = client_id


class Response:
    def __init__(self, body: str, status: int = 200):
        self.body = body
        self.status = status


def html(content: str) -> Response:
    return Response(content)


class SDCP:
    def __init__(
        self,
        server_id: str,
        server_key: str,
        main_server_host: str = "127.0.0.1",
        main_server_port: int = 300,
    ):
        self.server_id = server_id
        self.server_key = base64.urlsafe_b64decode(server_key + "==")
        self.main_server_host = main_server_host
        self.main_server_port = main_server_port
        self._routes: dict[str, Callable[[Request], Response]] = {}
        self._not_found_handler: Callable[[Request], Response] | None = None

    def route(self, path: str):
        def decorator(func: Callable[[Request], Response]):
            self._routes[path] = func
            return func
        return decorator

    def not_found(self, func: Callable[[Request], Response]):
        self._not_found_handler = func
        return func

    def _encrypt_and_send(self, conn: socket.socket, aesgcm: AESGCM, data: bytes) -> None:
        nonce = os.urandom(12)
        ciphertext = aesgcm.encrypt(nonce, data, None)
        payload = nonce + ciphertext
        conn.send(len(payload).to_bytes(4, "big") + payload)

    def _send_plain(self, conn: socket.socket, data: bytes) -> None:
        conn.send(len(data).to_bytes(4, "big") + data)

    def _recv_exact(self, conn: socket.socket, n: int) -> bytes:
        data = b""
        while len(data) < n:
            chunk = conn.recv(n - len(data))
            if not chunk:
                raise ConnectionError("Connection closed unexpectedly")
            data += chunk
        return data

    def _recv_and_decrypt(self, conn: socket.socket, aesgcm: AESGCM) -> bytes:
        length = int.from_bytes(self._recv_exact(conn, 4), "big")
        data = self._recv_exact(conn, length)
        nonce, ciphertext = data[:12], data[12:]
        return aesgcm.decrypt(nonce, ciphertext, None)

    def _recv_line(self, conn: socket.socket) -> str:
        data = b""
        while not data.endswith(b"\n"):
            chunk = conn.recv(1)
            if not chunk:
                break
            data += chunk
        return data.decode().strip()

    def _request_session_key(self, client_id: str) -> tuple[str, bytes]:
        aesgcm = AESGCM(self.server_key)

        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as conn:
            conn.connect((self.main_server_host, self.main_server_port))
            conn.send(self.server_id.encode() + b"\n")
            self._encrypt_and_send(conn, aesgcm, json.dumps({
                "role": "server",
                "client_id": client_id,
            }).encode())
            response = json.loads(self._recv_and_decrypt(conn, aesgcm).decode())

        return response["session_id"], base64.b64decode(response["session_key"])

    def _handle_client(self, conn: socket.socket, addr: Any) -> None:
        try:
            client_id = self._recv_line(conn)

            session_id, session_key = self._request_session_key(client_id)
            session_aesgcm = AESGCM(session_key)

            self._send_plain(conn, json.dumps({
                "status": "key_ready",
                "session_id": session_id,
                "main_server_host": self.main_server_host,
                "main_server_port": self.main_server_port,
                "server_id": self.server_id,
            }).encode())

            raw_path = self._recv_and_decrypt(conn, session_aesgcm).decode().strip()
            path = raw_path if raw_path else "/"

            request = Request(path=path, client_id=client_id)

            handler = self._routes.get(path)
            if handler is None:
                for route_path, route_handler in self._routes.items():
                    if route_path.endswith("*") and path.startswith(route_path[:-1]):
                        handler = route_handler
                        break

            if handler is None:
                if self._not_found_handler:
                    response = self._not_found_handler(request)
                else:
                    response = Response("<h1>404 Not Found</h1>", status=404)
            else:
                response = handler(request)

            self._encrypt_and_send(conn, session_aesgcm, response.body.encode())

        except Exception as e:
            print(f"[-] Error handling client {addr}: {e}")
        finally:
            conn.close()

    def run(self, host: str = "0.0.0.0", port: int = 299, debug: bool = False):
        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind((host, port))
        server.listen()
        print(f"[+] SDCP server running on {host}:{port}")

        while True:
            conn, addr = server.accept()
            if debug:
                print(f"[+] Connection from {addr}")
            threading.Thread(target=self._handle_client, args=(conn, addr), daemon=True).start()