"""
JEDI Comprehensive Benchmarks — Performance, throughput, stress tests.
"""

import sys, os, time, json, random, hashlib, statistics, threading
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from jedi.core.engine import JEDIEngine, ThreatLevel
from jedi.core.mission import Mission, MissionStatus
from jedi.core.nanobot import Nanobot, NanobotType, NanobotState
from jedi.legal.gate import LegalGate, AuthorizationLevel
from jedi.swarm.coordinator import SwarmCoordinator, ConsensusProtocol, SharedWisdom
from jedi.command.control import MissionControl
from jedi.comms.channel import CommsChannel
from jedi.memory.helix import HelixMemoryStore
from jedi.nanobots.red.striker import StrikerNanobot
from jedi.nanobots.blue.guardian import GuardianNanobot
from jedi.nanobots.black.ghost import GhostNanobot
from jedi.nanobots.purple.synapse import SynapseNanobot
from jedi.nanobots.yellow.auditor import AuditorNanobot
from jedi.nanobots.green.architect import ArchitectNanobot


class Benchmark:
    def __init__(self):
        self.results = {}
    
    def time_it(self, name, func, iterations=100):
        times = []
        for _ in range(iterations):
            start = time.perf_counter()
            result = func()
            end = time.perf_counter()
            times.append((end - start) * 1000)  # ms
        
        self.results[name] = {
            "iterations": iterations,
            "min_ms": round(min(times), 3),
            "max_ms": round(max(times), 3),
            "avg_ms": round(statistics.mean(times), 3),
            "median_ms": round(statistics.median(times), 3),
            "p99_ms": round(sorted(times)[int(len(times) * 0.99)], 3),
            "total_ms": round(sum(times), 3),
            "ops_per_sec": round(1000 / statistics.mean(times), 1),
        }
        return self.results[name]
    
    def report(self):
        print(f"\n{'='*70}")
        print(f"  JEDI BENCHMARK RESULTS")
        print(f"{'='*70}")
        print(f"  {'Test':<45} {'ops/s':>8} {'avg ms':>8} {'p99 ms':>8}")
        print(f"  {'-'*45} {'-'*8} {'-'*8} {'-'*8}")
        for name, data in self.results.items():
            print(f"  {name:<45} {data['ops_per_sec']:>8.1f} {data['avg_ms']:>8.3f} {data['p99_ms']:>8.3f}")
        print(f"{'='*70}\n")
        return self.results


def bench_engine_init(b):
    def fn():
        JEDIEngine()
    return b.time_it("Engine initialization", fn, 100)

def bench_threat_assessment(b):
    engine = JEDIEngine()
    def fn():
        engine.assess_threat({
            "indicators": [{"type": "lateral_movement", "severity": 10}],
            "confidence": 0.5
        })
    return b.time_it("Threat assessment", fn, 500)

def bench_situation_report(b):
    engine = JEDIEngine()
    def fn():
        engine.get_situation_report()
    return b.time_it("Situation report", fn, 500)

def bench_nanobot_create(b):
    def fn():
        Nanobot(NanobotType.GHOST, "bench-mission", {"address": "10.0.0.1"})
    return b.time_it("Nanobot creation", fn, 200)

def bench_nanobot_heartbeat(b):
    bot = Nanobot(NanobotType.SENTINEL, "bench-mission", {"address": "10.0.0.1"})
    def fn():
        bot.heartbeat()
    return b.time_it("Nanobot heartbeat", fn, 1000)

def bench_striker_exploit(b):
    s = StrikerNanobot("bench", {"address": "10.0.0.1"})
    def fn():
        s.exploit({"type": "sql_injection", "severity": 8})
    return b.time_it("Striker exploit", fn, 200)

def bench_guardian_detect(b):
    g = GuardianNanobot("bench", {"address": "10.0.0.1"})
    baseline = {"network_traffic": 100, "failed_auth": 0, "processes": ["nginx"]}
    current = {"network_traffic": 500, "failed_auth": 20, "processes": ["nginx", "nc", "wget"]}
    def fn():
        g.detect_anomaly(baseline, current)
    return b.time_it("Guardian anomaly detect", fn, 500)

def bench_ghost_infiltrate(b):
    g = GhostNanobot("bench", {"address": "10.0.0.1"})
    def fn():
        g.infiltrate({"method": "supply_chain", "os": "linux"})
    return b.time_it("Ghost infiltration", fn, 200)

def bench_ghost_attribution(b):
    g = GhostNanobot("bench", {"address": "10.0.0.1"})
    def fn():
        g.collect_attribution({
            "address": "10.0.0.1", "username": "root", "timezone": "US/Eastern",
            "installed_tools": ["cobalt_strike", "mimikatz"], "c2_beacons": ["cobalt_strike"]
        })
    return b.time_it("Ghost attribution", fn, 200)

def bench_auditor_compliance(b):
    a = AuditorNanobot("bench", {"address": "10.0.0.1"})
    def fn():
        a.check_compliance("nist_csf", {"logging_active": True, "mfa_enabled": True})
    return b.time_it("Auditor compliance check", fn, 500)

def bench_architect_review(b):
    a = ArchitectNanobot("bench", {"address": "10.0.0.1"})
    code = "import os\nos.system(user_input)\neval(code)\nimport pickle\npickle.loads(data)"
    def fn():
        a.review_code(code, "python")
    return b.time_it("Architect code review", fn, 300)

def bench_synapse_coordinate(b):
    s = SynapseNanobot("bench", {"address": "10.0.0.1"})
    def fn():
        s.coordinate_attack_defense(
            {"steps": ["exploit", "escalate", "exfiltrate"]},
            {"steps": ["detect", "contain", "recover"]}
        )
    return b.time_it("Synapse coordination", fn, 300)

def bench_legal_auth_create(b):
    gate = LegalGate()
    def fn():
        gate.create_authorization(
            "OP-BENCH", AuthorizationLevel.OFFEND, "US-FEDERAL",
            "AUTH-BENCH", "pentest", {"address": "192.168.1.1"}
        )
    return b.time_it("Legal auth creation", fn, 500)

def bench_legal_verify(b):
    gate = LegalGate()
    auth = gate.create_authorization(
        "OP-BENCH", AuthorizationLevel.OFFEND, "US-FEDERAL",
        "AUTH-BENCH", "pentest", {"address": "192.168.1.1"}
    )
    mission_config = {
        "type": "pentest", "target": {"address": "192.168.1.1"},
        "authorization": auth,
        "roe": {"proportionality_acknowledged": True, "distinction_acknowledged": True, "necessity_acknowledged": True}
    }
    def fn():
        gate.verify_authorization(mission_config)
    return b.time_it("Legal verification", fn, 1000)

def bench_ledger_log(b):
    engine = JEDIEngine()
    def fn():
        engine.ledger.log("bench_event", {"data": "test"})
    return b.time_it("Ledger log entry", fn, 2000)

def bench_ledger_verify(b):
    engine = JEDIEngine()
    for i in range(100):
        engine.ledger.log(f"event_{i}", {"data": i})
    def fn():
        engine.ledger.verify_integrity()
    return b.time_it("Ledger verify (100 entries)", fn, 200)

def bench_swarm_add(b):
    swarm = SwarmCoordinator("bench-swarm")
    counter = [0]
    def fn():
        counter[0] += 1
        swarm.add_member(f"BENCH-{counter[0]:06d}", "scout")
    return b.time_it("Swarm member add", fn, 500)

def bench_swarm_task(b):
    swarm = SwarmCoordinator("bench-swarm")
    for i in range(10):
        swarm.add_member(f"B-{i}", ["scout","guardian","striker"][i%3])
    def fn():
        swarm.queue_task({"type": "recon", "target": "10.0.0.1"})
    return b.time_it("Swarm task queue", fn, 500)

def bench_consensus(b):
    cp = ConsensusProtocol(threshold=0.6)
    proposals = [
        {"bot_id": f"B-{i}", "action": random.choice(["recon", "defend", "attack"]),
         "confidence": random.uniform(0.5, 1.0)}
        for i in range(5)
    ]
    def fn():
        cp.reach_consensus(proposals, min_voters=3)
    return b.time_it("Consensus protocol", fn, 1000)

def bench_comms_c2(b):
    ch = CommsChannel("bench")
    def fn():
        ch.send_c2("GHOST-001", {"intel": "test_data"})
    return b.time_it("Comms C2 send", fn, 1000)

def bench_comms_p2p(b):
    ch = CommsChannel("bench")
    def fn():
        ch.send_p2p("BOT-001", "BOT-002", {"data": "test"})
    return b.time_it("Comms P2P send", fn, 1000)

def bench_memory_store(b):
    mem = HelixMemoryStore("/tmp/jedi_bench_memory")
    counter = [0]
    def fn():
        counter[0] += 1
        mem.store(f"key_{counter[0]}", {"data": f"value_{counter[0]}"})
    return b.time_it("Memory store", fn, 1000)

def bench_memory_retrieve(b):
    mem = HelixMemoryStore("/tmp/jedi_bench_memory")
    for i in range(1000):
        mem.store(f"key_{i}", {"data": i})
    def fn():
        mem.retrieve(f"key_{random.randint(0, 999)}")
    return b.time_it("Memory retrieve", fn, 2000)

def bench_memory_search(b):
    mem = HelixMemoryStore("/tmp/jedi_bench_memory")
    for i in range(100):
        mem.store(f"threat_{i}", {"description": f"threat pattern {i}"}, "semantic")
    def fn():
        mem.search("threat")
    return b.time_it("Memory search (100 entries)", fn, 200)

def bench_mission_create(b):
    mc = MissionControl()
    def fn():
        mc.new_mission(
            f"Bench-{random.randint(0,9999)}", "recon",
            {"address": "10.0.0.1"}, ["test"],
            "OP-BENCH", "US-FEDERAL", "AUTH-BENCH"
        )
    return b.time_it("Mission creation", fn, 100)

def bench_dashboard(b):
    mc = MissionControl()
    def fn():
        mc.dashboard()
    return b.time_it("Dashboard render", fn, 500)

def bench_swarm_intel_collect(b):
    swarm = SwarmCoordinator("bench")
    for i in range(20):
        swarm.add_member(f"B-{i}", "scout")
    def fn():
        swarm.collect_intel(f"B-{random.randint(0,19)}", {"type": "scan", "result": "open_port"})
    return b.time_it("Swarm intel collection", fn, 1000)

def bench_pheromone(b):
    sw = SharedWisdom()
    counter = [0]
    def fn():
        counter[0] += 1
        sw.add_pheromone(f"trail_{counter[0] % 10}", 1.0)
    return b.time_it("Pheromone trail", fn, 2000)

def bench_batch_nanobot_lifecycle(b):
    """Full lifecycle: create -> heartbeat -> recon -> report -> self-destruct"""
    def fn():
        bot = Nanobot(NanobotType.SENTINEL, "bench", {"address": "10.0.0.1"})
        bot.heartbeat()
        bot.recon({"open_ports": [80, 443], "services": ["nginx"]})
        bot.report()
        bot.self_destruct()
    return b.time_it("Full nanobot lifecycle", fn, 200)

def bench_concurrent_operations(b):
    """Test thread safety under concurrent load"""
    engine = JEDIEngine()
    results = []
    
    def worker():
        for _ in range(10):
            engine.assess_threat({"indicators": [{"type": "suspicious_process", "severity": 5}], "confidence": 0.5})
            engine.ledger.log("concurrent_event", {"thread": threading.current_thread().name})
    
    def fn():
        threads = [threading.Thread(target=worker) for _ in range(10)]
        for t in threads: t.start()
        for t in threads: t.join()
    return b.time_it("Concurrent ops (10 threads)", fn, 50)


if __name__ == "__main__":
    b = Benchmark()
    
    print("Running JEDI benchmarks...")
    print("(This may take a minute)")
    print()
    
    # Core system
    bench_engine_init(b)
    bench_threat_assessment(b)
    bench_situation_report(b)
    
    # Nanobot ops
    bench_nanobot_create(b)
    bench_nanobot_heartbeat(b)
    bench_batch_nanobot_lifecycle(b)
    
    # Color team ops
    bench_striker_exploit(b)
    bench_guardian_detect(b)
    bench_ghost_infiltrate(b)
    bench_ghost_attribution(b)
    bench_auditor_compliance(b)
    bench_architect_review(b)
    bench_synapse_coordinate(b)
    
    # Legal
    bench_legal_auth_create(b)
    bench_legal_verify(b)
    
    # Ledger
    bench_ledger_log(b)
    bench_ledger_verify(b)
    
    # Swarm
    bench_swarm_add(b)
    bench_swarm_task(b)
    bench_consensus(b)
    bench_swarm_intel_collect(b)
    bench_pheromone(b)
    
    # Comms
    bench_comms_c2(b)
    bench_comms_p2p(b)
    
    # Memory
    bench_memory_store(b)
    bench_memory_retrieve(b)
    bench_memory_search(b)
    
    # Mission Control
    bench_mission_create(b)
    bench_dashboard(b)
    
    # Stress
    bench_concurrent_operations(b)
    
    results = b.report()
    
    # Save results
    with open("benchmark_results.json", "w") as f:
        json.dump(results, f, indent=2)
    print("Results saved to benchmark_results.json")
    
    # Summary
    total_ops = sum(r["ops_per_sec"] for r in results.values())
    print(f"\nTotal aggregate throughput: {total_ops:,.0f} operations/sec")
    print(f"Fastest: {min(results.values(), key=lambda x: x['avg_ms'])['avg_ms']:.3f} ms")
    print(f"Slowest: {max(results.values(), key=lambda x: x['avg_ms'])['avg_ms']:.3f} ms")