#!/usr/bin/env python3
"""
E2E Validation for LockStock Idempotency Migration
---------------------------------------------------
Tests the belt-and-suspenders idempotency mechanism against the live server.

Run with:
    python e2e_validation.py

Requirements:
    - httpx
    - Live server deployed with idempotency support
"""

import httpx
import uuid
import time
import sys
from typing import Optional, Dict, Any

# Configuration
API_ENDPOINT = "https://lockstock-api-i9kp.onrender.com"
API_KEY = "lsk_admin_0f89affdbfd62c68c7186bf91eb283969f7a163f5834cbf08923eaa9bac97cd5954f63cf7a07ca65168d3cfd54cec72357952c7462e21b5539828a6d788f1d58"
AGENT_ID = "agent_b3484fec_genesisprovision"


def headers():
    return {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }


def print_result(name: str, passed: bool, details: str = ""):
    status = "PASS" if passed else "FAIL"
    print(f"\n{'=' * 60}")
    print(f"[{status}] {name}")
    print(f"{'=' * 60}")
    if details:
        print(details)


def print_failure(response: httpx.Response):
    print(f"  Status: {response.status_code}")
    print(f"  Headers: {dict(response.headers)}")
    try:
        print(f"  Body: {response.json()}")
    except:
        print(f"  Body (raw): {response.text}")


def scenario_1_state_endpoint() -> Optional[int]:
    """
    Scenario 1: GET /api/guard/agent/{agent_id}/state
    Assert response contains required fields. Return last_sequence.
    """
    print("\n" + "=" * 60)
    print("SCENARIO 1: State Endpoint")
    print("=" * 60)

    with httpx.Client(timeout=30.0) as client:
        resp = client.get(
            f"{API_ENDPOINT}/api/guard/agent/{AGENT_ID}/state",
            headers=headers()
        )

        if resp.status_code != 200:
            print_result("State Endpoint", False, f"Expected 200, got {resp.status_code}")
            print_failure(resp)
            return None

        data = resp.json()
        required_fields = ["agent_id", "last_hash", "last_sequence", "state_matrix", "last_server_timestamp", "source"]
        missing = [f for f in required_fields if f not in data]

        if missing:
            print_result("State Endpoint", False, f"Missing fields: {missing}")
            print(f"  Got: {data}")
            return None

        sequence = data["last_sequence"]
        print_result("State Endpoint", True, f"Current sequence: {sequence}")
        print(f"  agent_id: {data['agent_id']}")
        print(f"  last_hash: {data['last_hash'][:32]}...")
        print(f"  source: {data['source']}")
        return sequence


def scenario_2_normal_stamp(baseline_sequence: int) -> bool:
    """
    Scenario 2: Normal stamp with idempotency_key and expected_sequence
    """
    print("\n" + "=" * 60)
    print("SCENARIO 2: Normal Stamp")
    print("=" * 60)

    idempotency_key = str(uuid.uuid4())
    expected_sequence = baseline_sequence + 1

    print(f"  idempotency_key: {idempotency_key}")
    print(f"  expected_sequence: {expected_sequence}")

    with httpx.Client(timeout=30.0) as client:
        resp = client.post(
            f"{API_ENDPOINT}/v1/stamp",
            json={
                "agent_id": AGENT_ID,
                "task": "heartbeat",
                "idempotency_key": idempotency_key,
                "expected_sequence": expected_sequence
            },
            headers=headers()
        )

        if resp.status_code != 200:
            print_result("Normal Stamp", False, f"Expected 200, got {resp.status_code}")
            print_failure(resp)
            return False

        data = resp.json()
        if not data.get("authorized"):
            print_result("Normal Stamp", False, f"Not authorized: {data}")
            return False

        actual_sequence = data.get("sequence")
        if actual_sequence != expected_sequence:
            print_result("Normal Stamp", False,
                f"Sequence mismatch: expected {expected_sequence}, got {actual_sequence}")
            return False

        print_result("Normal Stamp", True, f"Authorized, sequence={actual_sequence}")
        print(f"  state_hash: {data.get('state_hash', 'N/A')[:32]}...")
        return True


def scenario_3_critical_test(baseline_sequence: int) -> bool:
    """
    Scenario 3: THE CRITICAL TEST - Cached retry after lost response

    1. Send stamp with 1ms timeout (server processes, response lost)
    2. Retry with SAME idempotency_key (should return cached response)
    3. Verify sequence advanced by exactly 1, not 2
    """
    print("\n" + "=" * 60)
    print("SCENARIO 3: CRITICAL TEST - Cached Retry After Lost Response")
    print("=" * 60)

    idempotency_key = str(uuid.uuid4())
    expected_sequence = baseline_sequence + 1

    print(f"  Pre-test baseline sequence: {baseline_sequence}")
    print(f"  idempotency_key: {idempotency_key}")
    print(f"  expected_sequence: {expected_sequence}")

    # Step 1: Send with 1ms timeout (will timeout but server may process)
    print("\n  Step 1: Sending with 1ms timeout (simulating lost response)...")
    first_request_timed_out = False
    try:
        with httpx.Client(timeout=0.001) as client:  # 1ms timeout
            resp = client.post(
                f"{API_ENDPOINT}/v1/stamp",
                json={
                    "agent_id": AGENT_ID,
                    "task": "heartbeat",
                    "idempotency_key": idempotency_key,
                    "expected_sequence": expected_sequence
                },
                headers=headers()
            )
            # If we get here, the timeout was too generous
            print(f"    Unexpected: Got response with 1ms timeout")
            print(f"    Status: {resp.status_code}, Data: {resp.json()}")
    except httpx.TimeoutException:
        first_request_timed_out = True
        print(f"    Timeout as expected (server may have processed)")
    except Exception as e:
        print(f"    Error: {type(e).__name__}: {e}")
        first_request_timed_out = True  # Network errors are also "lost response"

    # Give server a moment to process if it received the request
    time.sleep(0.5)

    # Step 2: Retry with same idempotency_key and normal timeout
    print("\n  Step 2: Retrying with same idempotency_key (30s timeout)...")
    with httpx.Client(timeout=30.0) as client:
        resp = client.post(
            f"{API_ENDPOINT}/v1/stamp",
            json={
                "agent_id": AGENT_ID,
                "task": "heartbeat",
                "idempotency_key": idempotency_key,
                "expected_sequence": expected_sequence
            },
            headers=headers()
        )

        if resp.status_code != 200:
            print_result("Critical Test", False, f"Retry failed with {resp.status_code}")
            print_failure(resp)
            return False

        retry_data = resp.json()
        retry_sequence = retry_data.get("sequence")
        print(f"    Retry response: authorized={retry_data.get('authorized')}, sequence={retry_sequence}")

    # Step 3: Check current state - sequence should have advanced by exactly 1
    print("\n  Step 3: Verifying chain state...")
    with httpx.Client(timeout=30.0) as client:
        resp = client.get(
            f"{API_ENDPOINT}/api/guard/agent/{AGENT_ID}/state",
            headers=headers()
        )

        if resp.status_code != 200:
            print_result("Critical Test", False, f"State check failed")
            print_failure(resp)
            return False

        final_state = resp.json()
        final_sequence = final_state["last_sequence"]

        sequence_delta = final_sequence - baseline_sequence
        print(f"    Baseline sequence: {baseline_sequence}")
        print(f"    Final sequence: {final_sequence}")
        print(f"    Delta: {sequence_delta}")

        if sequence_delta == 1:
            print_result("Critical Test", True,
                "Chain advanced by exactly 1 (idempotency worked!)")
            return True
        elif sequence_delta == 2:
            print_result("Critical Test", False,
                "MIGRATION BROKEN: Chain advanced by 2 (double-advance detected!)")
            return False
        elif sequence_delta == 0:
            print_result("Critical Test", False,
                "Chain did not advance (request never reached server)")
            return False
        else:
            print_result("Critical Test", False,
                f"Unexpected delta: {sequence_delta}")
            return False


def scenario_4_sequence_mismatch(current_sequence: int) -> bool:
    """
    Scenario 4: 409 on intentionally wrong expected_sequence
    """
    print("\n" + "=" * 60)
    print("SCENARIO 4: Sequence Mismatch (409 Expected)")
    print("=" * 60)

    wrong_sequence = 1  # Intentionally wrong

    print(f"  Current sequence: {current_sequence}")
    print(f"  Sending expected_sequence: {wrong_sequence} (wrong)")

    with httpx.Client(timeout=30.0) as client:
        resp = client.post(
            f"{API_ENDPOINT}/v1/stamp",
            json={
                "agent_id": AGENT_ID,
                "task": "heartbeat",
                "idempotency_key": str(uuid.uuid4()),
                "expected_sequence": wrong_sequence
            },
            headers=headers()
        )

        if resp.status_code == 409:
            data = resp.json()
            print_result("Sequence Mismatch", True,
                f"Got 409 as expected: {data.get('reason', 'N/A')}")
            print(f"  current_sequence in response: {data.get('current_sequence', 'N/A')}")
            return True
        else:
            print_result("Sequence Mismatch", False,
                f"Expected 409, got {resp.status_code}")
            print_failure(resp)
            return False


def scenario_5_idempotent_replay() -> bool:
    """
    Scenario 5: Idempotent replay - same request returns identical response
    """
    print("\n" + "=" * 60)
    print("SCENARIO 5: Idempotent Replay")
    print("=" * 60)

    # Get current sequence
    with httpx.Client(timeout=30.0) as client:
        resp = client.get(
            f"{API_ENDPOINT}/api/guard/agent/{AGENT_ID}/state",
            headers=headers()
        )
        baseline_sequence = resp.json()["last_sequence"]

    idempotency_key = str(uuid.uuid4())
    expected_sequence = baseline_sequence + 1

    print(f"  idempotency_key: {idempotency_key}")
    print(f"  expected_sequence: {expected_sequence}")

    # First request
    print("\n  First request...")
    with httpx.Client(timeout=30.0) as client:
        resp1 = client.post(
            f"{API_ENDPOINT}/v1/stamp",
            json={
                "agent_id": AGENT_ID,
                "task": "heartbeat",
                "idempotency_key": idempotency_key,
                "expected_sequence": expected_sequence
            },
            headers=headers()
        )

        if resp1.status_code != 200:
            print_result("Idempotent Replay", False,
                f"First request failed: {resp1.status_code}")
            print_failure(resp1)
            return False

        data1 = resp1.json()
        seq1 = data1.get("sequence")
        hash1 = data1.get("state_hash")
        print(f"    sequence={seq1}, state_hash={hash1[:32] if hash1 else 'N/A'}...")

    # Second request with SAME idempotency_key
    print("\n  Second request (same idempotency_key)...")
    with httpx.Client(timeout=30.0) as client:
        resp2 = client.post(
            f"{API_ENDPOINT}/v1/stamp",
            json={
                "agent_id": AGENT_ID,
                "task": "heartbeat",
                "idempotency_key": idempotency_key,
                "expected_sequence": expected_sequence
            },
            headers=headers()
        )

        if resp2.status_code != 200:
            print_result("Idempotent Replay", False,
                f"Second request failed: {resp2.status_code}")
            print_failure(resp2)
            return False

        data2 = resp2.json()
        seq2 = data2.get("sequence")
        hash2 = data2.get("state_hash")
        print(f"    sequence={seq2}, state_hash={hash2[:32] if hash2 else 'N/A'}...")

    # Verify responses match and sequence didn't advance
    if seq1 != seq2:
        print_result("Idempotent Replay", False,
            f"Sequences differ: {seq1} vs {seq2}")
        return False

    if hash1 != hash2:
        print_result("Idempotent Replay", False,
            f"Hashes differ: {hash1} vs {hash2}")
        return False

    # Verify sequence in state endpoint matches
    with httpx.Client(timeout=30.0) as client:
        resp = client.get(
            f"{API_ENDPOINT}/api/guard/agent/{AGENT_ID}/state",
            headers=headers()
        )
        final_sequence = resp.json()["last_sequence"]

    if final_sequence != expected_sequence:
        print_result("Idempotent Replay", False,
            f"Chain advanced unexpectedly: expected {expected_sequence}, got {final_sequence}")
        return False

    print_result("Idempotent Replay", True,
        "Both requests returned identical response, sequence stable")
    return True


def main():
    print("""
╔═══════════════════════════════════════════════════════════════╗
║          LOCKSTOCK IDEMPOTENCY E2E VALIDATION                 ║
╠═══════════════════════════════════════════════════════════════╣
║  Endpoint: {:<49} ║
║  Agent:    {:<49} ║
╚═══════════════════════════════════════════════════════════════╝
""".format(API_ENDPOINT, AGENT_ID))

    results = {}

    # Scenario 1: State endpoint
    baseline = scenario_1_state_endpoint()
    results["1_state_endpoint"] = baseline is not None

    if baseline is None:
        print("\n\nABORTING: Cannot proceed without baseline sequence")
        sys.exit(1)

    # Scenario 2: Normal stamp
    results["2_normal_stamp"] = scenario_2_normal_stamp(baseline)

    # Get updated baseline after scenario 2
    with httpx.Client(timeout=30.0) as client:
        resp = client.get(
            f"{API_ENDPOINT}/api/guard/agent/{AGENT_ID}/state",
            headers=headers()
        )
        baseline = resp.json()["last_sequence"]

    # Scenario 3: THE CRITICAL TEST
    results["3_CRITICAL_cached_retry"] = scenario_3_critical_test(baseline)

    # Get updated baseline after scenario 3
    with httpx.Client(timeout=30.0) as client:
        resp = client.get(
            f"{API_ENDPOINT}/api/guard/agent/{AGENT_ID}/state",
            headers=headers()
        )
        current_sequence = resp.json()["last_sequence"]

    # Scenario 4: Sequence mismatch
    results["4_sequence_mismatch"] = scenario_4_sequence_mismatch(current_sequence)

    # Scenario 5: Idempotent replay
    results["5_idempotent_replay"] = scenario_5_idempotent_replay()

    # Summary
    print("\n" + "=" * 60)
    print("VALIDATION SUMMARY")
    print("=" * 60)

    all_passed = True
    for name, passed in results.items():
        status = "PASS" if passed else "FAIL"
        marker = "  " if passed else ">>"
        print(f"{marker} {name}: {status}")
        if not passed:
            all_passed = False

    print("\n" + "=" * 60)
    if all_passed:
        print("ALL TESTS PASSED - Ready for PyPI publish")
    else:
        print("SOME TESTS FAILED - DO NOT PUBLISH")
    print("=" * 60)

    sys.exit(0 if all_passed else 1)


if __name__ == "__main__":
    main()
