import importlib
import json
import warnings
from abc import abstractmethod
from typing import Any

from langchain_core.language_models import BaseChatModel
from langchain_core.language_models.llms import LLM
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
from langchain_core.output_parsers import BaseOutputParser

from lfx.base.constants import STREAM_INFO_TEXT
from lfx.custom.custom_component.component import Component
from lfx.field_typing import LanguageModel
from lfx.inputs.inputs import BoolInput, InputTypes, MessageInput, MultilineInput
from lfx.log.logger import logger
from lfx.schema.message import Message
from lfx.schema.properties import Usage
from lfx.schema.token_usage import extract_usage_from_message
from lfx.template.field.base import Output
from lfx.utils.constants import MESSAGE_SENDER_AI

# Enabled detailed thinking for NVIDIA reasoning models.
#
# Models are trained with this exact string. Do not update.
DETAILED_THINKING_PREFIX = "detailed thinking on\n\n"


def _normalize_message_content(content: Any) -> Any:
    """Flatten LangChain content blocks into a plain string.

    Gemini 3 (and other modern LangChain chat models) return ``AIMessage.content``
    as a list of content blocks such as
    ``[{"type": "text", "text": "...", "thought_signature": "..."}]`` instead of a
    plain string. Join the ``text`` blocks and drop non-text blocks so the result
    can populate ``Message.text``. Non-list inputs are returned unchanged.
    """
    if not isinstance(content, list):
        return content
    parts: list[str] = []
    for block in content:
        if isinstance(block, str):
            parts.append(block)
        elif isinstance(block, dict) and block.get("type") == "text":
            text = block.get("text")
            if isinstance(text, str):
                parts.append(text)
    return "".join(parts)


class LCModelComponent(Component):
    display_name: str = "Model Name"
    description: str = "Model Description"
    trace_type = "llm"
    metadata = {
        "keywords": [
            "model",
            "llm",
            "language model",
            "large language model",
        ],
    }

    # Optional output parser to pass to the runnable. Subclasses may allow the user to input an `output_parser`
    output_parser: BaseOutputParser | None = None

    _base_inputs: list[InputTypes] = [
        MessageInput(name="input_value", display_name="Input"),
        MultilineInput(
            name="system_message",
            display_name="System Message",
            info="System message to pass to the model.",
            advanced=False,
        ),
        BoolInput(name="stream", display_name="Stream", info=STREAM_INFO_TEXT, advanced=True),
    ]

    outputs = [
        Output(display_name="Model Response", name="text_output", method="text_response"),
        Output(display_name="Language Model", name="model_output", method="build_model"),
    ]

    def _get_exception_message(self, e: Exception):
        return str(e)

    def supports_tool_calling(self, model: LanguageModel) -> bool:
        try:
            # Check if the bind_tools method is the same as the base class's method
            if model.bind_tools is BaseChatModel.bind_tools:
                return False

            def test_tool(x: int) -> int:
                return x

            model_with_tool = model.bind_tools([test_tool])
            return hasattr(model_with_tool, "tools") and len(model_with_tool.tools) > 0
        except (AttributeError, TypeError, ValueError):
            return False

    def _validate_outputs(self) -> None:
        # At least these two outputs must be defined
        required_output_methods = ["text_response", "build_model"]
        output_names = [output.name for output in self.outputs]
        for method_name in required_output_methods:
            if method_name not in output_names:
                msg = f"Output with name '{method_name}' must be defined."
                raise ValueError(msg)
            if not hasattr(self, method_name):
                msg = f"Method '{method_name}' must be defined."
                raise ValueError(msg)

    async def text_response(self) -> Message:
        output = self.build_model()
        result = await self.get_chat_result(
            runnable=output, stream=self.stream, input_value=self.input_value, system_message=self.system_message
        )
        self.status = result
        return result

    def get_result(self, *, runnable: LLM, stream: bool, input_value: str):
        """Retrieves the result from the output of a Runnable object.

        Args:
            runnable (Runnable): The runnable to retrieve the result from.
            stream (bool): Indicates whether to use streaming or invocation mode.
            input_value (str): The input value to pass to the output object.

        Returns:
            The result obtained from the output object.
        """
        try:
            if stream:
                result = runnable.stream(input_value)
            else:
                message = runnable.invoke(input_value)
                result = message.content if hasattr(message, "content") else message
                self.status = result
        except Exception as e:
            if message := self._get_exception_message(e):
                raise ValueError(message) from e
            raise

        return result

    def build_status_message(self, message: AIMessage):
        """Builds a status message from an AIMessage object.

        Args:
            message (AIMessage): The AIMessage object to build the status message from.

        Returns:
            The status message.
        """
        if message.response_metadata:
            # Build a well formatted status message
            content = message.content
            response_metadata = message.response_metadata
            openai_keys = ["token_usage", "model_name", "finish_reason"]
            inner_openai_keys = ["completion_tokens", "prompt_tokens", "total_tokens"]
            anthropic_keys = ["model", "usage", "stop_reason"]
            inner_anthropic_keys = ["input_tokens", "output_tokens"]
            if all(key in response_metadata for key in openai_keys) and all(
                key in response_metadata["token_usage"] for key in inner_openai_keys
            ):
                token_usage = response_metadata["token_usage"]
                status_message = {
                    "tokens": {
                        "input": token_usage["prompt_tokens"],
                        "output": token_usage["completion_tokens"],
                        "total": token_usage["total_tokens"],
                        "stop_reason": response_metadata["finish_reason"],
                        "response": content,
                    }
                }

            elif all(key in response_metadata for key in anthropic_keys) and all(
                key in response_metadata["usage"] for key in inner_anthropic_keys
            ):
                usage = response_metadata["usage"]
                status_message = {
                    "tokens": {
                        "input": usage["input_tokens"],
                        "output": usage["output_tokens"],
                        "stop_reason": response_metadata["stop_reason"],
                        "response": content,
                    }
                }
            else:
                status_message = f"Response: {content}"  # type: ignore[assignment]
        else:
            status_message = f"Response: {message.content}"  # type: ignore[assignment]
        return status_message

    def extract_usage(self, message: AIMessage) -> Usage | None:
        """Extract token usage from an AIMessage's response metadata.

        Args:
            message: The AIMessage with response_metadata containing usage info.

        Returns:
            Usage with input_tokens, output_tokens, total_tokens or None if not available.
        """
        return extract_usage_from_message(message)

    async def get_chat_result(
        self,
        *,
        runnable: LanguageModel,
        stream: bool,
        input_value: str | Message,
        system_message: str | None = None,
    ) -> Message:
        # NVIDIA reasoning models use detailed thinking
        if getattr(self, "detailed_thinking", False):
            system_message = DETAILED_THINKING_PREFIX + (system_message or "")

        return await self._get_chat_result(
            runnable=runnable,
            stream=stream,
            input_value=input_value,
            system_message=system_message,
        )

    async def _get_chat_result(
        self,
        *,
        runnable: LanguageModel,
        stream: bool,
        input_value: str | Message,
        system_message: str | None = None,
    ) -> Message:
        """Get chat result from a language model.

        This method handles the core logic of getting a response from a language model,
        including handling different input types, streaming, and error handling.

        Args:
            runnable (LanguageModel): The language model to use for generating responses
            stream (bool): Whether to stream the response
            input_value (str | Message): The input to send to the model
            system_message (str | None, optional): System message to prepend. Defaults to None.

        Returns:
            The model response, either as a Message object or raw content

        Raises:
            ValueError: If the input message is empty or if there's an error during model invocation
        """
        messages: list[BaseMessage] = []
        if not input_value and not system_message:
            msg = "The message you want to send to the model is empty."
            raise ValueError(msg)
        system_message_added = False
        message = None
        if input_value:
            if isinstance(input_value, Message):
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    if "prompt" in input_value:
                        prompt = input_value.load_lc_prompt()
                        if system_message:
                            prompt.messages = [
                                SystemMessage(content=system_message),
                                *prompt.messages,  # type: ignore[has-type]
                            ]
                            system_message_added = True
                        runnable = prompt | runnable
                    else:
                        messages.append(input_value.to_lc_message(self.name))
            else:
                messages.append(HumanMessage(content=input_value))

        if system_message and not system_message_added:
            messages.insert(0, SystemMessage(content=system_message))
        inputs: list | dict = messages or {}
        lf_message = None
        usage_data = None
        try:
            # TODO: Depreciated Feature to be removed in upcoming release
            if hasattr(self, "output_parser") and self.output_parser is not None:
                runnable |= self.output_parser

            runnable = runnable.with_config(
                {
                    "run_name": self.display_name,
                    "project_name": self.get_project_name(),
                    "callbacks": self.get_langchain_callbacks(),
                }
            )
            if stream:
                lf_message, result, message = await self._handle_stream(runnable, inputs)
            else:
                message = await runnable.ainvoke(inputs)
                result = message.content if hasattr(message, "content") else message
                result = _normalize_message_content(result)
            if isinstance(message, AIMessage):
                status_message = self.build_status_message(message)
                self.status = status_message
                # Extract usage for the response
                usage_data = self.extract_usage(message)
                if usage_data:
                    self._token_usage = usage_data
            elif isinstance(result, dict):
                result = json.dumps(message, indent=4)
                self.status = result
            else:
                self.status = result
        except Exception as e:
            if message := self._get_exception_message(e):
                raise ValueError(message) from e
            raise

        if lf_message:
            if lf_message.properties and lf_message.properties.usage:
                self._token_usage = lf_message.properties.usage
            return lf_message

        # Create message with usage data if available
        msg = Message(text=result)
        if usage_data:
            msg.properties.usage = usage_data
        return msg

    async def _handle_stream(self, runnable, inputs):
        """Handle streaming responses from the language model.

        Args:
            runnable: The language model configured for streaming
            inputs: The inputs to send to the model

        Returns:
            tuple: (Message object if connected to chat output, model result, AIMessage or None)
        """
        lf_message = None
        ai_message = None
        if self.is_connected_to_chat_output():
            # Add a Message
            if hasattr(self, "graph"):
                session_id = self.graph.session_id
            elif hasattr(self, "_session_id"):
                session_id = self._session_id
            else:
                session_id = None
            # Streaming requires both a session_id and an event_manager:
            #   - session_id is required so astore_message validation passes when send_message
            #     persists the placeholder Message.
            #   - event_manager is required so the chunk iterator that backs Message.text gets
            #     drained; without one, no consumer iterates astream(), the iterator is stored
            #     verbatim, and downstream readers see empty text.
            # If either is missing, fall back to a non-streaming ainvoke.
            event_manager = getattr(self, "_event_manager", None)
            if session_id and event_manager:
                model_message = Message(
                    text=runnable.astream(inputs),
                    sender=MESSAGE_SENDER_AI,
                    sender_name="AI",
                    properties={"icon": self.icon, "state": "partial"},
                    session_id=session_id,
                )
                model_message.properties.source = self._build_source(self._id, self.display_name, self)
                lf_message = await self.send_message(model_message)
                result = lf_message.text or ""
            else:
                missing = []
                if not session_id:
                    missing.append("session_id")
                if not event_manager:
                    missing.append("event_manager")
                component_label = getattr(self, "display_name", None) or getattr(self, "_id", "<unknown>")
                logger.warning(
                    f"Streaming fallback to ainvoke for component '{component_label}' "
                    f"(id={getattr(self, '_id', '<unknown>')}): missing {', '.join(missing)}. "
                    "UI will not see token-by-token streaming for this run."
                )
                ai_message = await runnable.ainvoke(inputs)
                result = ai_message.content if hasattr(ai_message, "content") else ai_message
                result = _normalize_message_content(result)
        else:
            ai_message = await runnable.ainvoke(inputs)
            result = ai_message.content if hasattr(ai_message, "content") else ai_message
            result = _normalize_message_content(result)
        return lf_message, result, ai_message

    @abstractmethod
    def build_model(self) -> LanguageModel:  # type: ignore[type-var]
        """Implement this method to build the model."""

    def get_llm(self, provider_name: str, model_info: dict[str, dict[str, str | list[InputTypes]]]) -> LanguageModel:
        """Get LLM model based on provider name and inputs.

        Args:
            provider_name: Name of the model provider (e.g., "OpenAI", "Azure OpenAI")
            inputs: Dictionary of input parameters for the model
            model_info: Dictionary of model information

        Returns:
            Built LLM model instance
        """
        try:
            if provider_name not in [model.get("display_name") for model in model_info.values()]:
                msg = f"Unknown model provider: {provider_name}"
                raise ValueError(msg)

            # Find the component class name from MODEL_INFO in a single iteration
            component_info, module_name = next(
                ((info, key) for key, info in model_info.items() if info.get("display_name") == provider_name),
                (None, None),
            )
            if not component_info:
                msg = f"Component information not found for {provider_name}"
                raise ValueError(msg)
            component_inputs = component_info.get("inputs", [])
            # Get the component class from the models module
            # Ensure component_inputs is a list of the expected types
            if not isinstance(component_inputs, list):
                component_inputs = []

            import warnings

            with warnings.catch_warnings():
                warnings.filterwarnings(
                    "ignore", message="Support for class-based `config` is deprecated", category=DeprecationWarning
                )
                warnings.filterwarnings("ignore", message="Valid config keys have changed in V2", category=UserWarning)
                models_module = importlib.import_module("lfx.components.models")
                component_class = getattr(models_module, str(module_name))
                component = component_class()

            return self.build_llm_model_from_inputs(component, component_inputs)
        except Exception as e:
            msg = f"Error building {provider_name} language model"
            raise ValueError(msg) from e

    def build_llm_model_from_inputs(
        self, component: Component, inputs: list[InputTypes], prefix: str = ""
    ) -> LanguageModel:
        """Build LLM model from component and inputs.

        Args:
            component: LLM component instance
            inputs: Dictionary of input parameters for the model
            prefix: Prefix for the input names
        Returns:
            Built LLM model instance
        """
        # Ensure prefix is a string
        prefix = prefix or ""
        # Filter inputs to only include valid component input names
        input_data = {
            str(component_input.name): getattr(self, f"{prefix}{component_input.name}", None)
            for component_input in inputs
        }

        return component.set(**input_data).build_model()
