"""Test neuronlens common functions using the neuronlens package."""
import os
import sys
from pathlib import Path
import neuronlens
import glob

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

api_key = os.getenv("API_KEY") or os.getenv("NEURONLENS_API_KEY")
if not api_key:
    print("⚠ Warning: No API_KEY found. Using dummy key for local testing.")
    api_key = "test_key_local"

engine = neuronlens.Engine(api_key=api_key)

def get_model_cache_path(model_id: str) -> str:
    cache_base = os.getenv("HF_HOME") or os.getenv("TRANSFORMERS_CACHE") or os.path.expanduser("~/.cache/huggingface/hub")
    model_dir = model_id.replace("/", "--")
    matches = glob.glob(os.path.join(cache_base, f"models--{model_dir}", "snapshots", "*"))
    return matches[0] if matches else model_id

def get_sae_base():
    """Get SAE base path using same logic as test_slm.py"""
    # Try config.paths first (same as test_slm.py uses)
    try:
        # Add parent directory to path to import config
        parent_dir = os.path.dirname(os.path.dirname(__file__))
        if parent_dir not in sys.path:
            sys.path.insert(0, parent_dir)
        from config.paths import get_sae_models_base
        sae_base = get_sae_models_base()
        if sae_base:
            return sae_base
    except ImportError:
        pass
    
    # Fallback: same logic as test_slm.py
    sae_base = os.getenv("SAE_MODELS_BASE")
    if sae_base:
        return sae_base
    
    # Fallback: try cloud_api structure (same as test_slm.py)
    # From pip/ directory, go up 2 levels to sae_app_v3/, then up 1 more to InterpUseCases_v2/
    cloud_api_sae = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "cloud_api", "sae_models")
    if os.path.exists(cloud_api_sae):
        return cloud_api_sae
    
    return None

def test_model_cache_status():
    """Test engine.model_cache_status()"""
    print("1. model_cache_status()")
    status = engine.model_cache_status()
    loaded = status.get('loaded_models', [])
    count = len(loaded) if isinstance(loaded, list) else loaded
    print(f"   ✓ Models: {count}, Device: {status.get('device', 'unknown')}")
    return True

def test_extract_top_features():
    """Test engine.extract_top_features()"""
    print("\n2. extract_top_features()")
    
    sae_base = get_sae_base()
    model_path = get_model_cache_path("ProsusAI/finbert")
    
    if not sae_base:
        print("   ⚠ SKIPPED: SAE_MODELS_BASE not set")
        return False, None
    
    sae_path = os.path.join(sae_base, "finbert", "ae.pt")
    if not os.path.exists(sae_path):
        print(f"   ⚠ SKIPPED: SAE file not found: {sae_path}")
        return False, None
    
    # Optional config (same logic as api_functions.py analyze_sentiment)
    sae_config_path = None
    config_base = os.getenv("CONFIG_BASE")
    if not config_base:
        # Try cloud_api/configs first (3 levels up from api_key_util/pip/)
        config_base = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "cloud_api", "configs")
        if not os.path.exists(config_base):
            # Fallback to config/ (2 levels up from api_key_util/pip/)
            config_base = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config")
            if not os.path.exists(config_base):
                config_base = None
    if config_base:
        sae_config_path = os.path.join(config_base, "finbert_sae_config.json")
        if not os.path.exists(sae_config_path):
            sae_config_path = None  # Optional config
    
    result = engine.extract_top_features(
        "Strong quarterly earnings exceeded expectations",
        model_path,
        sae_path,
        10,
        top_n=5,
        sae_config_path=sae_config_path
    )
    
    top_features = result.get('top_features', [])
    print(f"   ✓ Features: {len(top_features)}, Device: {result.get('device', 'unknown')}")
    if top_features:
        top3 = [(f.get('feature_id'), round(f.get('activation', 0), 4)) for f in top_features[:3]]
        print(f"   Top 3: {top3}")
    return True, model_path

def test_unload_model(model_path):
    """Test engine.unload_model()"""
    print("\n3. unload_model()")
    
    if not model_path:
        print("   ⚠ SKIPPED: No model to unload (extract_top_features failed)")
        return False
    
    result = engine.unload_model(model_path)
    if result.get('success'):
        print(f"   ✓ Model unloaded")
        return True
    else:
        print(f"   ✗ Failed to unload model")
        return False

def test_unload_all_models():
    """Test engine.unload_all_models()"""
    print("\n4. unload_all_models()")
    
    result = engine.unload_all_models()
    # Check for success in various possible response formats
    if result.get('success') or result.get('status') == 'success' or 'success' in str(result).lower():
        print(f"   ✓ All models unloaded")
        return True
    else:
        print(f"   ⚠ Response: {result}")
        return True  # Still consider it a pass if it doesn't error

def test_common():
    """Test all 4 neuronlens common functions"""
    print("=== Testing NeuronLens Common Functions ===\n")
    
    results = []
    
    # Test all 4 functions
    results.append(("model_cache_status", test_model_cache_status()))
    success, model_path = test_extract_top_features()
    results.append(("extract_top_features", success))
    results.append(("unload_model", test_unload_model(model_path)))
    results.append(("unload_all_models", test_unload_all_models()))
    
    # Summary
    print("\n" + "="*60)
    print("=== Test Summary ===")
    for name, success in results:
        status = "✓ PASS" if success else "✗ FAIL"
        print(f"{status}: {name}")
    
    all_passed = all(r[1] for r in results)
    print("\n" + ("✓ All tests passed!" if all_passed else "✗ Some tests failed"))

if __name__ == "__main__":
    test_common()
