#!/usr/bin/env python3
import sys
import libvirt
import time
import re
import os
import requests
import json
import psutil
from socket import gethostbyaddr, gethostname
from xmltodict import parse
from multiprocessing import Process, Manager, Queue
from queue import Empty

# Constants
CACHE_TTL_SECONDS = 3600  # 1 hour
WORKER_POLL_INTERVAL = 5  # seconds between queue checks
PROMETHEUS_TIMEOUT = 120  # seconds


def fetch_storage_metrics_sync(storage):
    """Fetch and parse ZFS metrics from a storage's Prometheus exporter (blocking)."""
    metrics_data = {}
    try:
        r = requests.get(f"http://{storage}:9134/metrics", timeout=PROMETHEUS_TIMEOUT)
        if r.status_code == 200:
            # Parse zfs_dataset_used_bytes{datacenter="dc",dataset="one/345",pool="tank"} 857088.0
            # Also handles scientific notation like 3.702859823232e+012
            pattern = re.compile(
                r'zfs_dataset_used_bytes\{[^}]*dataset="([^"]+)"[^}]*pool="([^"]+)"[^}]*\}\s+(\d+(?:\.\d+)?(?:[eE][+-]?\d+)?)'
            )
            for line in r.text.splitlines():
                m = pattern.match(line)
                if m:
                    dataset = m.group(1)
                    pool = m.group(2)
                    size_bytes = int(float(m.group(3)))
                    # Store as pool/dataset -> size
                    full_path = f"{pool}/{dataset}"
                    metrics_data[full_path] = size_bytes
    except Exception as e:
        print(f"[worker] Failed to fetch metrics from {storage}: {e}", file=sys.stderr)
    return metrics_data


def storage_worker(cache, cache_timestamps, request_queue, stop_event):
    """
    Background worker that fetches storage metrics.

    Listens to request_queue for storage names, fetches metrics from Prometheus,
    and populates the shared cache.
    """
    pending_storages = set()

    while not stop_event.is_set():
        # Drain all pending requests from queue
        try:
            while True:
                storage = request_queue.get_nowait()
                if storage and storage != "localhost":
                    pending_storages.add(storage)
        except Empty:
            pass

        # Process pending storages
        if pending_storages:
            storages_to_fetch = list(pending_storages)
            pending_storages.clear()

            for storage in storages_to_fetch:
                # Check if cache is still valid (another worker might have updated it)
                now = time.time()
                if storage in cache_timestamps:
                    if now - cache_timestamps[storage] < CACHE_TTL_SECONDS:
                        continue

                print(f"[worker] Fetching metrics from {storage}", file=sys.stderr)
                metrics_data = fetch_storage_metrics_sync(storage)

                # Update shared cache
                if metrics_data:
                    cache[storage] = metrics_data
                    cache_timestamps[storage] = now
                    print(f"[worker] Cached {len(metrics_data)} datasets from {storage}", file=sys.stderr)
                else:
                    # Store empty dict to avoid repeated requests for broken storages
                    cache[storage] = {}
                    cache_timestamps[storage] = now
                    print(f"[worker] No metrics from {storage}, caching empty", file=sys.stderr)

        # Sleep before next poll
        time.sleep(WORKER_POLL_INTERVAL)


def get_disk_size_from_cache(cache, storage, zfspath):
    """Get disk size in bytes from shared cache (non-blocking)."""
    if not storage or not zfspath or storage == "localhost":
        return None

    if storage not in cache:
        print(f"{storage} not found in cache")
        return None

    metrics = cache[storage]
    print(f"trying to get {zfspath} from metrics")
    print(metrics)
    return metrics.get(zfspath)


def normalize_zfspath(zfspath):
    """
    Normalize zfspath to match prometheus cache key format.

    Input:  'slc-stor1.slc.efscloud.net/slc-fd1-vm2-ss.b.1/one/123' (from VM disk path)
    Output: 'slc-fd1-vm2-ss.b.1/one/123' (prometheus format: pool/dataset)

    The zfspath includes the NFS hostname prefix which we need to strip.
    Hostname is detected by checking for FQDN pattern (ends with .net, .com, etc.)
    """
    if not zfspath:
        return zfspath

    # zfspath format: hostname.domain.tld/pool/dataset/...
    # We need to strip the hostname part (first component before /)
    parts = zfspath.split('/', 1)
    if len(parts) == 2:
        first = parts[0]
        # Check if first part is an FQDN (hostname.domain.tld)
        # FQDNs end with known TLDs like .net, .com, .org, etc.
        if first.endswith('.net') or first.endswith('.com') or first.endswith('.org') or first.endswith('.local'):
            return parts[1]
    return zfspath


def get_disk_or_vm_size_from_cache(cache, storage, zfspath):
    """
    Get size for persistent disk or VM dataset for non-persistent disk (non-blocking).

    Returns: (disk_size, vm_dataset_size, needs_refresh)
    - Persistent disk: (size, None, False)
    - Non-persistent disk: (None, vm_dataset_size, False)
    - Not found in cache: (None, None, True) - signals cache should be refreshed
    - Storage not cached: (None, None, False) - will be handled separately
    """
    if not storage or not zfspath or storage == "localhost":
        return None, None, False

    if storage not in cache:
        print(f"{storage} not found in cache")
        return None, None, False  # Storage not cached yet, handled by is_storage_cached

    metrics = cache[storage]

    # Normalize zfspath to match prometheus cache key format
    normalized = normalize_zfspath(zfspath)

    # Exact match = persistent disk
    if normalized in metrics:
        print(f"{normalized} found in cache")
        return metrics[normalized], None, False

    # Check for non-persistent: parent is VM dataset, child is just diskid
    parent = os.path.dirname(normalized)
    child = os.path.basename(normalized)

    if parent and child and parent in metrics:
        print(f"{parent} found in cache")
        # Non-persistent disk - return VM dataset size
        return None, metrics[parent], False
    print(f"neiter {normalized} nor {parent} not found in cache")
    print(metrics)

    # Storage is cached but this zfspath is not found - likely a new VM
    # Signal that cache should be refreshed
    return None, None, True


def is_storage_cached(cache, cache_timestamps, storage):
    """Check if storage metrics are in cache and not expired."""
    if not storage or storage == "localhost":
        return True  # localhost doesn't need caching

    if storage not in cache or storage not in cache_timestamps:
        return False

    now = time.time()
    return now - cache_timestamps[storage] < CACHE_TTL_SECONDS


def request_storage_update(request_queue, storage, cache_timestamps=None, force=False):
    """
    Add storage to the request queue for background fetching.

    If force=True and cache_timestamps is provided, invalidate the cache entry
    to ensure the worker refreshes it even if it's not expired.
    """
    if storage and storage != "localhost":
        # If forced refresh, invalidate the timestamp
        if force and cache_timestamps is not None and storage in cache_timestamps:
            try:
                del cache_timestamps[storage]
            except:
                pass  # Race condition, ignore

        try:
            request_queue.put_nowait(storage)
        except:
            pass  # Queue full, worker will catch up


def send_info(me, vmc_data_payload, servers):
    for server in servers:
        try:
            r = requests.post(
                f'https://{server}/zzinfo/api/v2/vms/{me}',
                json=vmc_data_payload,
                timeout=10
            )
            if not r.status_code == 200:
                print(r.text, file=sys.stderr)
        except Exception as e:
            print(f"Failed to send to {server}: {e}", file=sys.stderr)


def map_domstate(code):
    if code == 1:
        return "running"
    if code == 2 or code == 3 or code == 7:
        return "paused"
    if code == 4 or code == 5 or code == 6:
        return "shutoff"
    return "unknown"


def guess_product(name):
    if re.match(r'^(aa).+$', name):
        return {'os': 'windows', 'product': 'aa'}

    if re.match(r'^(sp-|spc-).+$', name):
        return {'os': 'windows', 'product': 'sp'}

    if re.match(r'^(ve|slc-ve-).+$', name):
        return {'os': 'windows', 'product': 've'}

    if re.match(r'^(dr-).+$', name):
        return {'os': 'windows', 'product': 'dr'}

    if re.match(r'(w|slc-win-|slc-addc-|slc-vac-|efs-win|efs-addc-|efs-wsus).+$', name):
        return {'os': 'windows', 'product': 'unknown'}

    if re.match(r'^cfus-prd-365-\d+$', name):
        return {'os': 'windows', 'product': 'cf'}

    if re.match(r'^rb-prd-.+\d+$', name):
        return {'os': 'linux', 'product': 'rb', 'role': 'balancer'}

    if re.match(r'^rb-.+\d+$', name):
        return {'os': 'linux', 'product': 'rb', 'role': 'vault'}

    if re.match(r'^rb-.+-mgmt$', name):
        return {'os': 'linux', 'product': 'rb', 'role': 'gmp'}

    if re.match(r'^cf(us)?-prd-index\d+$', name):
        return {'os': 'linux', 'product': 'cf', 'role': 'esnode'}

    if re.match(r'^cf(us)?-prd-index(m|p)\d+$', name):
        return {'os': 'linux', 'product': 'cf', 'role': 'esmaster'}

    if re.match(r'^(cf(us)?-).+$', name):
        return {'os': 'linux', 'product': 'cf'}

    if re.match(r'^(axc(ca|ams|atl)?-|(ylw|atl|ams|slc)-axc-).+$', name):
        return {'os': 'linux', 'product': 'axc'}

    if re.match(r'^(anc-(prd|ylw|atl|ams|slc)).+$', name):
        return {'os': 'linux', 'product': 'anc'}

    if re.match(r'^(efs-(prd|dev)-|slc-prd).+$', name):
        return {'os': 'linux', 'product': 'unknown'}

    return {'os': 'unknown', 'product': 'unknown'}


def get_blk_info(xml, dc):
    ret = {}

    # Fill target
    ret["target"] = xml['target']['@dev']

    # Fill source & storage & pool
    if xml['@type'] == 'file' and not xml['@device'] == 'cdrom':
        srcXML = xml['source']
        source = srcXML['@file']
        ret["file"] = source
        m = re.match(r'^/nfs/([^/]+)/', source)
        if m:
            ret["storage"] = m.group(1).split('.')[0]
            ret["pool"] = m.group(1)
            ret["zfspath"] = os.path.dirname(source.replace("/nfs", ""))[1:]
        else:
            ret["storage"] = "localhost"
            ret["pool"] = ""
            ret["zfspath"] = ""
    elif xml['@type'] == 'network' and xml['source']['@protocol'] == 'iscsi':
        srcXML = xml['source']
        source = srcXML['@name']
        ret["iqn"] = source
        storage = srcXML["host"]["@name"]
        if re.match(r'^\d+\.\d+\.\d+\.\d+$', storage):
            storage = str(gethostbyaddr(storage)[0])
        ret["storage"] = storage.split('.')[0]
        ret["zfspath"] = ""
        ret["pool"] = ""
    else:
        return None
    return ret


def get_vmc_details(conn, vmc_name, vms_list):
    cpu_count = os.cpu_count()
    mem_mb = psutil.virtual_memory().total // (1024*1024)

    cluster_id = None
    cluster_name = None

    if vms_list and vms_list[0].get("cluster_id") is not None:
        cluster_id = vms_list[0]["cluster_id"]
    if vms_list and vms_list[0].get("cluster") is not None:
        cluster_name = vms_list[0]["cluster"]

    return {
        "name": vmc_name,
        "cpu_count": cpu_count,
        "mem_mb": mem_mb,
        "cluster_id": cluster_id,
        "cluster_name": cluster_name
    }


def collect(conn, cache, cache_timestamps, request_queue):
    """Collect VM info from libvirt, using cached disk sizes."""
    ret = list()
    domains = conn.listAllDomains(0)

    _domain = {}

    me = conn.getHostname().split('.')[0]
    dc = me[:3]

    if len(domains) == 0:
        return ret

    for dom in domains:
        tags = {}
        xml = parse(dom.XMLDesc())
        (_domain["state"], _domain["maxmem"],
         _domain["mem"], _domain["vcpus"], _domain["cputime"]) = dom.info()
        state, reason = dom.state()
        tags["vmname"] = xml['domain']['name']
        tags["state"] = map_domstate(int(_domain["state"]))
        tags["mem_mb"] = int(_domain["mem"] / 1024)
        tags["vcpu"] = _domain["vcpus"]
        tags["datacenter"] = dc
        vmn = re.match(r'^one-(\d+)$', tags["vmname"])
        if vmn:
            return None  # Don't export ONe VMs ATM
            tags["vmname"] = xml['domain']['title']
            tags["vmid"] = int(vmn.group(1))
            tags["vmm"] = "one"
        else:
            tags["vmm"] = "spxops"
            tags["vmid"] = -1

        product_tags = guess_product(tags["vmname"])
        tags["vmc"] = me
        tags = {**tags, **product_tags}

        tags["nets"] = list()
        if 'interface' in xml['domain']['devices']:
            interfaces = xml['domain']['devices']['interface']
            i = 0
            if isinstance(interfaces, list):
                for _if in interfaces:
                    _tmp = dict()
                    _tmp["mac"] = _if["mac"]["@address"]
                    _tmp["bridge"] = _if["source"]["@bridge"]
                    _tmp["index"] = i
                    tags["nets"].append(_tmp)
                    i = i + 1
            else:
                _if = interfaces
                _tmp = dict()
                _tmp["mac"] = _if["mac"]["@address"]
                _tmp["bridge"] = _if["source"]["@bridge"]
                _tmp["index"] = i
                tags["nets"].append(_tmp)

        tags["disks"] = list()
        disks = xml['domain']['devices']['disk']
        i = 0
        disk_list = disks if isinstance(disks, list) else [disks]

        for _disk in disk_list:
            _d = get_blk_info(_disk, dc)
            if _d is not None:
                _d["index"] = i
                if "file" in _d and _d["file"].endswith("DISABLED"):
                    tags["state"] = "disabled"

                storage = _d.get("storage")
                zfspath = _d.get("zfspath")

                if storage and storage != "localhost":
                    # Check if storage is cached, request async fetch if not
                    if not is_storage_cached(cache, cache_timestamps, storage):
                        request_storage_update(request_queue, storage)

                # Get disk size from cache (non-blocking)
                disk_sz, vm_sz, needs_refresh = get_disk_or_vm_size_from_cache(cache, storage, zfspath)
                if disk_sz is not None:
                    _d["sz"] = disk_sz
                elif vm_sz is not None:
                    tags["sz"] = vm_sz
                elif needs_refresh:
                    # Disk not found in cache - likely a new VM, force refresh
                    request_storage_update(request_queue, storage, cache_timestamps, force=True)

                tags["disks"].append(_d)
            i = i + 1

        ret.append(tags)

    return ret


def is_opennebula():
    with open("/etc/libvirt/qemu.conf", "r") as f:
        for line in f:
            if line.startswith("user"):
                u = line.split("=")[1]
                if u.strip("\r\n\t \"'") == "oneadmin":
                    return True
    return False


def discovery(service: str, domain):
    from dns import resolver

    a = resolver.resolve(f'_{service}._tcp.{domain}', 'SRV')

    for i in a:
        yield i.target.to_text().rstrip('.')  # , i.port)


def main_collect(cache, cache_timestamps, request_queue):
    """Main collection routine that uses shared cache."""
    conn = None
    try:
        conn = libvirt.open("qemu:///system")
    except libvirt.libvirtError as e:
        print(repr(e), file=sys.stderr)
        return

    me = conn.getHostname().split('.')[0]
    dc = me[:3]
    servers = list(discovery("zzinfo", f"{dc}.efscloud.net"))

    if is_opennebula():
        try:
            r = requests.get(f"https://vmm.{dc}.efscloud.net/axcapi/v2/vmc/e/{me}", timeout=10)
            if r.status_code == 200:
                ret = json.loads(r.text)

                # Populate disk sizes from shared cache
                for vm in ret:
                    for disk in vm.get("disks", []):
                        storage = disk.get("storage")
                        zfspath = disk.get("zfspath")

                        if not storage or storage == "localhost":
                            continue

                        # Check if storage is cached, request async fetch if not
                        if not is_storage_cached(cache, cache_timestamps, storage):
                            request_storage_update(request_queue, storage)

                        # Get from cache (non-blocking)
                        disk_sz, vm_sz, needs_refresh = get_disk_or_vm_size_from_cache(cache, storage, zfspath)
                        if disk_sz is not None:
                            disk["sz"] = disk_sz
                        elif vm_sz is not None:
                            vm["sz"] = vm_sz
                        elif needs_refresh:
                            # Disk not found in cache - likely a new VM, force refresh
                            request_storage_update(request_queue, storage, cache_timestamps, force=True)

                vmc_details = get_vmc_details(conn, me, ret)
                vmc_data_payload = {"vmc_info": vmc_details, "vms": ret}
                send_info(me, vmc_data_payload, servers)
            else:
                print(r.text, file=sys.stderr)
        except Exception as e:
            print(f"OpenNebula fetch error: {e}", file=sys.stderr)
    else:
        ret = collect(conn, cache, cache_timestamps, request_queue)
        if ret is not None:
            vmc_details = get_vmc_details(conn, me, ret)
            vmc_data_payload = {"vmc_info": vmc_details, "vms": ret}
            send_info(me, vmc_data_payload, servers)


def alert(alertname):
    me = gethostname().split('.')[0]
    dc = me[:3]
    alerts = [{
        "labels": {
            "alertname": alertname,
            "severity": "info",
            "host": me,
            "team": "ceunix",
            "product": "zzinfo",
            "datacenter": dc,
        }
    }]
    try:
        requests.post(
            f"http://{dc}-prd-alerts1.{dc}.efscloud.net:9093/api/v2/alerts",
            json=alerts,
            timeout=5
        )
    except:
        pass


def run_daemon():
    """
    Main daemon loop with background worker for storage metrics.

    Architecture:
    - Main loop: collects VM info every 30s, reads disk sizes from cache
    - Background worker: listens to request queue, fetches storage metrics, populates cache
    - Cache TTL: 1 hour
    """
    # Create shared state using Manager
    manager = Manager()
    cache = manager.dict()  # storage_name -> {zfspath: size_bytes}
    cache_timestamps = manager.dict()  # storage_name -> timestamp
    request_queue = manager.Queue(maxsize=300)
    stop_event = manager.Event()

    # Start background worker
    worker = Process(target=storage_worker, args=(cache, cache_timestamps, request_queue, stop_event))
    worker.daemon = True
    worker.start()
    print("[main] Started storage metrics worker", file=sys.stderr)

    try:
        while True:
            # Run collection in subprocess with timeout
            p = Process(target=main_collect, args=(cache, cache_timestamps, request_queue))
            p.start()
            p.join(45)

            if p.is_alive():
                p.terminate()
                p.join(5)  # Give it a moment to terminate
                alert("zzagent.collect.timeout")
                print("[main] Collection timeout, process terminated", file=sys.stderr)

            time.sleep(30)
    except KeyboardInterrupt:
        print("[main] Shutting down...", file=sys.stderr)
    finally:
        stop_event.set()
        worker.join(timeout=5)
        if worker.is_alive():
            worker.terminate()


if __name__ == "__main__":
    run_daemon()
