"""
NeuronLens Engine class.
Wraps all frontend functions from sae_app_v4.api_functions with API key management.
Works standalone when sae_app_v4 is not available (PyPI installation).
"""
import sys
import os
import asyncio
from pathlib import Path
from typing import Dict, List, Optional, Any
import requests

# Add sae_app_v4 parent directory to path so we can import from it
# neuronlens is in sae_app_v4/pip, so parent.parent is sae_app_v4
_interp_dir = Path(__file__).parent.parent.parent
if str(_interp_dir) not in sys.path:
    sys.path.insert(0, str(_interp_dir))

# Import from local package (use relative imports for development)
# When installed as package, these will be available as neuronlens.config, etc.
from .config import get_default_api_url, get_api_timeout, get_api_key_from_env
from .api_client import APIClient

# Try to import base functions from sae_app_v4 (for local development)
# If not available, we'll make HTTP requests directly (for PyPI installation)
_USE_SAE_APP_FUNCTIONS = False
try:
    from sae_app_v4.api_functions import (
        analyze_agent as _analyze_agent,
        analyze_sentiment as _analyze_sentiment,
        analyze_trading as _analyze_trading,
        analyze_reasoning as _analyze_reasoning,
        analyze_hallucination as _analyze_hallucination,
        search_features as _search_features,
        steer_features as _steer_features,
        extract_top_features as _extract_top_features_async,
    )
    from sae_app_v4.backend_functions import extract_top_features as _extract_top_features_sync
    from sae_app_v4.services.model_loader import (
        unload_model as _unload_model,
        clear_model_cache as _clear_model_cache,
        get_model_cache_status as _get_model_cache_status
    )
    _USE_SAE_APP_FUNCTIONS = True
except ImportError:
    # sae_app_v4 not available - will use direct HTTP requests
    _USE_SAE_APP_FUNCTIONS = False
    import requests


class Engine:
    """
    NeuronLens Engine for accessing interpretability analysis functions.
    
    Usage:
        # With API key (required)
        engine = Engine(api_key="nlive_your_key")
        result = engine.analyze_agent("Get Tesla stock price")
        
        # API key from environment
        engine = Engine()  # Reads from NEURONLENS_API_KEY env var
    """
    
    def __init__(self, api_key: Optional[str] = None, api_url: Optional[str] = None):
        """
        Initialize NeuronLens Engine.
        
        Args:
            api_key: Required API key. If None, checks NEURONLENS_API_KEY environment variable.
                    API key is required for all functions.
            api_url: Optional API URL. If None, uses default based on api_key.
                    For cloud mode, should be the cloud API URL (frontend server on port 8001).
                    For local mode, defaults to http://localhost:8001
        """
        # Get API key from parameter or environment
        if api_key is None:
            api_key = get_api_key_from_env()
        
        # Validate API key is provided (required for all functions)
        if not api_key:
            raise ValueError(
                "API key is required. Please provide api_key parameter or set NEURONLENS_API_KEY environment variable."
            )
        
        self.api_key = api_key
        
        # Determine API URL
        if api_url is None:
            api_url = get_default_api_url(self.api_key)
            if api_url is None and self.api_key:
                raise ValueError(
                    "API key provided but no API URL found. "
                    "Please provide api_url parameter or set NEURONLENS_API_URL environment variable."
                )
        
        self.api_url = api_url
        self.timeout = get_api_timeout()
        
        # Create API client for request injection
        self._api_client = APIClient(api_key=self.api_key)
    
    def _call_with_api_key(self, func, *args, **kwargs):
        """
        Call a function with API key injection via requests.post patching.
        Handles both sync and async functions.
        
        Args:
            func: Function to call
            *args: Positional arguments
            **kwargs: Keyword arguments (api_url will be set if not provided)
        
        Returns:
            Function result
        """
        # Set api_url if not provided
        if 'api_url' not in kwargs:
            kwargs['api_url'] = self.api_url
        
        # If using sae_app_v4 functions, patch requests.post to inject API key
        if _USE_SAE_APP_FUNCTIONS:
            with self._api_client.patch_requests():
                # Check if function is async
                if asyncio.iscoroutinefunction(func):
                    return asyncio.run(func(*args, **kwargs))
                else:
                    return func(*args, **kwargs)
        else:
            # Direct HTTP request - API key already injected via api_client
            return func(*args, **kwargs)
    
    def _make_request(self, endpoint: str, json_data: dict, method: str = "POST") -> dict:
        """
        Make HTTP request directly to API endpoint with API key injection.
        Used when sae_app_v4 is not available (PyPI installation).
        
        Args:
            endpoint: API endpoint path (e.g., "/frontend/analyze_sentiment")
            json_data: Request JSON data
            method: HTTP method (POST or GET)
        
        Returns:
            Response JSON as dictionary
        """
        url = f"{self.api_url}{endpoint}"
        headers = {}
        if self.api_key:
            headers['X-API-Key'] = self.api_key
        
        if method == "GET":
            response = requests.get(url, headers=headers, timeout=self.timeout)
        else:
            response = requests.post(url, json=json_data, headers=headers, timeout=self.timeout)
        response.raise_for_status()
        return response.json()
    
    def analyze_agent(self, query: str, response: Optional[str] = None,
                      model_path: Optional[str] = None, sae_path: Optional[str] = None,
                      layer: int = 19, api_url: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        Analyze agent query with tool-intent tracing and R/A/G alignment.
        
        Args:
            query: User query to analyze
            response: Optional agent response (if None, will generate)
            model_path: Optional model path override
            sae_path: Optional SAE path override
            layer: Layer number (default: 19)
            api_url: Optional API URL override
            **kwargs: Additional arguments passed to base function
        
        Returns:
            Dictionary with alignment_status, tool_called, intent_scores, top_features, etc.
        """
        if _USE_SAE_APP_FUNCTIONS:
            return self._call_with_api_key(
                _analyze_agent, 
                query, 
                response=response,
                model_path=model_path,
                sae_path=sae_path,
                layer=layer,
                api_url=api_url,
                **kwargs
            )
        else:
            return self._make_request(
                "/frontend/analyze_agent", 
                {
                    "query": query,
                    "response": response,
                    "model_path": model_path,
                    "sae_path": sae_path,
                    "layer": layer,
                    **kwargs
                }
            )
    
    def analyze_sentiment(self, text: str, top_n: int = 10, model_path: Optional[str] = None,
                         sae_path: Optional[str] = None, layer: int = 10, 
                         api_url: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        Analyze sentiment using SLM Lens with SAE feature attribution.
        
        Args:
            text: Text to analyze
            top_n: Number of top features to return (only used in HTTP mode; direct mode uses hardcoded 10)
            model_path: Optional model path override
            sae_path: Optional SAE path override
            layer: Layer number (default: 10)
            api_url: Optional API URL override
            **kwargs: Additional arguments passed to base function
        
        Returns:
            Dictionary with sentiment (predicted_label, confidence, probabilities) and top_features
        """
        if _USE_SAE_APP_FUNCTIONS:
            # Note: _analyze_sentiment doesn't accept top_n parameter (hardcoded to 10)
            return self._call_with_api_key(
                _analyze_sentiment, 
                text, 
                model_path=model_path,
                sae_path=sae_path,
                layer=layer,
                api_url=api_url,
                **kwargs
            )
        else:
            return self._make_request(
                "/frontend/analyze_sentiment",
                {
                    "text": text,
                    "top_n": top_n,
                    "model_path": model_path,
                    "sae_path": sae_path,
                    "layer": layer,
                    **kwargs
                }
            )
    
    def analyze_trading(self, text: str, ticker: Optional[str] = None,
                        model_path: Optional[str] = None, sae_path: Optional[str] = None,
                        layer: int = 19, api_url: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        Extract trading signals from news text.
        
        Args:
            text: News text to analyze
            ticker: Optional stock ticker symbol
            model_path: Optional model path override
            sae_path: Optional SAE path override
            layer: Layer number (default: 19)
            api_url: Optional API URL override
            **kwargs: Additional arguments passed to base function
        
        Returns:
            Dictionary with extracted signals
        """
        if _USE_SAE_APP_FUNCTIONS:
            return self._call_with_api_key(
                _analyze_trading, 
                text, 
                ticker=ticker,
                model_path=model_path,
                sae_path=sae_path,
                layer=layer,
                api_url=api_url,
                **kwargs
            )
        else:
            return self._make_request(
                "/frontend/analyze_trading",
                {
                    "text": text,
                    "ticker": ticker,
                    "model_path": model_path,
                    "sae_path": sae_path,
                    "layer": layer,
                    **kwargs
                }
            )
    
    def analyze_reasoning(self, prompt: str, model_path: Optional[str] = None, 
                         sae_path: Optional[str] = None, layer: int = 28, 
                         expected_buckets: Optional[List[str]] = None, 
                         api_url: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        Analyze reasoning with CoT faithfulness analysis.
        
        Args:
            prompt: Prompt to analyze
            model_path: Optional model path override
            sae_path: Optional SAE path override
            layer: Layer number (default: 28)
            expected_buckets: Optional list of expected reasoning buckets
            api_url: Optional API URL override
            **kwargs: Additional arguments passed to base function
        
        Returns:
            Dictionary with reasoning analysis results
        """
        if _USE_SAE_APP_FUNCTIONS:
            return self._call_with_api_key(
                _analyze_reasoning, 
                prompt, 
                model_path=model_path,
                sae_path=sae_path,
                layer=layer,
                expected_buckets=expected_buckets,
                api_url=api_url,
                **kwargs
            )
        else:
            return self._make_request(
                "/frontend/analyze_reasoning",
                {
                    "prompt": prompt,
                    "model_path": model_path,
                    "sae_path": sae_path,
                    "layer": layer,
                    "expected_buckets": expected_buckets,
                    **kwargs
                }
            )
    
    def analyze_hallucination(self, prompt: str, max_tokens: int = 256, 
                             temperature: float = 0.7, api_url: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        Analyze hallucination with token-level scoring.
        
        Args:
            prompt: Prompt to analyze
            max_tokens: Maximum tokens to generate (default: 256)
            temperature: Generation temperature (default: 0.7)
            api_url: Optional API URL override
            **kwargs: Additional arguments passed to base function
        
        Returns:
            Dictionary with token_details, hallucination scores, etc.
        """
        if _USE_SAE_APP_FUNCTIONS:
            return self._call_with_api_key(
                _analyze_hallucination, 
                prompt, 
                max_tokens=max_tokens, 
                temperature=temperature,
                api_url=api_url,
                **kwargs
            )
        else:
            return self._make_request(
                "/frontend/analyze_hallucination",
                {
                    "prompt": prompt,
                    "max_tokens": max_tokens,
                    "temperature": temperature,
                    **kwargs
                }
            )
    
    def search_features(self, query: str, k: int = 10, layer: int = 16, 
                       api_url: Optional[str] = None, **kwargs) -> List[Dict[str, Any]]:
        """
        Search for features using semantic similarity.
        
        Args:
            query: Search query
            k: Number of features to return (default: 10)
            layer: Layer to search in (default: 16)
            api_url: Optional API URL override
            **kwargs: Additional arguments passed to base function
        
        Returns:
            List of dictionaries with feature_id, score, and label
        """
        if _USE_SAE_APP_FUNCTIONS:
            return self._call_with_api_key(_search_features, query, k=k, layer=layer, api_url=api_url, **kwargs)
        else:
            result = self._make_request(
                "/frontend/search_features",
                {"query": query, "k": k, "layer": layer, **kwargs}
            )
            # API returns a list directly, not a dict with "features" key
            if isinstance(result, list):
                return result
            return result.get("features", [])
    
    def steer_features(self, prompt: str, features: List[Dict[str, Any]], 
                      model: str = "meta-llama/Llama-2-7b-hf", api_url: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        Apply feature steering and generate text.
        
        Args:
            prompt: Prompt to generate from
            features: List of feature dictionaries with "id" and "magnitude" keys
                     Magnitude range: -1.0 to +1.0 (positive = steer towards, negative = steer away)
            model: Model to use for generation (default: "meta-llama/Llama-2-7b-hf")
            api_url: Optional API URL override
            **kwargs: Additional arguments passed to base function
        
        Returns:
            Dictionary with success, original_text, and steered_text
        """
        if _USE_SAE_APP_FUNCTIONS:
            return self._call_with_api_key(_steer_features, prompt, features, model=model, api_url=api_url, **kwargs)
        else:
            return self._make_request(
                "/frontend/steer_features",
                {"prompt": prompt, "features": features, "model": model, **kwargs}
            )
    
    # Common backend wrapper functions
    def extract_top_features(self, text: str, model_path: str, sae_path: str, 
                            layer: int, top_n: int = 10, use_hf_model: bool = False,
                            use_hf_sae: bool = False, sae_config_path: Optional[str] = None,
                            api_url: Optional[str] = None, **kwargs) -> Dict[str, Any]:
        """
        Extract top SAE features from text.
        
        Args:
            text: Input text
            model_path: Path to model (local or HuggingFace)
            sae_path: Path to SAE (local or HuggingFace)
            layer: Layer number
            top_n: Number of top features to return (default: 10)
            use_hf_model: Whether to use HuggingFace model (default: False)
            use_hf_sae: Whether to use HuggingFace SAE (default: False)
            sae_config_path: Optional SAE config path
            api_url: Optional API URL override
            **kwargs: Additional arguments
        
        Returns:
            Dictionary with features and top_features
        """
        # Always use HTTP mode for extract_top_features (like analyze_sentiment does)
        # Match backend ExtractTopFeaturesRequest format exactly
        # Frontend wrapper proxies to backend, so we can use frontend endpoint in HTTP mode
        if _USE_SAE_APP_FUNCTIONS:
            # Use async extract_top_features from api_functions which calls backend directly
            # It automatically calculates use_hf_model/use_hf_sae, so we don't pass them
            backend_api_url = api_url
            if not backend_api_url:
                # Convert frontend URL to backend URL if needed (api_functions uses backend)
                if self.api_url and ":8001" in self.api_url:
                    backend_api_url = self.api_url.replace(":8001", ":8000")
                else:
                    # Default to backend URL
                    backend_api_url = os.getenv("NEURONLENS_API_URL", "http://localhost:8000")
            
            return self._call_with_api_key(
                _extract_top_features_async,
                text, model_path, sae_path, layer,
                top_n=top_n, sae_config_path=sae_config_path, api_url=backend_api_url
            )
        else:
            # HTTP mode: use frontend endpoint (which proxies to backend, same as test_common.py)
            # Build request matching backend ExtractTopFeaturesRequest format
            request_data = {
                "text": text,
                "model_path": model_path,
                "sae_path": sae_path,
                "layer": layer,
                "top_n": top_n,
                "use_hf_model": use_hf_model,
                "use_hf_sae": use_hf_sae,
            }
            if sae_config_path:
                request_data["sae_config_path"] = sae_config_path
            
            return self._make_request(
                "/frontend/common/extract_top_features",
                request_data
            )
    
    def unload_model(self, model_path: str, device: Optional[str] = None) -> Dict[str, Any]:
        """
        Unload a specific model from cache.
        
        Args:
            model_path: Path to model to unload
            device: Optional device (if None, unloads from all devices)
        
        Returns:
            Dictionary with success status
        """
        if _USE_SAE_APP_FUNCTIONS:
            # Direct mode: call function directly
            _unload_model(model_path, device=device)
            return {"success": True, "model_path": model_path}
        else:
            # HTTP mode: make request to frontend
            return self._make_request(
                "/frontend/common/unload_model",
                {"model_path": model_path, "device": device}
            )
    
    def unload_all_models(self) -> Dict[str, Any]:
        """
        Unload all models from cache.
        
        Returns:
            Dictionary with success status
        """
        if _USE_SAE_APP_FUNCTIONS:
            # Direct mode: call function directly
            _clear_model_cache()
            return {"success": True}
        else:
            # HTTP mode: make request to frontend
            return self._make_request(
                "/frontend/common/unload_all_models",
                {}
            )
    
    def model_cache_status(self) -> Dict[str, Any]:
        """
        Get model cache status.
        
        Returns:
            Dictionary with cache status information
        """
        if _USE_SAE_APP_FUNCTIONS:
            # Direct mode: call function directly
            return _get_model_cache_status()
        else:
            # HTTP mode: make GET request to frontend
            return self._make_request(
                "/frontend/common/model_cache_status",
                {},
                method="GET"
            )
