import logging
import os
from typing import Generator, Optional

from langchain_openai import ChatOpenAI
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_aws import ChatBedrock
from langchain_core.messages import HumanMessage

from sonika_ai_toolkit.utilities.types import ILanguageModel


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
        """
        from pydantic import SecretStr
        self.model = ChatOpenAI(
            temperature=temperature, 
            model=model_name, 
            api_key=SecretStr(api_key), 
            stream_usage=True
        )
        self.supports_thinking = False

    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
        """
        res = self.model.predict(prompt)
        return str(res)
    
    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])
        content = response.content
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts = []
            for p in content:
                if isinstance(p, str):
                    parts.append(p)
                elif isinstance(p, dict):
                    if p.get("type") != "thinking":
                        parts.append(str(p.get("text") or p.get("content") or ""))
            return "\n".join(parts).strip()
        return str(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]):
            content = chunk.content
            if isinstance(content, str):
                yield content
            elif isinstance(content, list):
                parts = []
                for p in content:
                    if isinstance(p, str):
                        parts.append(p)
                    elif isinstance(p, dict):
                        if p.get("type") != "thinking":
                            parts.append(str(p.get("text") or p.get("content") or ""))
                yield "\n".join(parts)
            else:
                yield str(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=model_name,
            region_name=region_name,
            model_kwargs={"temperature": temperature}
        )
        self.supports_thinking = False

    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])
        content = response.content
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts = [str(p.get("text") or p.get("content") or p if isinstance(p, dict) else p) for p in content if not (isinstance(p, dict) and p.get("type") == "thinking")]
            return "\n".join(parts).strip()
        return str(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]):
            content = chunk.content
            if isinstance(content, str):
                yield content
            elif isinstance(content, list):
                parts = [str(p.get("text") or p.get("content") or p if isinstance(p, dict) else p) for p in content if not (isinstance(p, dict) and p.get("type") == "thinking")]
                yield "\n".join(parts)
            else:
                yield str(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,
        thinking_budget: Optional[int] = None,
    ):
        """
        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
            thinking_budget (int, optional): Presupuesto de tokens para modelos con thinking nativo.
        """
        logger = logging.getLogger(__name__)
        self.supports_thinking = any(
            marker in model_name.lower()
            for marker in ["-thinking-", "thinking-exp"]
        )

        effective_temperature = temperature
        if self.supports_thinking and temperature != 1:
            logger.warning(
                "Gemini thinking models require temperature=1. Overriding provided value %s.",
                temperature,
            )
            effective_temperature = 1.0

        model_kwargs: dict = {}
        if self.supports_thinking:
            budget = thinking_budget if thinking_budget is not None else 8192
            model_kwargs["thinking_budget"] = budget
            model_kwargs["include_thoughts"] = True

        self.model = ChatGoogleGenerativeAI(
            temperature=effective_temperature,
            model=model_name,
            google_api_key=api_key,
            **model_kwargs,
        )

    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])
        content = response.content
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts = [str(p.get("text") or p.get("content") or p if isinstance(p, dict) else p) for p in content if not (isinstance(p, dict) and p.get("type") == "thinking")]
            return "\n".join(parts).strip()
        return str(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]):
            content = chunk.content
            if isinstance(content, str):
                yield content
            elif isinstance(content, list):
                parts = [str(p.get("text") or p.get("content") or p if isinstance(p, dict) else p) for p in content if not (isinstance(p, dict) and p.get("type") == "thinking")]
                yield "\n".join(parts)
            else:
                yield str(content)


class _DeepSeekReasonerChatModel(ChatOpenAI):
    """ChatOpenAI subclass that captures reasoning_content for DeepSeek R1."""

    def _create_chat_result(self, response, generation_info=None):
        result = super()._create_chat_result(response, generation_info)
        try:
            if not isinstance(response, dict):
                response = response.model_dump()
            for i, choice in enumerate(response.get("choices", [])):
                reasoning = (choice.get("message") or {}).get("reasoning_content")
                if reasoning and i < len(result.generations):
                    result.generations[i].message.additional_kwargs["reasoning_content"] = reasoning
        except Exception:
            pass
        return result


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
        """
        is_reasoner = model_name == "deepseek-reasoner" or "r1" in model_name.lower()
        model_class = _DeepSeekReasonerChatModel if is_reasoner else ChatOpenAI
        from pydantic import SecretStr
        self.model = model_class(
            temperature=temperature,
            model=model_name,
            api_key=SecretStr(api_key),
            base_url="https://api.deepseek.com"
        )
        self.supports_thinking = is_reasoner

    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])
        content = response.content
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts = [str(p.get("text") or p.get("content") or p if isinstance(p, dict) else p) for p in content if not (isinstance(p, dict) and p.get("type") == "thinking")]
            return "\n".join(parts).strip()
        return str(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]):
            content = chunk.content
            if isinstance(content, str):
                yield content
            elif isinstance(content, list):
                parts = [str(p.get("text") or p.get("content") or p if isinstance(p, dict) else p) for p in content if not (isinstance(p, dict) and p.get("type") == "thinking")]
                yield "\n".join(parts)
            else:
                yield str(content)
