import os
import yaml
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, Tuple
from dotenv import load_dotenv

from aeiva.config.base_config import BaseConfig


@dataclass
class LLMGatewayConfig(BaseConfig):
    """
    Configuration for the Language Model (LLM).
    """

    llm_model_name: Optional[str] = field(
        default='gpt-4',
        metadata={"help": "The name of the LLM model to use (e.g., 'gpt-4', 'gpt-3.5-turbo')."}
    )
    llm_api_key: Optional[str] = field(
        default=None,
        metadata={"help": "The API key for authentication with the LLM provider."}
    )
    llm_base_url: Optional[str] = field(
        default=None,
        metadata={"help": "The base URL for API requests to the LLM provider."}
    )
    llm_api_version: Optional[str] = field(
        default=None,
        metadata={"help": "The version of the LLM API to use."}
    )
    llm_embedding_model: Optional[str] = field(
        default=None,
        metadata={"help": "The embedding model to use for tasks requiring embeddings."}
    )
    llm_timeout: Optional[int] = field(
        default=30,
        metadata={"help": "The timeout in seconds for API requests."}
    )
    llm_max_input_tokens: Optional[int] = field(
        default=4096,
        metadata={"help": "The maximum number of input tokens allowed in a request."}
    )
    llm_max_output_tokens: Optional[int] = field(
        default=1024,
        metadata={"help": "The maximum number of output tokens generated by the LLM."}
    )
    llm_temperature: Optional[float] = field(
        default=0.7,
        metadata={"help": "Sampling temperature for response variability (range: 0.0 - 1.0)."}
    )
    llm_top_p: Optional[float] = field(
        default=0.9,
        metadata={"help": "Nucleus sampling probability for token selection (range: 0.0 - 1.0)."}
    )
    llm_num_retries: Optional[int] = field(
        default=3,
        metadata={"help": "The number of times to retry failed API requests."}
    )
    llm_retry_backoff_factor: Optional[float] = field(
        default=0.5,
        metadata={"help": "Factor for exponential backoff between retries."}
    )
    llm_retry_on_status: Optional[Tuple[int, ...]] = field(
        default=(429, 500, 502, 503, 504),
        metadata={"help": "HTTP status codes that should trigger a retry."}
    )
    llm_use_async: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to use asynchronous API calls."}
    )
    llm_stream: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to enable streaming responses from the LLM."}
    )
    llm_logging_level: Optional[str] = field(
        default='INFO',
        metadata={"help": "Logging level for the LLM module (e.g., 'DEBUG', 'INFO')."}
    )
    llm_additional_params: Optional[Dict[str, Any]] = field(
        default_factory=dict,
        metadata={"help": "Additional parameters to pass to the LLM API."}
    )

    def __post_init__(self):
        super().__post_init__()
        # Load API keys from the configuration file if not provided
        if not self.llm_api_key:
            self.load_api_key()

    def load_api_key(self):
        config_path = os.path.join(os.path.dirname(__file__), '../../../configs/llm_api_keys.yaml')
        try:
            with open(config_path, 'r') as f:
                keys = yaml.safe_load(f)
                self.llm_api_key = keys.get('openai_api_key')
        except FileNotFoundError:
            raise FileNotFoundError('API keys file not found.')
        except Exception as e:
            raise e

    def to_dict(self):
        return {
            key: ('******' if key == 'llm_api_key' and value else value)
            for key, value in self.__dict__.items()
            if not key.startswith('_')
        }