import os
import time
import datetime
import sys
import json
import shlex

# Ensure we can import from common
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))

from collectors.common.es_client import get_es_client
from collectors.common.opnsense_client import get_opnsense_client
from collectors.common.nmap_parser import run_nmap_scan
from collectors.common.logging_config import setup_logging

logger = setup_logging("nmap_collector")

def get_now_iso():
    return datetime.datetime.now(datetime.timezone.utc).isoformat()

def chunk_list(lst, n):
    for i in range(0, len(lst), n):
        yield lst[i:i + n]

def should_scan_vlan(vlan, allowlist):
    if not allowlist:
        return True
    name = (vlan.get("name") or "").strip()
    key = (vlan.get("key") or "").strip()
    return name in allowlist or key in allowlist

def build_discovery_update_action(host_id, mac, ip, hostname, vendor, ts_iso):
    mac_norm = mac.lower() if mac else None
    upsert_host = {
        "host": {
            "id": host_id,
            "macs": [mac_norm] if mac_norm else [],
            "ips": [ip] if ip else [],
            "name": hostname,
            "hostnames": [hostname] if hostname else [],
            "vendor": vendor,
            "sources": ["nmap-discovery"],
            "last_seen": ts_iso,
            "first_seen": ts_iso
        }
    }

    script_source = """
        if (ctx._source.host == null) { ctx._source.host = [:]; }
        if (ctx._source.host.macs == null) { ctx._source.host.macs = []; }
        if (ctx._source.host.ips == null) { ctx._source.host.ips = []; }
        if (ctx._source.host.hostnames == null) { ctx._source.host.hostnames = []; }
        if (ctx._source.host.sources == null) { ctx._source.host.sources = []; }

        if (params.mac != null && !ctx._source.host.macs.contains(params.mac)) {
            ctx._source.host.macs.add(params.mac);
        }
        if (params.ip != null && !ctx._source.host.ips.contains(params.ip)) {
            ctx._source.host.ips.add(params.ip);
        }
        if (params.hostname != null && !ctx._source.host.hostnames.contains(params.hostname)) {
            ctx._source.host.hostnames.add(params.hostname);
        }
        if (!ctx._source.host.sources.contains(params.source_tag)) {
            ctx._source.host.sources.add(params.source_tag);
        }
        ctx._source.host.last_seen = params.ts;
        if (ctx._source.host.name == null && params.hostname != null) {
            ctx._source.host.name = params.hostname;
        }
        if (params.vendor != null && (ctx._source.host.vendor == null || ctx._source.host.vendor == \"\")) {
            ctx._source.host.vendor = params.vendor;
        }
    """

    return {
        "_index": "network-hosts",
        "_op_type": "update",
        "_id": host_id,
        "script": {
            "source": script_source,
            "lang": "painless",
            "params": {
                "mac": mac_norm,
                "ip": ip,
                "hostname": hostname,
                "vendor": vendor,
                "ts": ts_iso,
                "source_tag": "nmap-discovery"
            }
        },
        "upsert": upsert_host
    }

def run_vlan_discovery(es, opnsense_client, discovery_args, vlan_filter):
    networks = opnsense_client.get_vlan_networks()
    if not networks:
        logger.info("VLAN discovery skipped: OPNsense returned no interfaces.")
        return

    scoped_networks = [n for n in networks if should_scan_vlan(n, vlan_filter)]
    if not scoped_networks:
        logger.info("VLAN discovery skipped: no interfaces matched NMAP_DISCOVERY_VLANS.")
        return

    actions = []
    today = datetime.datetime.now().strftime("%Y.%m.%d")
    event_index = f"network-events-{today}"

    for vlan in scoped_networks:
        cidr = vlan.get("cidr")
        if not cidr:
            continue
        logger.info(f"VLAN discovery scan for {vlan.get('name')} ({cidr})")
        scan_ts = get_now_iso()
        scan_id = f"nmap_discovery_{vlan.get('name')}_{scan_ts}"
        results = run_nmap_scan([cidr], discovery_args)

        for res in results:
            ip = res.get("ip")
            if not ip:
                continue

            mac = res.get("mac")
            hostname = res.get("hostname")
            vendor = res.get("vendor")
            host_id = f"mac:{mac.lower()}" if mac else None

            event_doc = {
                "@timestamp": scan_ts,
                "source": "nmap-discovery",
                "scan_id": scan_id,
                "vlan": vlan.get("name"),
                "cidr": cidr,
                "host": {
                    "id": host_id,
                    "ip": ip,
                    "mac": mac,
                    "hostname": hostname,
                    "vendor": vendor
                }
            }
            actions.append({
                "_index": event_index,
                "_op_type": "index",
                "_source": event_doc
            })

            if host_id:
                actions.append(
                    build_discovery_update_action(host_id, mac, ip, hostname, vendor, scan_ts)
                )

    if actions:
        logger.info(f"VLAN discovery produced {len(actions)} Elasticsearch actions.")
        es.bulk_index(actions)
    else:
        logger.info("VLAN discovery finished with no hosts discovered.")

def main():
    es = get_es_client()
    opnsense_client = get_opnsense_client()
    
    interval = int(os.getenv("NMAP_INTERVAL_SECONDS", "300"))
    full_batch_size = int(os.getenv("NMAP_BATCH_SIZE", "10"))
    quick_batch_size = int(os.getenv("NMAP_QUICK_BATCH_SIZE", "30"))
    port_range = os.getenv("NMAP_PORT_RANGE", "1-1024") # Full scan range
    discovery_enabled = os.getenv("NMAP_DISCOVERY_ENABLED", "false").lower() == "true"
    discovery_interval = int(os.getenv("NMAP_DISCOVERY_INTERVAL_SECONDS", "3600"))
    discovery_vlan_filter = [v.strip() for v in os.getenv("NMAP_DISCOVERY_VLANS", "").split(",") if v.strip()]
    discovery_extra_args = os.getenv("NMAP_DISCOVERY_EXTRA_ARGS", "-sn -n").strip()
    if discovery_extra_args:
        discovery_extra_args = shlex.split(discovery_extra_args)
    else:
        discovery_extra_args = ["-sn", "-n"]
    discovery_last_run = time.time() - discovery_interval if discovery_enabled else 0.0
    full_interval = int(os.getenv("NMAP_FULL_INTERVAL_SECONDS", "86400"))
    quick_extra_str = os.getenv("NMAP_QUICK_EXTRA_ARGS", "-sS --top-ports 100 -T4 --open -Pn").strip()
    quick_extra_args = shlex.split(quick_extra_str) if quick_extra_str else ["-sS", "--top-ports", "100", "-T4", "--open", "-Pn"]
    last_full_scan = time.time()
    
    # Construct base nmap args
    # -sV for service version, -O for OS detection (requires root usually), --open to only show open
    # We run as root in docker (usually) or need capabilities.
    extra_args = ["-sV", "--open"]
    
    # Check if port_range looks like a range or specific ports
    if port_range:
        extra_args.extend(["-p", port_range])
    
    # Add user provided extra args
    user_args = os.getenv("NMAP_EXTRA_ARGS", "")
    if user_args:
        extra_args.extend(user_args.split())

    logger.info("Starting Nmap collector loop...")

    while True:
        try:
            start_time = time.time()
            ts_iso = get_now_iso()
            now = time.time()
            use_full_scan = (now - last_full_scan) >= full_interval
            scan_type = "full" if use_full_scan else "quick"
            scan_id = f"nmap_{scan_type}_{ts_iso}"
            current_batch_size = full_batch_size if use_full_scan else quick_batch_size
            scan_args = extra_args if use_full_scan else quick_extra_args

            if use_full_scan:
                last_full_scan = now
                logger.info("Running scheduled full service scan.")
            else:
                logger.info("Running quick common-port sweep.")

            if discovery_enabled and (time.time() - discovery_last_run) >= discovery_interval:
                run_vlan_discovery(es, opnsense_client, discovery_extra_args, discovery_vlan_filter)
                discovery_last_run = time.time()
            
            # 1. Get targets from ES
            # We only want hosts that have an IP.
            hosts = es.search_hosts(index="network-hosts", size=1000)
            
            # Extract IPs to scan. Map IP -> Host ID to correlate back
            targets = []
            ip_to_host_id = {}
            
            for h in hosts:
                # h is {"host": {...}, "ports": [...]}
                host_info = h.get("host", {})
                hid = host_info.get("id")
                ips = host_info.get("ips", [])
                
                if not hid or not ips:
                    continue
                    
                # Pick the "best" IP? Or scan all? 
                # Scaning all might be duplicate work if they point to same box.
                # Let's pick the first one for now.
                target_ip = ips[0] 
                targets.append(target_ip)
                ip_to_host_id[target_ip] = hid

            logger.info(f"Found {len(targets)} targets to scan ({scan_type}).")

            total_processed = 0
            logger.info(f"Scanning {scan_type} run with {len(targets)} targets.")
            scan_results = run_nmap_scan(targets, scan_args)
            actions = []
            today = datetime.datetime.now().strftime("%Y.%m.%d")
            event_index = f"network-events-{today}"

            for res in scan_results:
                ip = res.get("ip")
                if not ip or ip not in ip_to_host_id:
                    continue
                
                hid = ip_to_host_id[ip]
                total_processed += 1
                
                for p in res["ports"]:
                    p["last_seen"] = ts_iso
                    p["last_scan_id"] = scan_id
                
                event_doc = {
                    "@timestamp": ts_iso,
                    "source": "nmap",
                    "scan_id": scan_id,
                    "host": {"id": hid, "ip": ip},
                    "ports": res["ports"],
                    "os": res.get("os_match")
                }
                actions.append({
                    "_index": event_index,
                    "_op_type": "index",
                    "_source": event_doc
                })
                
                script_source = """
                        if (ctx._source.host == null) { ctx._source.host = [:]; }
                        if (ctx._source.host.sources == null) { ctx._source.host.sources = []; }
                        if (!ctx._source.host.sources.contains('nmap')) {
                            ctx._source.host.sources.add('nmap');
                        }
                        ctx._source.host.last_seen = params.ts;
                        
                        if (params.os != null) {
                            ctx._source.host.os = params.os;
                        }

                        if (ctx._source.ports == null) {
                            ctx._source.ports = [];
                        }

                        for (new_p in params.new_ports) {
                            boolean found = false;
                            for (old_p in ctx._source.ports) {
                                if (old_p.port == new_p.port && old_p.proto == new_p.proto) {
                                    old_p.last_seen = params.ts;
                                    old_p.state = new_p.state;
                                    old_p.service = new_p.service;
                                    old_p.last_scan_id = params.scan_id;
                                    found = true;
                                    break;
                                }
                            }
                            if (!found) {
                                new_p.first_seen = params.ts;
                                ctx._source.ports.add(new_p);
                            }
                        }
                    """
                    
                actions.append({
                    "_index": "network-hosts",
                    "_op_type": "update",
                    "_id": hid,
                    "script": {
                        "source": script_source,
                        "lang": "painless",
                        "params": {
                            "ts": ts_iso,
                            "os": res.get("os_match"),
                            "new_ports": res["ports"],
                            "scan_id": scan_id
                        }
                    }
                })

                for p in res["ports"]:
                    svc_id = f"{hid}:{p['proto']}:{p['port']}"
                    svc_script = """
                        ctx._source.last_seen = params.ts;
                        ctx._source.state = params.state;
                        ctx._source.service = params.service;
                        if (ctx._source.first_seen == null) {
                            ctx._source.first_seen = params.ts;
                        }
                    """
                    actions.append({
                        "_index": "network-services",
                        "_op_type": "update",
                        "_id": svc_id,
                        "script": {
                            "source": svc_script,
                            "lang": "painless",
                            "params": {
                                "ts": ts_iso,
                                "state": p["state"],
                                "service": p["service"]
                            }
                        },
                        "upsert": {
                            "host_id": hid,
                            "host_ip": ip,
                            "port": p["port"],
                            "proto": p["proto"],
                            "service": p["service"],
                            "state": p["state"],
                            "last_seen": ts_iso,
                            "first_seen": ts_iso,
                            "sources": ["nmap"]
                        }
                    })

            if actions:
                es.bulk_index(actions)

            elapsed = time.time() - start_time
            sleep_time = max(0, interval - elapsed)
            logger.info(f"Nmap {scan_type} cycle done. Scanned {total_processed} hosts in {elapsed:.2f}s. Sleeping {sleep_time:.2f}s")
            time.sleep(sleep_time)

        except Exception as e:
            logger.error(f"Error in Nmap loop: {e}")
            time.sleep(10)

if __name__ == "__main__":
    main()