import os
import logging
from typing import List, Optional, Union, Any, Dict
from dataclasses import dataclass

try:
    import voyageai
except ImportError:
    # We'll handle this gracefully if the user doesn't have it installed
    voyageai = None

@dataclass
class VoyageConfig:
    """Configuration for Voyage AI helper."""
    api_key: Optional[str] = None
    embedding_model: str = "voyage-3.5"
    base_url: Optional[str] = None
    timeout: float = 60.0
    enable_logging: bool = True
    log_level: str = "INFO"

    def __post_init__(self):
        if not self.api_key:
            self.api_key = os.getenv("VOYAGE_API_KEY")

class VoyageHelper:
    """Helper for Voyage AI embeddings."""
    
    def __init__(self, config: Optional[VoyageConfig] = None, **kwargs):
        if config is None:
            # Allow passing parameters directly via kwargs
            self.config = VoyageConfig(
                api_key=kwargs.get("api_key"),
                embedding_model=kwargs.get("embedding_model", "voyage-3.5"),
                base_url=kwargs.get("base_url")
            )
        else:
            self.config = config
            
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(self.config.log_level)
        
        if voyageai is None:
            self.logger.error("voyageai package not found. Please install with: pip install voyageai")
            self.client = None
        else:
            self.client = voyageai.Client(api_key=self.config.api_key)

    def _setup_client(self):
        """Re-initialize Voyage AI client."""
        if voyageai:
            try:
                self.client = voyageai.Client(api_key=self.config.api_key)
            except Exception as e:
                self.logger.error(f"Failed to re-initialize Voyage client: {e}")

    def _retry(self, fn: Callable, *args, **kwargs):
        """Retry wrapper with self-healing."""
        import time
        max_retries = 3
        delay = 1.0
        for attempt in range(max_retries + 1):
            try:
                return fn(*args, **kwargs)
            except AttributeError as e:
                if "'NoneType' object has no attribute 'request'" in str(e):
                    self.logger.warning(f"Voyage client internal state is broken (attempt {attempt + 1}), re-initializing...")
                    self._setup_client()
                    if attempt == max_retries:
                        raise
                    continue
                raise
            except Exception as e:
                if attempt == max_retries:
                    raise
                self.logger.warning(f"Voyage API error: {e}. Retrying in {delay}s...")
                time.sleep(delay)
                delay *= 2

    def embed(self, texts: List[str], model: Optional[str] = None, **kwargs) -> List[List[float]]:
        """Generate embeddings for a list of texts."""
        if self.client is None:
            raise ImportError("voyageai package not found.")
            
        model = model or self.config.embedding_model
        input_type = kwargs.get("input_type", "document")
        
        from .tracing_helper import trace_span
        with trace_span("voyage_embed", {"model": model, "input_type": input_type}):
            def _make_request():
                return self.client.embed(
                    texts,
                    model=model,
                    input_type=input_type,
                    **kwargs
                ).embeddings
    
            try:
                return self._retry(_make_request)
            except Exception as e:
                self.logger.error(f"Error generating Voyage AI embeddings: {e}")
                raise

    def embed_single(self, text: str, model: Optional[str] = None, **kwargs) -> List[float]:
        """Generate embedding for a single text."""
        embeddings = self.embed([text], model=model, **kwargs)
        return embeddings[0] if embeddings else []

    def get_embeddings(self, texts: List[str], model: Optional[str] = None, **kwargs) -> List[List[float]]:
        """Alias for embed for compatibility with other helpers."""
        return self.embed(texts, model=model, **kwargs)

    def get_embedding(self, text: str, model: Optional[str] = None, **kwargs) -> List[float]:
        """Alias for embed_single for compatibility."""
        return self.embed_single(text, model=model, **kwargs)
