import os
from langchain_openai import ChatOpenAI
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_aws import ChatBedrock
from sonika_ai_toolkit.utilities.types import ILanguageModel
from langchain_core.messages import HumanMessage
from typing import Generator


class OpenAILanguageModel(ILanguageModel):
    """
    Clase que implementa la interfaz ILanguageModel para interactuar con los modelos de lenguaje de OpenAI.
    Proporciona funcionalidades para generar respuestas y contar tokens.
    """

    def __init__(self, api_key: str, model_name: str = "gpt-4o-mini", temperature: float = 0.7):
        """
        Inicializa el modelo de lenguaje de OpenAI.
        
        Args:
            api_key (str): Clave API de OpenAI
            model_name (str): Nombre del modelo a utilizar
            temperature (float): Temperatura para la generación de respuestas
        """
        self.model = ChatOpenAI(temperature=temperature, model_name=model_name, api_key=api_key)

    def predict(self, prompt: str) -> str:
        """
        Genera una respuesta basada en el prompt proporcionado.
        
        Args:
            prompt (str): Texto de entrada para generar la respuesta
            
        Returns:
            str: Respuesta generada por el modelo
        """
        return self.model.predict(prompt)
    
    def invoke(self, prompt: str) -> str:
        """
        Invokes the language model with a given prompt and returns the generated response.

        Args:
            prompt (str): The input text to be processed by the language model.

        Returns:
            str: The response generated by the language model based on the provided prompt.
        """
        message = HumanMessage(content=prompt)
        response = self.model.invoke([message])
        return response.content
    
    def stream_response(self, prompt: str) -> Generator[str, None, None]:
        """
        Genera una respuesta en streaming basada en el prompt proporcionado.

        Args:
            prompt (str): Texto de entrada para generar la respuesta
        
        Yields:
            str: Fragmentos de la respuesta generada por el modelo en tiempo real
        """
        message = HumanMessage(content=prompt)
        for chunk in self.model.stream([message]):
            yield chunk.content


class BedrockLanguageModel(ILanguageModel):
    """
    Clase que implementa la interfaz ILanguageModel para interactuar con los modelos de Amazon Bedrock.
    Proporciona funcionalidades para generar respuestas y contar tokens.
    """

    def __init__(self, api_key: str, region_name: str, model_name: str = "amazon.nova-micro-v1:0", temperature: float = 0.7):
        """
        Inicializa el modelo de lenguaje de Amazon Bedrock.

        Args:
            api_key (str): API Key de Amazon Bedrock (AWS_BEARER_TOKEN_BEDROCK)
            region_name (str): AWS Region (ej: us-east-1)
            model_name (str): ID del modelo en Bedrock (ej: amazon.nova-micro-v1:0)
            temperature (float): Temperatura para la generación de respuestas
        """
        # Configurar la variable de entorno necesaria para langchain-aws
        os.environ["AWS_BEARER_TOKEN_BEDROCK"] = api_key

        self.model = ChatBedrock(
            model_id=model_name,
            region_name=region_name,
            model_kwargs={"temperature": temperature}
        )

    def predict(self, prompt: str) -> str:
        """
        Genera una respuesta basada en el prompt proporcionado.

        Args:
            prompt (str): Texto de entrada para generar la respuesta

        Returns:
            str: Respuesta generada por el modelo
        """
        return self.model.predict(prompt)

    def invoke(self, prompt: str) -> str:
        """
        Invokes the language model with a given prompt and returns the generated response.

        Args:
            prompt (str): The input text to be processed by the language model.

        Returns:
            str: The response generated by the language model based on the provided prompt.
        """
        message = HumanMessage(content=prompt)
        response = self.model.invoke([message])
        return response.content

    def stream_response(self, prompt: str) -> Generator[str, None, None]:
        """
        Genera una respuesta en streaming basada en el prompt proporcionado.

        Args:
            prompt (str): Texto de entrada para generar la respuesta

        Yields:
            str: Fragmentos de la respuesta generada por el modelo en tiempo real
        """
        message = HumanMessage(content=prompt)
        for chunk in self.model.stream([message]):
            yield chunk.content


class GeminiLanguageModel(ILanguageModel):
    """
    Clase que implementa la interfaz ILanguageModel para interactuar con los modelos de lenguaje de Gemini (Google).
    Proporciona funcionalidades para generar respuestas y contar tokens.
    """

    def __init__(self, api_key: str, model_name: str = "gemini-3-flash-preview", temperature: float = 0.7):
        """
        Inicializa el modelo de lenguaje de Gemini.

        Args:
            api_key (str): Clave API de Google Gemini
            model_name (str): Nombre del modelo a utilizar
            temperature (float): Temperatura para la generación de respuestas
        """
        self.model = ChatGoogleGenerativeAI(
            temperature=temperature,
            model=model_name,
            google_api_key=api_key
        )

    def predict(self, prompt: str) -> str:
        """
        Genera una respuesta basada en el prompt proporcionado.

        Args:
            prompt (str): Texto de entrada para generar la respuesta

        Returns:
            str: Respuesta generada por el modelo
        """
        return self.model.predict(prompt)

    def invoke(self, prompt: str) -> str:
        """
        Invokes the language model with a given prompt and returns the generated response.

        Args:
            prompt (str): The input text to be processed by the language model.

        Returns:
            str: The response generated by the language model based on the provided prompt.
        """
        message = HumanMessage(content=prompt)
        response = self.model.invoke([message])
        return response.content

    def stream_response(self, prompt: str) -> Generator[str, None, None]:
        """
        Genera una respuesta en streaming basada en el prompt proporcionado.

        Args:
            prompt (str): Texto de entrada para generar la respuesta

        Yields:
            str: Fragmentos de la respuesta generada por el modelo en tiempo real
        """
        message = HumanMessage(content=prompt)
        for chunk in self.model.stream([message]):
            yield chunk.content


class DeepSeekLanguageModel(ILanguageModel):
    """
    Clase que implementa la interfaz ILanguageModel para interactuar con los modelos de lenguaje de DeepSeek.
    Proporciona funcionalidades para generar respuestas y contar tokens.
    """

    def __init__(self, api_key: str, model_name: str = "deepseek-chat", temperature: float = 0.7):
        """
        Inicializa el modelo de lenguaje de DeepSeek.

        Args:
            api_key (str): Clave API de DeepSeek
            model_name (str): Nombre del modelo a utilizar
            temperature (float): Temperatura para la generación de respuestas
        """
        self.model = ChatOpenAI(
            temperature=temperature,
            model_name=model_name,
            api_key=api_key,
            base_url="https://api.deepseek.com"
        )

    def predict(self, prompt: str) -> str:
        """
        Genera una respuesta basada en el prompt proporcionado.

        Args:
            prompt (str): Texto de entrada para generar la respuesta

        Returns:
            str: Respuesta generada por el modelo
        """
        return self.model.predict(prompt)

    def invoke(self, prompt: str) -> str:
        """
        Invokes the language model with a given prompt and returns the generated response.

        Args:
            prompt (str): The input text to be processed by the language model.

        Returns:
            str: The response generated by the language model based on the provided prompt.
        """
        message = HumanMessage(content=prompt)
        response = self.model.invoke([message])
        return response.content

    def stream_response(self, prompt: str) -> Generator[str, None, None]:
        """
        Genera una respuesta en streaming basada en el prompt proporcionado.

        Args:
            prompt (str): Texto de entrada para generar la respuesta

        Yields:
            str: Fragmentos de la respuesta generada por el modelo en tiempo real
        """
        message = HumanMessage(content=prompt)
        for chunk in self.model.stream([message]):
            yield chunk.content
