import os
import sys
import time
import json
import random
import unittest
import numpy as np
from pathlib import Path
from unittest.mock import MagicMock

# Correct sys.path
ROOT_DIR = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT_DIR / "nautilus_dolphin"))
sys.path.insert(0, str(ROOT_DIR))

from nautilus_dolphin.nautilus.ob_features import OBFeatureEngine, NEUTRAL_PLACEMENT
from nautilus_dolphin.nautilus.ob_provider import OBSnapshot

# Mock memory monitor
try:
    import psutil
    _HAS_PSUTIL = True
except ImportError:
    _HAS_PSUTIL = False

class FastMockProvider:
    def __init__(self):
        self.snap = None
    def get_snapshot(self, asset, ts):
        return self.snap
    def get_assets(self):
        return ["BTCUSDT", "ETHUSDT", "SOLUSDT", "BNBUSDT", "XRPUSDT"]

class ExtremeStressTester:
    def __init__(self, iterations=180000, fuzz_rate=0.05):
        self.iterations = iterations
        self.fuzz_rate = fuzz_rate
        self.assets = ["BTCUSDT", "ETHUSDT", "SOLUSDT", "BNBUSDT", "XRPUSDT"]
        self.provider = FastMockProvider()
        self.engine = OBFeatureEngine(self.provider)
        
    def create_valid_snap(self, asset):
        return OBSnapshot(
            timestamp=time.time(),
            asset=asset,
            bid_notional=np.random.rand(5) * 10000,
            ask_notional=np.random.rand(5) * 10000,
            bid_depth=np.random.rand(5),
            ask_depth=np.random.rand(5)
        )

    def create_fuzzed_snap(self, asset):
        # Return various 'bad' flavors
        choice = random.choice(["NONE", "EMPTY_LISTS", "NAN_VALS", "WRONG_TYPE"])
        if choice == "NONE":
            return None
        if choice == "EMPTY_LISTS":
            return OBSnapshot(time.time(), asset, np.array([]), np.array([]), np.array([]), np.array([]))
        if choice == "NAN_VALS":
            return OBSnapshot(time.time(), asset, np.array([np.nan]*5), np.array([np.nan]*5), np.array([0.0]*5), np.array([0.0]*5))
        if choice == "WRONG_TYPE":
            # Test shape mismatch or zeroing
            return OBSnapshot(time.time(), asset, np.array([1.0, 2.0]), np.array([1.0]), np.array([0]), np.array([0]))
        return None

    def run(self):
        # 180,000 iterations = 30 minutes of 100Hz
        self.iterations = 180000 
        print(f"STARTING EXTREME STRESS: {self.iterations} iterations, {self.fuzz_rate*100}% fuzz rate")
        
        process = psutil.Process(os.getpid()) if _HAS_PSUTIL else None
        mem_start = process.memory_info().rss / (1024*1024) if process else 0
        
        start_time = time.time()
        fuzz_counts = 0
        success_counts = 0
        
        for i in range(self.iterations):
            for asset in self.assets:
                if random.random() < self.fuzz_rate:
                    snap = self.create_fuzzed_snap(asset)
                    fuzz_counts += 1
                else:
                    snap = self.create_valid_snap(asset)
                
                self.provider.snap = snap
                
            try:
                self.engine.step_live(self.assets, i)
                success_counts += 1
            except Exception as e:
                # Malformed data might trigger errors in kernels if not jitted with guards
                # but we handle it gracefully
                fuzz_counts += 1
                
            if i % 20000 == 0 and i > 0:
                mem_now = process.memory_info().rss / (1024*1024) if process else 0
                print(f" - Progress: {i}/{self.iterations} | Mem: {mem_now:.2f}MB | Fuzz: {fuzz_counts}")

        end_time = time.time()
        duration = end_time - start_time
        mem_end = process.memory_info().rss / (1024*1024) if process else 0
        
        print("\n" + "="*60)
        print("EXTREME STRESS COMPLETED SUCCESSFULY")
        print(f"Duration: {duration:.2f}s (~{self.iterations/duration:.2f} Hz)")
        print(f"Fuzz Events Handled: {fuzz_counts}")
        print(f"Final Memory Delta: {mem_end - mem_start:.2f}MB")
        print("="*60)
        
        if mem_end - mem_start > 500: # 500MB allowance for long run + buffers
             print("CRITICAL: Significant memory leak detected.")
             sys.exit(2)
        
        sys.exit(0)

if __name__ == "__main__":
    tester = ExtremeStressTester(iterations=100000, fuzz_rate=0.10)
    tester.run()