"""
API Integration Testing for CircuitScope - Mechanistic Interpretability Platform
Tests:
1. Health endpoint (/api/health)
2. Live Activation Patching endpoint (/api/patch)
3. Attention extraction endpoint (/api/attention)
4. Input validation and error handling
"""
import requests
import json
import sys

BACKEND_URL = "http://localhost:8000"

def test_health():
    print("=" * 50)
    print("TEST 1: Health Check")
    print("=" * 50)
    try:
        r = requests.get(f"{BACKEND_URL}/api/health", timeout=10)
        data = r.json()
        print(f"  Status: {r.status_code}")
        print(f"  Response: {json.dumps(data, indent=2)}")
        assert r.status_code == 200
        assert data["status"] == "healthy"
        assert "model_loaded" in data
        assert "model_loading" in data
        assert "real_inference_active" in data
        print("  PASSED ✓\n")
        return True, data
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False, None

def test_patch(model_loaded):
    print("=" * 50)
    print("TEST 2: Live Activation Patching")
    print("=" * 50)
    try:
        payload = {
            "prompt": "When Mary and John went to the store, John gave a bottle of milk to"
        }
        r = requests.post(
            f"{BACKEND_URL}/api/patch",
            json=payload,
            timeout=30
        )
        data = r.json()
        print(f"  Status: {r.status_code}")
        assert r.status_code == 200
        assert "tokens" in data
        assert "layers" in data
        assert "values" in data
        assert "hotspots" in data
        assert "baseline" in data
        
        # Verify length of outputs
        assert len(data["layers"]) == 12
        assert len(data["values"]) == 12
        assert len(data["tokens"]) == len(data["values"][0])
        print("  PASSED ✓\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

def test_attention(model_loaded):
    print("=" * 50)
    print("TEST 3: Attention Weights Extraction")
    print("=" * 50)
    try:
        payload = {
            "prompt": "When Alice and Bob went to the park, Bob gave flowers to",
            "layer": 5,
            "head": 5
        }
        r = requests.post(
            f"{BACKEND_URL}/api/attention",
            json=payload,
            timeout=15
        )
        data = r.json()
        print(f"  Status: {r.status_code}")
        assert r.status_code == 200
        assert "tokens" in data
        assert "weights" in data
        assert data["layer"] == 5
        assert data["head"] == 5
        
        # Verify weight matrices dimensions
        n_tokens = len(data["tokens"])
        assert len(data["weights"]) == n_tokens
        assert len(data["weights"][0]) == n_tokens
        print("  PASSED ✓\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

def test_validation_error():
    print("=" * 50)
    print("TEST 4: Input Validation (Short/Invalid Parameters)")
    print("=" * 50)
    try:
        # Short prompt test
        payload = {"prompt": "Hi"}
        r = requests.post(
            f"{BACKEND_URL}/api/patch",
            json=payload,
            timeout=10
        )
        print(f"  Status (Short Prompt): {r.status_code}")
        assert r.status_code == 422  # Validation error
        
        # Invalid layer index test
        payload_attn = {
            "prompt": "Valid prompt here",
            "layer": 15,  # Out of range (max 11)
            "head": 0
        }
        r = requests.post(
            f"{BACKEND_URL}/api/attention",
            json=payload_attn,
            timeout=10
        )
        print(f"  Status (Invalid Layer): {r.status_code}")
        assert r.status_code == 422  # Validation error
        
        print("  PASSED ✓ (correctly rejected invalid payloads)\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

def test_sae():
    print("=" * 50)
    print("TEST 5: Dynamic Sparse Autoencoder (SAE) Activation")
    print("=" * 50)
    try:
        payload = {
            "prompt": "When the doctor analyzed the DNA sequence, he found a gene mutation."
        }
        r = requests.post(
            f"{BACKEND_URL}/api/sae/activate",
            json=payload,
            timeout=15
        )
        data = r.json()
        print(f"  Status: {r.status_code}")
        assert r.status_code == 200
        assert "tokens" in data
        assert "l0_sparsity" in data
        assert "explained_variance" in data
        assert "active_features" in data
        assert "real_inference" in data
        assert "feature_clusters" in data
        
        # Verify active features properties
        if len(data["active_features"]) > 0:
            first_feat = data["active_features"][0]
            assert "index" in first_feat
            assert "activation_value" in first_feat
            assert "label" in first_feat
            assert "category" in first_feat
            assert "tokens_fired" in first_feat
            
        print("  PASSED ✓\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

def test_attribution():
    print("=" * 50)
    print("TEST 6: Attribution Patching (Efficient Syed et al.)")
    print("=" * 50)
    try:
        payload = {
            "prompt": "When Mary and John went to the store, John gave a bottle of milk to"
        }
        r = requests.post(f"{BACKEND_URL}/api/attribution", json=payload, timeout=15)
        data = r.json()
        print(f"  Status: {r.status_code}")
        assert r.status_code == 200
        assert "tokens" in data
        assert "values" in data
        assert data["method"] == "attribution_patching"
        assert len(data["values"]) == 12
        print("  PASSED ✓\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

def test_logit_lens():
    print("=" * 50)
    print("TEST 7: Logit Lens Vocabulary Projection")
    print("=" * 50)
    try:
        payload = {
            "prompt": "When Mary and John went to the store, John gave a bottle of milk to"
        }
        r = requests.post(f"{BACKEND_URL}/api/logit-lens", json=payload, timeout=15)
        data = r.json()
        print(f"  Status: {r.status_code}")
        assert r.status_code == 200
        assert "tokens" in data
        assert "top_predictions" in data
        assert "target_token_probs" in data
        assert len(data["top_predictions"]) == 12
        print("  PASSED ✓\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

def test_steer():
    print("=" * 50)
    print("TEST 8: SAE Feature Steering")
    print("=" * 50)
    try:
        payload = {
            "prompt": "When Mary and John went to the store,",
            "feature_index": 2048,
            "steering_strength": 15.0,
            "layer": 6
        }
        r = requests.post(f"{BACKEND_URL}/api/steer", json=payload, timeout=15)
        data = r.json()
        print(f"  Status: {r.status_code}")
        assert r.status_code == 200
        assert "steered_completion" in data
        assert "original_completion" in data
        assert data["feature_index"] == 2048
        print("  PASSED ✓\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

def test_knockout():
    print("=" * 50)
    print("TEST 9: Attention Head Knockout")
    print("=" * 50)
    try:
        payload = {
            "prompt": "When Mary and John went to the store, John gave a bottle of milk to",
            "heads_to_knockout": [[5, 1], [6, 9]],
            "mode": "zero"
        }
        r = requests.post(f"{BACKEND_URL}/api/knockout", json=payload, timeout=15)
        data = r.json()
        print(f"  Status: {r.status_code}")
        assert r.status_code == 200
        assert "baseline_logit_diff" in data
        assert "knocked_logit_diff" in data
        assert "performance_retained_pct" in data
        print("  PASSED ✓\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

def test_compare_models():
    print("=" * 50)
    print("TEST 10: Pythia Cross-Model Comparison")
    print("=" * 50)
    try:
        payload = {
            "prompt": "When Mary and John went to the store, John gave a bottle of milk to"
        }
        r = requests.post(f"{BACKEND_URL}/api/compare-models", json=payload, timeout=45)
        data = r.json()
        print(f"  Status: {r.status_code}")
        assert r.status_code == 200
        assert "gpt2" in data
        assert "cross_model_correlation" in data
        print("  PASSED ✓\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

def test_validate_attribution():
    print("=" * 50)
    print("TEST 11: Pearson Linearity Validation Check")
    print("=" * 50)
    try:
        payload = {
            "prompt": "When Mary and John went to the store, John gave a bottle of milk to",
            "layer": 6,
            "head": 9
        }
        r = requests.post(f"http://localhost:8000/api/validate-attribution", json=payload, timeout=15)
        data = r.json()
        print(f"  Status: {r.status_code}")
        assert r.status_code == 200
        assert "pearson_r" in data
        assert "verdict" in data
        assert "interpretation" in data
        assert "taylor_values" in data
        assert "causal_values" in data
        assert data["checked_head"] == "L6H9"
        assert len(data["taylor_values"]) == 12
        assert len(data["causal_values"]) == 12
        print("  PASSED ✓\n")
        return True
    except Exception as e:
        print(f"  FAILED ✗: {e}\n")
        return False

if __name__ == "__main__":
    results = []
    
    # 1. Health check
    health_passed, health_data = test_health()
    results.append(("Health Check", health_passed))
    
    model_loaded = False
    if health_passed and health_data:
        model_loaded = health_data.get("model_loaded", False)
        
    # 2. Patching endpoint
    results.append(("Activation Patching", test_patch(model_loaded)))
    
    # 3. Attention endpoint
    results.append(("Attention Weights", test_attention(model_loaded)))
    
    # 4. Input validation
    results.append(("Validation Checks", test_validation_error()))
    
    # 5. SAE activation endpoint
    results.append(("Dynamic SAE Latent Extractor", test_sae()))
    
    # 6. Attribution patching
    results.append(("Attribution Patching", test_attribution()))
    
    # 7. Logit Lens
    results.append(("Logit Lens", test_logit_lens()))
    
    # 8. Feature Steering
    results.append(("SAE Feature Steering", test_steer()))
    
    # 9. Attention Knockout
    results.append(("Attention Knockout", test_knockout()))
    
    # 10. Cross-Model comparison
    results.append(("Cross-Model Comparison", test_compare_models()))
    
    # 11. Pearson Linearity Validation Check
    results.append(("Attribution Linearity Validation", test_validate_attribution()))
    
    print("\n" + "=" * 50)
    print("SUMMARY")
    print("=" * 50)
    all_passed = True
    for name, passed in results:
        status = "PASSED ✓" if passed else "FAILED ✗"
        print(f"  {name}: {status}")
        if not passed:
            all_passed = False
            
    print()
    if all_passed:
        print("ALL TESTS PASSED ✓")
        sys.exit(0)
    else:
        print("SOME TESTS FAILED ✗")
        sys.exit(1)