"""Test neuronlens search_features and steer_features functions."""
import os
import sys
from pathlib import Path
import neuronlens

env_path = Path(__file__).parent.parent / ".env"
if env_path.exists():
    try:
        from dotenv import load_dotenv
        load_dotenv(env_path, override=True)
    except ImportError:
        pass

engine = neuronlens.Engine(api_key=os.getenv("API_KEY") or os.getenv("NEURONLENS_API_KEY"))

print("=== Testing NeuronLens Search & Steer ===\n")

# Test search_features
features = engine.search_features("financial analysis", k=3, layer=16)
print(f"✓ Search: {len(features)} features found")
if features:
    print(f"✓ Top feature: {features[0].get('label', 'N/A')[:40]}")

# Test steer_features (with longer timeout for model loading)
if features:
    print(f"Testing steer with feature: {features[0].get('id')}")
    # Use feature_id instead of id if available, or convert id format
    feature_id = features[0].get('feature_id') or features[0].get('id', '').replace('f_', '')
    try:
        result = engine.steer_features(
            "What is inflation?", 
            [{"id": features[0]["id"], "magnitude": 0.5}], 
            model="meta-llama/Llama-2-7b-hf"
        )
        print(f"✓ Steer: has_steering={result.get('has_steering', False)}")
        if result.get('original_text'):
            print(f"✓ Original text: {result['original_text'][:50]}...")
        if result.get('steered_text'):
            print(f"✓ Steered text: {result['steered_text'][:50]}...")
    except Exception as e:
        print(f"⚠ Steer timeout/error (this is expected if model needs to load): {type(e).__name__}")
        print(f"  Note: Steering requires model loading which can take 30+ seconds")

engine.unload_all_models()
print("\n=== Search & Steer test passed ===")
