# DimyServer.py - Backend server for DIMY protocol

import socket
import threading
import http.server
import json
from pybloom_live import BloomFilter
from http.server import BaseHTTPRequestHandler, HTTPServer
import pickle
import os

CBF_STORAGE = []  # Stores tuples of (client_id, CBF) uploaded by diagnosed users
TCP_PORT = 55000  # For QBF queries
HTTP_PORT = 8000  # For CBF uploads

if os.getenv("DIMY_SERVER_TCP_PORT"):
    TCP_PORT = int(os.getenv("DIMY_SERVER_TCP_PORT"))
    print(f"[ENV] Using TCP port from environment: {TCP_PORT}")

class CBFUploadHandler(BaseHTTPRequestHandler):
    def do_POST(self):
        print(f"[HTTP] Incoming POST to {self.path}")

        if self.path == '/upload_cbf':
            content_length = int(self.headers['Content-Length'])
            post_data = self.rfile.read(content_length)

            try:
                data = json.loads(post_data)
                cbf_hex = data.get("cbf")

                # Debug:show preview of received CBF data
                print(f"[DEBUG] Received CBF hex: {cbf_hex[:60]}...")

                cbf_bytes = bytes.fromhex(cbf_hex)
                cbf = pickle.loads(cbf_bytes)
                client_id = data.get("client_id", "unknown")

                CBF_STORAGE.append((client_id, cbf))
                print(f"[HTTP] CBF uploaded and stored. It contains approx. {cbf.count} items")
                if hasattr(cbf, 'bitarray'):
                    print("[CBF DEBUG] Bitarray size:", len(cbf.bitarray))

                self.send_response(200)
                self.end_headers()
                self.wfile.write(b"CBF upload successful")
            except Exception as e:
                print(f"[UPLOAD ERROR] {e}")
                self.send_response(400)
                self.end_headers()
                self.wfile.write(f"Error: {str(e)}".encode())


def start_http_server():
    server = HTTPServer(("", HTTP_PORT), CBFUploadHandler)
    print(f"[HTTP] CBF upload server running on port {HTTP_PORT}...")
    server.serve_forever()


def handle_client_connection(client_socket):
    try:
        qbf_data = b""
        while True:
            part = client_socket.recv(4096)
            if not part:
                break
            qbf_data += part

        payload = pickle.loads(qbf_data)
        qbf = pickle.loads(bytes.fromhex(payload["qbf"]))
        encids = payload.get("encids", [])
        client_id = payload.get("client_id", "unknown")
        print("\n" + "/" * 30 + f" CLIENT {client_id[:8]} " + "/" * 30 + "\n")


        print(f"[DEBUG] Received {len(qbf_data)} bytes of QBF data")
        print("[DEBUG] QBF deserialized successfully")
        print("[DEBUG] QBF received — count not available due to pybloom serialization")

        # Dump content-like preview if known (optional and not guaranteed accurate)
        if hasattr(qbf, 'bitarray'):
            print("[QBF DEBUG] Bitarray size:", len(qbf.bitarray))

        if not CBF_STORAGE:
            print("[DEBUG] No CBFs stored yet — auto 'not matched'")
        else:
            for i, (cbf_owner, cbf) in enumerate(CBF_STORAGE):
                print(f"[DEBUG] CBF {i} contains approx. {cbf.count} items")
                if hasattr(cbf, 'bitarray'):
                    print(f"[CBF {i} DEBUG] Bitarray size:", len(cbf.bitarray))
                if cbf_owner == client_id:
                    print(f"[DEBUG] Skipping self CBF {i} (client_id match)")
                    continue
                print("[MATCH DEBUG] Checking EncIDs in QBF against CBF:")
                for encid in encids:
                    print(f"  → Checking {encid[:8]}...", end="")
                    try:
                        if encid.encode() in cbf:
                            print(" ✅ FOUND")
                        else:
                            print(" ❌ NOT FOUND")
                    except Exception as e:
                        print(f" [ERROR] {e}")

        matched = False
        for i, (cbf_owner, cbf) in enumerate(CBF_STORAGE):
            if cbf_owner == client_id:
                continue
            hit = any(encid.encode() in cbf for encid in encids)
            if hit:
                print(f"[MATCH DEBUG] QBF ∩ CBF {i}: match found ✅")
                matched = True
            else:
                print(f"[MATCH DEBUG] QBF ∩ CBF {i}: no match ❌")

        response = b"matched" if matched else b"not matched"
        if response == b"matched":
            print("🔴 [MATCH] Match result: matched")
        else:
            print("✅ [MATCH] Match result: not matched")

        client_socket.sendall(response)
        print(f"[TCP] Sent response: {response.decode()}")
        print("\n" + "=" * 70 + "\n")


    except Exception as inner_e:
        print(f"[MATCH ERROR] {inner_e}")
        client_socket.sendall(b"error")

    except Exception as e:
        print(f"[TCP ERROR] {e}")
    finally:
        client_socket.close()


def start_tcp_server():
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind(("", TCP_PORT))
    server_socket.listen()
    print(f"[TCP] Query server running on port {TCP_PORT}...")

    while True:
        client_sock, addr = server_socket.accept()
        print(f"[TCP] Connection from {addr}")
        threading.Thread(target=handle_client_connection, args=(client_sock,), daemon=True).start()


def main():
    threading.Thread(target=start_http_server, daemon=True).start()
    start_tcp_server()


if __name__ == '__main__':
    main()