#!/usr/bin/env python3
import argparse
import json
import math
import os
import random
import sys
import time

STATE_DIR = "/tmp/nems_demo_state"

def get_args():
    parser = argparse.ArgumentParser(description="NEMS Linux Demo Plugin - Dynamic Wave Engine")
    parser.add_argument("-H", "--host", default=os.getenv("NAGIOS_HOSTNAME", "demo-host"), help="Target Hostname")
    parser.add_argument("-t", "--type", choices=["host", "disk", "cpu", "memory", "ping", "http"], default="disk", help="Check scenario type")
    parser.add_argument("--force", choices=["ok", "warning", "critical"], help="Force specific state for live presentation")
    parser.add_argument("--reset", action="store_true", help="Reset state file to OK")
    return parser.parse_args()

def load_state(filepath):
    if os.path.exists(filepath):
        try:
            with open(filepath, "r") as f:
                return json.load(f)
        except Exception:
            pass
    return {}

def save_state(filepath, data):
    try:
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, "w") as f:
            json.dump(data, f)
    except Exception:
        pass

def get_fleet_target(now):
    # Oscillates target alert cap between 0 and 4 over a 45-minute sine wave cycle
    # Guarantees periods of 100% green quiet, followed by light issue clusters
    cycle = (now / 2700.0) * 2 * math.pi
    wave = (math.sin(cycle) + 1) / 2
    return int(round(wave * 4))

def count_active_incidents(now):
    active = 0
    if not os.path.exists(STATE_DIR):
        return 0
    for fname in os.listdir(STATE_DIR):
        if fname.endswith(".json"):
            try:
                with open(os.path.join(STATE_DIR, fname), "r") as f:
                    data = json.load(f)
                    inc = data.get("incident")
                    if inc and (now - inc["start_time"] < inc["duration"]):
                        active += 1
            except Exception:
                pass
    return active

def main():
    args = get_args()
    safe_host = "".join(c for c in args.host if c.isalnum() or c in ("-", "_"))
    state_file = os.path.join(STATE_DIR, f"{safe_host}_{args.type}.json")

    if args.reset:
        if os.path.exists(state_file):
            try:
                os.remove(state_file)
            except OSError:
                pass
        print(f"DEMO RESET - State cleared for {args.host} ({args.type}) | state=0")
        sys.exit(0)

    if args.force:
        status_map = {"ok": 0, "warning": 1, "critical": 2}
        status_level = status_map[args.force]
    else:
        now = time.time()
        state = load_state(state_file)
        incident = state.get("incident")

        # Resolve expired incident
        if incident and (now - incident["start_time"] >= incident["duration"]):
            incident = None

        if not incident:
            active_alerts = count_active_incidents(now)
            target_cap = get_fleet_target(now)
            
            # 1. Standard Incident Roll (respects dynamic wave ceiling)
            if active_alerts < target_cap and random.random() < 0.015:
                incident = {
                    "start_time": now,
                    "duration": random.randint(360, 900), # 6 to 15 minutes
                    "profile": "standard"
                }
            # 2. Surprise Flash Spike Roll (0.3% independent chance, ignores wave ceiling)
            elif random.random() < 0.003:
                incident = {
                    "start_time": now,
                    "duration": random.randint(90, 240), # Quick 1.5 to 4 minute burst
                    "profile": "flash"
                }
            # 3. Rare Flapping Incident Roll (0.2% independent chance)
            elif random.random() < 0.002:
                incident = {
                    "start_time": now,
                    "duration": random.randint(180, 360), # 3 to 6 minute oscillation
                    "profile": "flapping"
                }

        save_state(state_file, {"incident": incident, "last_check": now})

        # Calculate status based on active incident profile
        if not incident:
            status_level = 0
        else:
            elapsed = now - incident["start_time"]
            progress = elapsed / incident["duration"]
            profile = incident.get("profile", "standard")

            if profile == "flash":
                # Instant spike to CRITICAL, brief WARNING before clearing
                status_level = 2 if progress < 0.75 else 1
            elif profile == "flapping":
                # Rapidly toggles between WARNING and CRITICAL every 30 seconds
                status_level = 2 if int(elapsed // 30) % 2 == 0 else 1
            else:
                # Standard curve: OK -> WARNING -> CRITICAL -> WARNING -> OK
                if progress < 0.20 or progress > 0.80:
                    status_level = 1
                else:
                    status_level = 2

    # Nagios Host Check Scenario
    if args.type == "host":
        if status_level == 0:
            print(f"OK - {args.host} is UP (0.4ms response) | rta=0.4ms;200;500;0")
            sys.exit(0)
        else:
            print(f"CRITICAL - {args.host} Host Unreachable (100% packet loss) | rta=0ms;200;500;0")
            sys.exit(1)

    # Service Check Scenarios
    if args.type == "disk":
        if status_level == 0:
            used = random.randint(35, 65)
            msg = f"DISK OK - free space: / {100-used} GB ({used}% used) | /= {used}GB;80;90;0;100"
        elif status_level == 1:
            used = random.randint(81, 88)
            msg = f"DISK WARNING - free space: / {100-used} GB ({used}% used) | /= {used}GB;80;90;0;100"
        else:
            used = random.randint(93, 98)
            msg = f"DISK CRITICAL - free space: / {100-used} GB ({used}% used) | /= {used}GB;80;90;0;100"

    elif args.type == "cpu":
        if status_level == 0:
            l1, l5, l15 = round(random.uniform(0.1, 1.2), 2), round(random.uniform(0.1, 1.5), 2), round(random.uniform(0.1, 1.5), 2)
            msg = f"OK - load average: {l1:.2f}, {l5:.2f}, {l15:.2f} | load1={l1:.2f};5.00;10.00;0 load5={l5:.2f};4.00;8.00;0 load15={l15:.2f};3.00;6.00;0"
        elif status_level == 1:
            l1, l5, l15 = round(random.uniform(5.1, 7.8), 2), round(random.uniform(4.1, 6.0), 2), round(random.uniform(3.1, 4.5), 2)
            msg = f"WARNING - load average: {l1:.2f}, {l5:.2f}, {l15:.2f} | load1={l1:.2f};5.00;10.00;0 load5={l5:.2f};4.00;8.00;0 load15={l15:.2f};3.00;6.00;0"
        else:
            l1, l5, l15 = round(random.uniform(11.5, 18.2), 2), round(random.uniform(8.5, 14.0), 2), round(random.uniform(6.5, 10.0), 2)
            msg = f"CRITICAL - load average: {l1:.2f}, {l5:.2f}, {l15:.2f} | load1={l1:.2f};5.00;10.00;0 load5={l5:.2f};4.00;8.00;0 load15={l15:.2f};3.00;6.00;0"

    elif args.type == "memory":
        total_mb = 16384
        if status_level == 0:
            used_mb = random.randint(4000, 7500)
            pct = int((used_mb / total_mb) * 100)
            msg = f"OK - Memory usage {pct}% ({used_mb} MB / {total_mb} MB) | total={total_mb}MB used={used_mb}MB;13107;14745;0;{total_mb}"
        elif status_level == 1:
            used_mb = random.randint(13200, 14500)
            pct = int((used_mb / total_mb) * 100)
            msg = f"WARNING - Memory usage {pct}% ({used_mb} MB / {total_mb} MB) | total={total_mb}MB used={used_mb}MB;13107;14745;0;{total_mb}"
        else:
            used_mb = random.randint(15000, 16100)
            pct = int((used_mb / total_mb) * 100)
            msg = f"CRITICAL - Memory usage {pct}% ({used_mb} MB / {total_mb} MB) | total={total_mb}MB used={used_mb}MB;13107;14745;0;{total_mb}"

    elif args.type == "ping":
        if status_level == 0:
            rta, pl = round(random.uniform(2.0, 12.0), 1), 0
            msg = f"OK - Packet loss = {pl}%, RTA = {rta} ms | rta={rta}ms;100.0;300.0;0 pl={pl}%;20;50;0;100"
        elif status_level == 1:
            rta, pl = round(random.uniform(105.0, 220.0), 1), random.choice([0, 5, 10])
            msg = f"WARNING - Packet loss = {pl}%, RTA = {rta} ms | rta={rta}ms;100.0;300.0;0 pl={pl}%;20;50;0;100"
        else:
            rta, pl = round(random.uniform(350.0, 850.0), 1), random.choice([25, 50, 75])
            msg = f"CRITICAL - Packet loss = {pl}%, RTA = {rta} ms | rta={rta}ms;100.0;300.0;0 pl={pl}%;20;50;0;100"

    elif args.type == "http":
        if status_level == 0:
            resp = round(random.uniform(0.02, 0.18), 3)
            msg = f"HTTP OK: HTTP/1.1 200 OK - {resp} second response time | time={resp}s;2.00;5.00;0"
        elif status_level == 1:
            resp = round(random.uniform(2.10, 4.80), 3)
            msg = f"HTTP WARNING: HTTP/1.1 200 OK - {resp} second response time | time={resp}s;2.00;5.00;0"
        else:
            resp = round(random.uniform(5.10, 9.90), 3)
            msg = f"HTTP CRITICAL: HTTP/1.1 503 Service Unavailable - {resp} second response time | time={resp}s;2.00;5.00;0"

    print(msg)
    sys.exit(status_level)

if __name__ == "__main__":
    main()
