#!/usr/bin/env python3
"""
Axiom Zero — Real-World High-Concurrency Raspberry Pi 5 Continuous Stress Test Harness
Supports CLI arguments: --url, --concurrency, --duration (up to 86400s / 24h)
Runs continuously for the EXACT duration requested by the user.
"""

import argparse
import asyncio
import time
import json
import statistics
import aiohttp
import sys
from typing import List, Dict, Any

PAYLOAD_TEMPLATES = [
    {"customer_id": "cust_stress_test", "is_bot": False, "risk_score": 0.02, "user_agent": "Mozilla/5.0 (X11; Linux x86_64)"},
    {"customer_id": "cust_stress_test", "is_bot": True, "bot_tier": "Tier 1: cURL", "user_agent": "curl/7.88.1"},
    {"customer_id": "cust_stress_test", "is_bot": True, "bot_tier": "Tier 2: Headless", "cdp_detected": True},
    {"customer_id": "cust_stress_test", "is_bot": True, "bot_tier": "Tier 3: Stealth", "signals": {"automation": {"webdriver": True}}},
]

async def worker(worker_id: int, target_url: str, end_time: float, session: aiohttp.ClientSession, results: list, stop_event: asyncio.Event):
    request_idx = 0
    while time.perf_counter() < end_time and not stop_event.is_set():
        payload = PAYLOAD_TEMPLATES[(worker_id + request_idx) % len(PAYLOAD_TEMPLATES)]
        request_idx += 1
        start = time.perf_counter()
        try:
            async with session.post(target_url, json=payload, timeout=aiohttp.ClientTimeout(total=10.0)) as resp:
                status = resp.status
                await resp.text()
                elapsed = (time.perf_counter() - start) * 1000.0
                results.append((status == 200, elapsed, status))
        except Exception as e:
            elapsed = (time.perf_counter() - start) * 1000.0
            results.append((False, elapsed, str(e)))

async def main():
    parser = argparse.ArgumentParser(description="Axiom Zero Continuous Stress Harness")
    parser.add_argument("--url", default="http://127.0.0.1:8085/v1/score", help="Target URL")
    parser.add_argument("--concurrency", type=int, default=100, help="Number of concurrent workers")
    parser.add_argument("--duration", type=int, default=30, help="Duration in seconds (up to 86400s)")
    args = parser.parse_args()

    target_url = args.url
    concurrency = max(1, min(2000, args.concurrency))
    duration = max(1, min(86400, args.duration))

    print(f"\n========================================================")
    print(f"🔥 CONTINUOUS STRESS TEST LAUNCHED")
    print(f" Target:      {target_url}")
    print(f" Concurrency: {concurrency} Workers")
    print(f" Duration:    {duration:,} Seconds ({duration/3600:.2f} Hours)")
    print(f"========================================================\n")

    results = []
    stop_event = asyncio.Event()
    start_time = time.perf_counter()
    end_time = start_time + duration

    connector = aiohttp.TCPConnector(limit=concurrency * 2, force_close=False, enable_cleanup_closed=True)
    async with aiohttp.ClientSession(connector=connector) as session:
        workers = [worker(w_id, target_url, end_time, session, results, stop_event) for w_id in range(concurrency)]
        
        # Periodic report loop while running
        report_task = asyncio.create_task(periodic_report(start_time, duration, results, stop_event))
        
        await asyncio.gather(*workers, return_exceptions=True)
        stop_event.set()
        report_task.cancel()

    actual_duration = time.perf_counter() - start_time
    total_reqs = len(results)
    successes = [r[1] for r in results if r[0]]
    failures = [r for r in results if not r[0]]

    rps = total_reqs / actual_duration if actual_duration > 0 else 0
    success_rate = (len(successes) / total_reqs * 100.0) if total_reqs > 0 else 0.0
    mean_lat = statistics.mean(successes) if successes else 0.0
    
    if successes:
        s_sorted = sorted(successes)
        p50_lat = s_sorted[int(len(s_sorted) * 0.50)]
        p95_lat = s_sorted[int(len(s_sorted) * 0.95)]
        p99_lat = s_sorted[int(len(s_sorted) * 0.99)]
    else:
        p50_lat = p95_lat = p99_lat = 0.0

    print("\n========================================================")
    print("               FINAL STRESS TEST SUMMARY                ")
    print("========================================================")
    print(f" Duration Actual:   {actual_duration:.2f}s")
    print(f" Total Requests:    {total_reqs:,}")
    print(f" Throughput (RPS):  {rps:.2f} req/sec")
    print(f" Success Rate:      {success_rate:.2f}%")
    print(f" Latency Mean:      {mean_lat:.2f}ms")
    print(f" Latency P50:       {p50_lat:.2f}ms")
    print(f" Latency P95:       {p95_lat:.2f}ms")
    print(f" Latency P99:       {p99_lat:.2f}ms")
    print("========================================================\n")

async def periodic_report(start_time: float, target_duration: int, results: list, stop_event: asyncio.Event):
    last_count = 0
    while not stop_event.is_set():
        await asyncio.sleep(5)
        elapsed = time.perf_counter() - start_time
        curr_count = len(results)
        delta_reqs = curr_count - last_count
        last_count = curr_count
        interval_rps = delta_reqs / 5.0
        pct_done = min(100.0, (elapsed / target_duration) * 100.0)
        print(f" ⏱️ [{elapsed:6.1f}s / {target_duration}s ({pct_done:5.1f}%)] RPS: {interval_rps:7.1f} | Total Reqs: {curr_count:,}")

if __name__ == "__main__":
    asyncio.run(main())