#!/usr/bin/env python3
"""
Quick test to verify logits_for_choices() works.
This tests the engine function directly without starting the server.
"""

import sys
import os

# Set up PYTHONPATH like the server does
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, repo_root)
sys.path.insert(0, "/home/voidstr3m33/echo_labs/helix/proof/helix_proof_v2")

def test_logits_extraction():
    """Test that logits_for_choices returns 4 floats."""
    print("=== Testing logits_for_choices() ===\n")

    # Import the engine module
    try:
        from artifacts.context_pack.flowtorch import flowtorch_engine
        print("✅ Engine module imported")
    except Exception as e:
        print(f"❌ Failed to import engine: {e}")
        return False

    # Test with sample messages
    messages = [
        {"role": "user", "content": "What is 2+2? A)3 B)4 C)5 D)6"}
    ]

    print(f"Messages: {messages}\n")

    # Call logits_for_choices
    try:
        logits = flowtorch_engine.logits_for_choices(messages)
        print(f"✅ logits_for_choices() returned: {logits}")
    except Exception as e:
        print(f"❌ logits_for_choices() failed: {e}")
        import traceback
        traceback.print_exc()
        return False

    # Verify format
    if not isinstance(logits, list):
        print(f"❌ Expected list, got {type(logits)}")
        return False

    if len(logits) != 4:
        print(f"❌ Expected 4 logits, got {len(logits)}")
        return False

    if not all(isinstance(x, (int, float)) for x in logits):
        print(f"❌ Expected all floats, got {[type(x) for x in logits]}")
        return False

    print("✅ Logits format is correct: list of 4 floats\n")

    # Check if they're all zeros (fallback) or real values
    if all(x == 0.0 for x in logits):
        print("⚠️  Warning: All logits are 0.0 (using fallback mode)")
        print("   This is OK if HuggingFace transformers not available")
        print("   Real logits will work when HF model is loaded\n")
    else:
        print(f"✅ Real logits extracted! Values: {logits}")
        print("   (Non-zero means real model logits working)\n")

    return True

def test_serve_integration():
    """Test that serve_flowtorch.py can call the logits function."""
    print("=== Testing serve_flowtorch.py integration ===\n")

    try:
        from edge_app import serve_flowtorch
        print("✅ serve_flowtorch imported")
    except Exception as e:
        print(f"❌ Failed to import serve_flowtorch: {e}")
        return False

    # Test get_base_logits_for_letters
    messages = [
        {"role": "user", "content": "Test"}
    ]

    try:
        logits = serve_flowtorch.get_base_logits_for_letters(messages)
        print(f"✅ get_base_logits_for_letters() returned: {logits}")
    except Exception as e:
        print(f"❌ get_base_logits_for_letters() failed: {e}")
        import traceback
        traceback.print_exc()
        return False

    # Verify format
    if len(logits) != 4:
        print(f"❌ Expected 4 logits, got {len(logits)}")
        return False

    print("✅ Server integration working\n")
    return True

if __name__ == "__main__":
    print("Phase A1: Real Logits Extraction Test\n")
    print("=" * 50)
    print()

    # Run tests
    test1 = test_logits_extraction()
    test2 = test_serve_integration()

    print("=" * 50)
    if test1 and test2:
        print("\n✅ All tests PASSED")
        print("\nNext steps:")
        print("  1. Start server: FLOWTORCH_STEER_MODE=bias bash scripts/edge_app_run_flowtorch.sh")
        print("  2. Run smoke test: bash tests/test_bias_steering_smoke.sh")
        sys.exit(0)
    else:
        print("\n❌ Some tests FAILED")
        sys.exit(1)
