import queue
import threading
import time
import numpy as np

from db_writer import init_db, log_detection, update_active_flow_count
from inference_engine import classify_flow, THRESHOLD


SEQ_LEN = 8
FLOW_TIMEOUT = 30
REPEAT_LOG_INTERVAL_SEC = 5

active_flows = {}
flows_lock = threading.Lock()

def get_flow_key(packet):
    if not packet.haslayer("IP") or not packet.haslayer("TCP"):
        return None

    ip_layer = packet["IP"]
    tcp_layer = packet["TCP"]

    src_ip = ip_layer.src
    dst_ip = ip_layer.dst
    src_port = tcp_layer.sport
    dst_port = tcp_layer.dport

    if (src_ip, src_port) < (dst_ip, dst_port):
        forward = True
        key = (src_ip, dst_ip, src_port, dst_port)
    else:
        forward = False
        key = (dst_ip, src_ip, dst_port, src_port)

    return key, forward

def create_new_flow_state(timestamp, forward):
    return {
        "first_seen": timestamp,
        "last_seen": timestamp,
        "bytes_fwd": 0,
        "bytes_rev": 0,
        "packets_fwd": 0,
        "packets_rev": 0,
        "ppi_sizes": [],
        "ppi_directions": [],
        "ppi_times": [],
        "tcp_flags_fwd": 0,
        "tcp_flags_rev": 0,
        "final_features_sent": False,
        "client_is_forward": forward,
        "last_decision": None,
        "last_error": None,
        "last_repeat_logged_at": 0.0
    }

def compute_ppi_matrix(flow):
    seq = np.zeros((SEQ_LEN, 3), dtype=np.float64)

    sizes = flow["ppi_sizes"]
    directions = flow["ppi_directions"]
    times = flow["ppi_times"]

    n = len(sizes)

    for i in range(n):
        if i == 0:
            seq[i, 0] = 0.0
        else:
            seq[i, 0] = (times[i] - times[i - 1]) * 1000.0
        seq[i, 1] = directions[i]
        seq[i, 2] = sizes[i]

    return seq

def update_flow(packet):
    result = get_flow_key(packet)
    if result is None:
        return None

    key, forward = result
    timestamp = float(packet.time)
    packet_size = len(packet)
    with flows_lock:
        if key not in active_flows:
            active_flows[key] = create_new_flow_state(timestamp, forward)

        flow = active_flows[key]
        flow["last_seen"] = timestamp

        if forward:
            flow["bytes_fwd"] += packet_size
            flow["packets_fwd"] += 1
        else:
            flow["bytes_rev"] += packet_size
            flow["packets_rev"] += 1

        if len(flow["ppi_sizes"]) < SEQ_LEN:
            flow["ppi_sizes"].append(packet_size)
            is_client_packet = (forward == flow["client_is_forward"])
            flow["ppi_directions"].append(1 if is_client_packet else -1)
            flow["ppi_times"].append(timestamp)

        tcp_layer = packet["TCP"]
        flags_int = int(tcp_layer.flags)

        if forward:
            flow["tcp_flags_fwd"] |= flags_int
        else:
            flow["tcp_flags_rev"] |= flags_int
        return key

def check_flow_completion(key, packet):
    with flows_lock:
        if key not in active_flows:
            return

        flow = active_flows[key]
        ppi_count = len(flow["ppi_sizes"])

        # key = (ip_a, ip_b, port_a, port_b) بترتيب موحّد - منستخدمها للعرض بس
        src_ip, dst_ip, _, dst_port = key

        # ---------- تدفق سبق وتصنّف: نسجل تكرار برأيه القديم بدل ما نتجاهله بصمت ----------
        # محدود بفاصل زمني (REPEAT_LOG_INTERVAL_SEC) لتفادي إغراق قاعدة البيانات
        if flow["final_features_sent"]:
            now = flow["last_seen"]
            if now - flow["last_repeat_logged_at"] >= REPEAT_LOG_INTERVAL_SEC:
                flow["last_repeat_logged_at"] = now
                flow_duration_sec = flow["last_seen"] - flow["first_seen"]
                log_detection(
                    src_ip=src_ip,
                    dst_ip=dst_ip,
                    dst_port=dst_port,
                    stage="REPEATED",
                    decision=flow["last_decision"],
                    reconstruction_error=flow["last_error"],
                    threshold=THRESHOLD,
                    processing_latency_ms=None,
                    flow_duration_sec=flow_duration_sec
                )
            return

        if ppi_count == SEQ_LEN and not flow["final_features_sent"]:
            flow["final_features_sent"] = True
            ppi_matrix = compute_ppi_matrix(flow)
            decision, error, latency_ms = classify_flow(ppi_matrix)
            flow["last_decision"] = decision
            flow["last_error"] = error
            flow_duration_sec = flow["last_seen"] - flow["first_seen"]
            print(f"[FINAL DETECTION] Flow {key} reached {SEQ_LEN} packets (full window)")
            print(f"  Decision: {decision} | Reconstruction Error: {error:.4f} | Latency: {latency_ms:.3f}ms")
            print(f"  FULL RAW MATRIX:\n{ppi_matrix}")

            log_detection(
                src_ip=src_ip,
                dst_ip=dst_ip,
                dst_port=dst_port,
                stage="FINAL_CONFIRMED",
                decision=decision,
                reconstruction_error=error,
                threshold=THRESHOLD,
                processing_latency_ms=latency_ms,
                flow_duration_sec=flow_duration_sec
            )

        tcp_layer = packet["TCP"]
        flags_str = str(tcp_layer.flags)

        if ("F" in flags_str or "R" in flags_str) and not flow["final_features_sent"]:
            flow["final_features_sent"] = True
            ppi_matrix = compute_ppi_matrix(flow)
            decision, error, latency_ms = classify_flow(ppi_matrix)
            flow["last_decision"] = decision
            flow["last_error"] = error
            flow_duration_sec = flow["last_seen"] - flow["first_seen"]
            print(f"[FINAL DETECTION] Flow {key} ended early via FIN/RST at {ppi_count} packets")
            print(f"  Decision: {decision} | Reconstruction Error: {error:.4f} | Latency: {latency_ms:.3f}ms")

            log_detection(
                src_ip=src_ip,
                dst_ip=dst_ip,
                dst_port=dst_port,
                stage="FINAL_CONFIRMED",
                decision=decision,
                reconstruction_error=error,
                threshold=THRESHOLD,
                processing_latency_ms=latency_ms,
                flow_duration_sec=flow_duration_sec
            )

            del active_flows[key]

def check_timeouts():
    current_time = time.time()
    with flows_lock:
        expired_keys = []
        for key, flow in active_flows.items():
            if current_time - flow["last_seen"] > FLOW_TIMEOUT:
                expired_keys.append(key)

        for key in expired_keys:
            flow = active_flows[key]
            if not flow["final_features_sent"]:
                ppi_matrix = compute_ppi_matrix(flow)
                ppi_count = len(flow["ppi_sizes"])
                print(f"[FINAL DETECTION] Flow {key} timed out at {ppi_count} packets")
                print(f"  PPI matrix shape: {ppi_matrix.shape}")
                print(f"  First 3 rows:\n{ppi_matrix[:3]}")

                decision, error, latency_ms = classify_flow(ppi_matrix)
                flow["last_decision"] = decision
                flow["last_error"] = error
                flow_duration_sec = flow["last_seen"] - flow["first_seen"]
                src_ip, dst_ip, _, dst_port = key
                log_detection(
                    src_ip=src_ip,
                    dst_ip=dst_ip,
                    dst_port=dst_port,
                    stage="FINAL_CONFIRMED",
                    decision=decision,
                    reconstruction_error=error,
                    threshold=THRESHOLD,
                    processing_latency_ms=latency_ms,
                    flow_duration_sec=flow_duration_sec
                )

            del active_flows[key]

def process_packets(packet_queue, stop_event):
    while not stop_event.is_set():
        try:
            packet = packet_queue.get(timeout=1)
        except queue.Empty:
            continue

        try:
            key = update_flow(packet)
            if key is not None:
                check_flow_completion(key, packet)
        except Exception as e:
            print(f"[ERROR] Failed to process packet: {e}")

def timeout_checker_loop(stop_event):
    while not stop_event.is_set():
        time.sleep(1)
        try:
            check_timeouts()

            with flows_lock:
                current_count = len(active_flows)
            update_active_flow_count(current_count)

        except Exception as e:
            print(f"[ERROR] Timeout checker failed: {e}")

if __name__ == "__main__":
    from sniffer import sniffer, packet_queue

    init_db()

    stop_event = threading.Event()

    processing_thread = threading.Thread(
        target=process_packets,
        args=(packet_queue, stop_event)
    )
    processing_thread.start()

    timeout_thread = threading.Thread(
        target=timeout_checker_loop,
        args=(stop_event,)
    )
    timeout_thread.start()

    print("[*] Starting live packet capture and flow aggregation...")
    sniffer.start()

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n[*] Stopping...")
        sniffer.stop()
        stop_event.set()
        processing_thread.join()
        timeout_thread.join()
        print("[*] Stopped cleanly.")
