# src/sniffer.py
import asyncio
import logging
import threading
import time
import socket
import selectors
import errno
import struct
from typing import Dict, List, Optional, Any

# NOTE: ensure this path points to your shared runtime module
from src import shared_objects

from scapy.all import (
    Ether,
    ARP,
    IP,
    IPv6,
    TCP,
    UDP,
    ICMP,
    ICMPv6Unknown,
    Dot1Q,
    Raw,
)

from src.utilities.interface_bridge_helpers import (
    check_interface_exists,
    check_interface_up,
    get_bridge_ports_once,
)
from src.Models.etherType import EtherTypeEnum, ethertype_from_int
from src.Models.ip_protocol import IPProtocolEnum, protocol_from_number

# ---- Logging ----------------------------------------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("af_packet_sniffer")

# ---- Globals (kept minimal) -------------------------------------------
af_sockets: Dict[str, socket.socket] = {}
af_thread: Optional[threading.Thread] = None
af_stop_event: Optional[threading.Event] = None

fixed_bridge_ports: Dict[str, List[str]] = {}
current_bridge: Optional[str] = None

# small bounded buffer for packets produced before shared_objects is ready
_PACKET_BUFFER: List[Dict[str, Any]] = []
_BUFFER_CAPACITY = 20000

# Background asyncio loop used for internal tasks in this module (kept but not used for DB pool)
async_loop = asyncio.new_event_loop()


def _start_async_loop(loop: asyncio.AbstractEventLoop) -> None:
    asyncio.set_event_loop(loop)
    loop.run_forever()


# Start background loop in daemon thread immediately
threading.Thread(target=_start_async_loop, args=(async_loop,), daemon=True).start()


# -------------------------
# Helpers for buffer draining
# -------------------------
def drain_buffer_to_shared_db() -> None:
    """
    Attempt to schedule buffered packets for insertion on shared_objects.web_loop.
    Call this from main.py after shared_objects.db and shared_objects.web_loop are initialized.
    """
    try:
        web_loop = getattr(shared_objects, "web_loop", None)
        web_db = getattr(shared_objects, "db", None)
        if web_db is None or web_loop is None:
            return

        # schedule draining on the web loop to avoid blocking this thread
        def _drain() -> None:
            while _PACKET_BUFFER:
                pkt = _PACKET_BUFFER.pop(0)
                try:
                    asyncio.run_coroutine_threadsafe(web_db.insert_packet(pkt), web_loop)
                except Exception:
                    # re-buffer first element and stop to avoid busy loop
                    _PACKET_BUFFER.insert(0, pkt)
                    break

        try:
            web_loop.call_soon_threadsafe(_drain)
        except Exception:
            # fallback: run directly (best-effort)
            _drain()
    except Exception:
        logger.exception("Failed to drain packet buffer")


# -------------------------
# Packet parsing
# -------------------------
def _safe_get_attr(layer, attr: str):
    try:
        return getattr(layer, attr, None)
    except Exception:
        return None


def parse_packet(pkt, bridge: str) -> None:
    """
    Parse a scapy Packet object into a normalized dict and schedule DB insert.
    """
    pkt_iface = getattr(pkt, "sniffed_on", None)
    if not pkt_iface:
        return  # can't determine interface context

    logger.debug("Packet captured on %s (bridge %s)", pkt_iface, bridge)

    pkt_info: Dict[str, Any] = {
        "iface": pkt_iface,
        "length": len(pkt),
        "raw": bytes(pkt),
        "src_mac": None,
        "dst_mac": None,
        "eth_type_raw": None,
        "eth_type": EtherTypeEnum.UNKNOWN,
        "vlan_id": None,
        "protocol_raw": None,
        "protocol": IPProtocolEnum.UNKNOWN,
        "protocol_name": None,
        "src_ip": None,
        "dst_ip": None,
        "src_port": None,
        "dst_port": None,
    }

    # Ethernet layer
    if Ether in pkt:
        eth = pkt[Ether]
        pkt_info["src_mac"] = _safe_get_attr(eth, "src")
        pkt_info["dst_mac"] = _safe_get_attr(eth, "dst")

        # Base ethertype
        try:
            eth_type_raw = int(eth.type)
        except Exception:
            eth_type_raw = None

        # VLAN inner ethertype and vlan id if Dot1Q exists
        if Dot1Q in pkt:
            try:
                inner = int(pkt[Dot1Q].type)
                if inner:
                    eth_type_raw = inner
            except Exception:
                pass
            try:
                pkt_info["vlan_id"] = int(pkt[Dot1Q].vlan)
            except Exception:
                pkt_info["vlan_id"] = None

        if eth_type_raw is not None:
            pkt_info["eth_type_raw"] = eth_type_raw
            try:
                pkt_info["eth_type"] = ethertype_from_int(eth_type_raw)
            except Exception:
                pkt_info["eth_type"] = EtherTypeEnum.UNKNOWN

    # ARP
    if ARP in pkt:
        arp = pkt[ARP]
        pkt_info["protocol_name"] = "ARP"
        pkt_info["src_ip"] = _safe_get_attr(arp, "psrc")
        pkt_info["dst_ip"] = _safe_get_attr(arp, "pdst")
        pkt_info["src_port"] = None
        pkt_info["dst_port"] = None

    # IPv4
    if IP in pkt:
        ip = pkt[IP]
        pkt_info["src_ip"] = _safe_get_attr(ip, "src")
        pkt_info["dst_ip"] = _safe_get_attr(ip, "dst")

        try:
            proto_num = int(_safe_get_attr(ip, "proto"))
        except Exception:
            proto_num = None

        if proto_num is not None:
            pkt_info["protocol_raw"] = proto_num
            try:
                pkt_info["protocol"] = protocol_from_number(proto_num)
            except Exception:
                pkt_info["protocol"] = IPProtocolEnum.UNKNOWN

        if proto_num == 6 and TCP in pkt:
            pkt_info["protocol_name"] = "TCP"
            pkt_info["src_port"] = _safe_get_attr(pkt[TCP], "sport")
            pkt_info["dst_port"] = _safe_get_attr(pkt[TCP], "dport")
        elif proto_num == 17 and UDP in pkt:
            pkt_info["protocol_name"] = "UDP"
            pkt_info["src_port"] = _safe_get_attr(pkt[UDP], "sport")
            pkt_info["dst_port"] = _safe_get_attr(pkt[UDP], "dport")
        elif proto_num == 1 and ICMP in pkt:
            pkt_info["protocol_name"] = "ICMP"
        else:
            if pkt_info["protocol_name"] is None:
                pkt_info["protocol_name"] = f"IP_PROTO_{proto_num}" if proto_num is not None else None

    # IPv6
    if IPv6 in pkt:
        ip6 = pkt[IPv6]
        pkt_info["src_ip"] = _safe_get_attr(ip6, "src")
        pkt_info["dst_ip"] = _safe_get_attr(ip6, "dst")

        try:
            nh = int(_safe_get_attr(ip6, "nh"))
        except Exception:
            nh = None

        if nh is not None:
            pkt_info["protocol_raw"] = nh
            try:
                pkt_info["protocol"] = protocol_from_number(nh)
            except Exception:
                pkt_info["protocol"] = IPProtocolEnum.UNKNOWN

        if nh == 6 and TCP in pkt:
            pkt_info["protocol_name"] = "TCP"
            pkt_info["src_port"] = _safe_get_attr(pkt[TCP], "sport")
            pkt_info["dst_port"] = _safe_get_attr(pkt[TCP], "dport")
        elif nh == 17 and UDP in pkt:
            pkt_info["protocol_name"] = "UDP"
            pkt_info["src_port"] = _safe_get_attr(pkt[UDP], "sport")
            pkt_info["dst_port"] = _safe_get_attr(pkt[UDP], "dport")
        elif ICMPv6Unknown in pkt:
            pkt_info["protocol_name"] = "ICMPv6"
        else:
            if pkt_info["protocol_name"] is None:
                pkt_info["protocol_name"] = f"IPV6_PROTO_{nh}" if nh is not None else None

    # Raw fallback label
    if Raw in pkt and not pkt_info["protocol_name"]:
        pkt_info["protocol_name"] = "RAW"

    # Submit DB insert to shared web loop if available, otherwise buffer
    try:
        web_loop = getattr(shared_objects, "web_loop", None)
        web_db = getattr(shared_objects, "db", None)
        if web_db is not None and web_loop is not None:
            asyncio.run_coroutine_threadsafe(web_db.insert_packet(pkt_info), web_loop)
        else:
            # buffer (bounded) until the web app initializes
            _PACKET_BUFFER.append(pkt_info)
            if len(_PACKET_BUFFER) > _BUFFER_CAPACITY:
                # drop oldest packet
                _PACKET_BUFFER.pop(0)
    except Exception:
        logger.exception("Failed to schedule DB insert")


# -------------------------
# AF_PACKET socket utilities
# -------------------------
def _create_af_packet_socket(ifname: str, rx_buf: int = 4 * 1024 * 1024) -> Optional[socket.socket]:
    """
    Create and bind an AF_PACKET raw socket to interface.
    Non-blocking socket returned or None on failure.
    """
    try:
        s = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(0x0003))  # ETH_P_ALL
    except PermissionError:
        logger.exception("Permission denied creating AF_PACKET socket for %s", ifname)
        return None
    except Exception:
        logger.exception("Failed creating AF_PACKET socket for %s", ifname)
        return None

    try:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, rx_buf)
    except Exception:
        logger.debug("SO_RCVBUF set failed for %s (non-fatal)", ifname)

    # Try TPACKET_V3 best-effort
    try:
        SOL_PACKET = getattr(socket, "SOL_PACKET", 263)
        PACKET_VERSION = getattr(socket, "PACKET_VERSION", 10)
        TPACKET_V3 = 3
        s.setsockopt(SOL_PACKET, PACKET_VERSION, struct.pack("I", TPACKET_V3))
        logger.debug("Requested TPACKET_V3 on %s", ifname)
    except Exception:
        logger.debug("TPACKETv3 not available for %s", ifname)

    try:
        s.bind((ifname, 0))
    except OSError as e:
        logger.exception("Bind failed for %s: %s", ifname, e)
        s.close()
        return None

    s.setblocking(False)
    return s


def _close_socket(ifname: str) -> None:
    s = af_sockets.pop(ifname, None)
    if s:
        try:
            s.close()
        except Exception:
            pass


# -------------------------
# Main reader thread
# -------------------------
def afpacket_reader_loop(bridge: str, stop_event: threading.Event) -> None:
    """
    Multiplex AF_PACKET sockets using a selector and hand packets to parse_packet.
    Uses the fixed_bridge_ports snapshot to decide which interfaces to open.
    """
    logger.info("AF_PACKET reader starting for bridge %s", bridge)
    sel = selectors.DefaultSelector()

    # create sockets for fixed snapshot (if any)
    for iface in fixed_bridge_ports.get(bridge, []):
        if iface in af_sockets:
            continue
        if not check_interface_exists(iface):
            logger.warning("Snapshot port %s missing, skipping", iface)
            continue
        s = _create_af_packet_socket(iface)
        if s:
            af_sockets[iface] = s

    # register existing sockets
    for iface, s in list(af_sockets.items()):
        try:
            sel.register(s, selectors.EVENT_READ, data=iface)
        except Exception:
            logger.debug("Register failed for %s (continuing)", iface)

    try:
        while not stop_event.is_set():
            # ensure newly created sockets are registered
            for iface, s in list(af_sockets.items()):
                try:
                    if not any(k.fileobj is s for k in sel.get_map().values()):
                        sel.register(s, selectors.EVENT_READ, data=iface)
                except Exception:
                    pass

            try:
                events = sel.select(timeout=1.0)
            except Exception:
                logger.exception("Selector error")
                time.sleep(0.1)
                continue

            if not events:
                continue

            for key, _ in events:
                sock: socket.socket = key.fileobj
                iface: str = key.data
                try:
                    raw = sock.recv(65536)
                    if not raw:
                        continue
                except BlockingIOError:
                    continue
                except OSError as e:
                    if e.errno in (errno.ENODEV, errno.ENETDOWN, errno.EBADF):
                        logger.warning("Socket error on %s: %s — closing", iface, e)
                        try:
                            sel.unregister(sock)
                        except Exception:
                            pass
                        _close_socket(iface)
                        continue
                    else:
                        logger.exception("Recv error on %s", iface)
                        continue

                # parse with scapy
                try:
                    pkt = Ether(raw)
                    pkt.sniffed_on = iface
                    parse_packet(pkt, bridge)
                except Exception:
                    logger.exception("Failed to parse/process packet from %s", iface)
                    continue
    finally:
        logger.info("AF_PACKET reader stopping; cleaning up sockets")
        # unregister and close
        try:
            for key in list(sel.get_map().values()):
                try:
                    sel.unregister(key.fileobj)
                except Exception:
                    pass
        except Exception:
            pass

        for iface in list(af_sockets.keys()):
            _close_socket(iface)

        try:
            sel.close()
        except Exception:
            pass

        logger.info("AF_PACKET reader stopped")


# -------------------------
# Public API: start/stop/status
# -------------------------
def start_afpacket_sniffer(bridge: str) -> None:
    """
    Start the sniffer: snapshot ports once, spin up reader thread.
    """
    global af_thread, af_stop_event, current_bridge
    if af_thread and af_thread.is_alive():
        logger.info("Sniffer already running")
        return

    ports = get_bridge_ports_once(bridge)
    fixed_bridge_ports[bridge] = ports
    current_bridge = bridge
    logger.info("Fixed bridge ports for %s: %s", bridge, ports)

    af_stop_event = threading.Event()
    af_thread = threading.Thread(target=afpacket_reader_loop, args=(bridge, af_stop_event), daemon=True)
    af_thread.start()
    logger.info("AF_PACKET sniffer started")


def stop_afpacket_sniffer() -> None:
    """
    Stop the reader thread and clear snapshot. Also attempt to close DB pool.
    """
    global af_thread, af_stop_event, current_bridge
    if not af_thread:
        return

    if af_stop_event:
        af_stop_event.set()

    af_thread.join(timeout=2)
    af_thread = None
    af_stop_event = None

    if current_bridge:
        fixed_bridge_ports.pop(current_bridge, None)
    current_bridge = None

    # close db pool if shared.web_loop is available; otherwise leave to main
    try:
        web_loop = getattr(shared_objects, "web_loop", None)
        web_db = getattr(shared_objects, "db", None)
        if web_db is not None and web_loop is not None:
            asyncio.run_coroutine_threadsafe(web_db.close_pool(), web_loop)
    except Exception:
        logger.exception("Failed to schedule DB pool close")

    logger.info("AF_PACKET sniffer stopped")


def get_sniffer_status() -> Dict[str, Dict[str, object]]:
    """
    Return simple status per managed interface.
    """
    running = af_thread.is_alive() if af_thread else False
    out: Dict[str, Dict[str, object]] = {}
    for iface in list(af_sockets.keys()):
        out[iface] = {
            "running": running,
            "exists": check_interface_exists(iface),
            "up": check_interface_up(iface),
        }
    return out