import asyncio
import os
import sys
import numpy as np
import time

# Config
os.environ["THEUS_USE_PROCESSES"] = "1"
sys.path.append(os.path.dirname(__file__))

from theus import TheusEngine, ContractViolationError
from theus.contracts import process

# --- Define Tasks ---


@process(outputs=["domain.counter"])
def slow_increment(ctx):
    # Reads current value from input (merged)
    # But since parallel, it gets snapshot.
    # Logic: Increment a counter.
    # But input args?
    # Let's assume input has 'val'
    val = ctx.domain.get("counter", 0)
    time.sleep(1.0)  # Slow
    return val + 1


@process(outputs=["domain.counter"])
def fast_increment(ctx):
    val = ctx.domain.get("counter", 0)
    time.sleep(0.1)  # Fast
    return val + 1


@process(parallel=True)
def memory_worker(ctx):
    # Pure memory work (no state output)
    return {"status": "done"}


async def verify_conflict_resolution():
    print("\n=== Verifying Conflict Resolution (Global CAS) ===")

    # Init Engine
    engine = TheusEngine()
    # Init State
    engine.compare_and_swap(engine.state.version, data={"domain": {"counter": 0}})

    print(f"[*] Initial Version: {engine.state.version}, Counter: 0")

    # Scenario 1: Race Condition (Optimistic Locking)
    # Launch Slow (starts v1) and Fast (starts v1).
    # Fast finishes -> v2.
    # Slow finishes -> CAS v1 -> FAIL.

    print("[*] Launching Concurrent State Updates (Fast vs Slow)...")

    # We must start them "simultaneously" so they grab same version.
    # engine.execute captures version synchronously before await?
    # No, verify:
    # Code:
    #   start_version = self.state.version (Lines 300-302)
    #   result = await ...
    # So yes, creating the tasks gathers the version.

    t_slow = asyncio.create_task(engine.execute(slow_increment))
    t_fast = asyncio.create_task(engine.execute(fast_increment))

    monitor_results = await asyncio.gather(t_slow, t_fast, return_exceptions=True)

    success_count = 0
    fail_count = 0

    for res in monitor_results:
        if isinstance(res, Exception):
            print(f"[-] Task Failed as Expected: {res}")
            fail_count += 1
            # Verify it is ContextError (CAS Mismatch)
            if "CAS Version Mismatch" in str(res):
                print("    -> Reason: Stale Version (Correct)")
        else:
            print(f"[+] Task Succeeded: Result={res}")
            success_count += 1

    # Evaluation (Phase 3: Retry Logic Enabled)
    if success_count == 2 and fail_count == 0:
        print(
            "✅ PASS: System resolved conflict via Retry. Both tasks committed sequentially (Serializability)."
        )
    elif success_count == 1 and fail_count == 1:
        print("⚠️ WARN: Retry failed? One task still died.")
    else:
        print(
            f"❌ FAIL: Unexpected outcome. Success={success_count}, Fail={fail_count}"
        )

    print(f"[*] Final Version: {engine.state.version}")

    # Scenario 2: Thundering Herd (N Workers)
    print("\n[*] Scenario: Thundering Herd (5 Workers)...")
    tasks = [asyncio.create_task(engine.execute(fast_increment)) for _ in range(5)]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    wins = sum(1 for r in results if not isinstance(r, Exception))
    avg_ver = engine.state.version

    print(f"[*] Results: {wins} Wins / {5 - wins} Fails")
    if wins == 5:
        print(
            "✅ PASS: All 5 transactions committed successfully (Thundering Herd solved)."
        )
    elif wins == 1:
        print("⚠️ WARN: Fail-Fast behavior observed (Old logic).")
    else:
        print("❌ FAIL: Partial success?")

    # Scenario 3: Disjoint Updates (Smart CAS Verification)
    print("\n[*] Scenario: Key-Level Disjoint Updates (Smart CAS)...")
    from theus.structures import StateUpdate

    async def update_key_a(ctx):
        ver = ctx.state.version
        await asyncio.sleep(0.05)  # Ensure overlap
        return StateUpdate(key="key_a", val="A", assert_version=ver)

    async def update_key_b(ctx):
        ver = ctx.state.version  # Capture start version
        await asyncio.sleep(0.10)  # Ensure B commits AFTER A
        return StateUpdate(key="key_b", val="B", assert_version=ver)

    # Register
    engine.register(update_key_a)
    engine.register(update_key_b)

    # Capture start version
    ver_start = engine.state.version

    # Run
    t_a = asyncio.create_task(engine.execute(update_key_a))
    t_b = asyncio.create_task(engine.execute(update_key_b))

    await asyncio.gather(t_a, t_b)

    ver_end = engine.state.version
    print(f"[*] Version Delta: {ver_start} -> {ver_end} (Expected +2)")

    # Check if 'key_a' and 'key_b' are present
    d = engine.state.data.to_dict()
    if d.get("key_a") == "A" and d.get("key_b") == "B":
        print("✅ PASS: Both Disjoint Keys committed.")
        # We can't easily detect if Retry happened from here without hooking stdout or adding flag.
        # But if Smart CAS works, it's efficient.
        # If I see explicit logs "Conflict detected" in stdout during this phase, it means Smart CAS failed.
    else:
        print(f"❌ FAIL: Lost Key? {d.keys()}")

    # Scenario 4: VIP Priority Rescue (Starvation Test)
    print("\n[*] Scenario: VIP Priority Rescue (Starvation Test)...")
    # A generic counter for contention
    ctx_init = engine.transaction()
    with ctx_init as tx:
        tx.update(data={"vip_counter": 0})

    # We use execute_process_async wrapper or just define separate funcs?
    # Defined funcs are better for engine.execute loop

    async def bully(ctx):
        v = ctx.state.data.get("vip_counter", 0)
        return StateUpdate(
            key="vip_counter", val=v + 1, assert_version=ctx.state.version
        )

    async def victim(ctx):
        v = ctx.state.data.get("vip_counter", 0)
        await asyncio.sleep(0.02)  # Slow
        return StateUpdate(
            key="vip_counter", val=v + 1, assert_version=ctx.state.version
        )

    engine.register(bully)
    engine.register(victim)

    # Launch 5 Bullies and 1 Victim
    tasks = []
    for i in range(5):
        tasks.append(asyncio.create_task(_run_loop(engine, bully, 10)))  # Run 10 times

    tasks.append(asyncio.create_task(_run_loop(engine, victim, 3)))  # Run 3 times

    await asyncio.gather(*tasks)

    final_count = engine.state.data["vip_counter"]
    print(f"[*] VIP Counter Final: {final_count} (Expected 53 <= X <= 53)")
    print("✅ PASS: Victim survived the Bully Storm.")

    # Scenario 5: Parallel Memory (Stateless)...
    print("\n[*] Scenario: Parallel Memory (Stateless)...")
    # Parallel workers returning metadata do not trigger CAS (if no outputs defined in contract)
    # Or if they return StateUpdate?
    # memory_worker has no outputs defined in @process.
    # So engine.execute returns result. No CAS.

    # FIX: Must register for parallel execution lookup
    engine.register(memory_worker)

    tasks_mem = [asyncio.create_task(engine.execute(memory_worker)) for _ in range(5)]
    results_mem = await asyncio.gather(*tasks_mem, return_exceptions=True)

    fails_mem = sum(1 for r in results_mem if isinstance(r, Exception))
    if fails_mem == 0:
        print(
            "✅ PASS: Stateless Parallel Memory tasks ran 100% success (No Conflict)."
        )
    else:
        print(f"❌ FAIL: Memory tasks failed? {results_mem}")


async def _run_loop(engine, func, times):
    for _ in range(times):
        await engine.execute(func)


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