"""LangChain-based LLM implementation for Fluxibly.

Wraps LangChain ChatModels for multi-provider support.
Auto-detects provider from model name.
"""

from __future__ import annotations

from typing import Any

from fluxibly.llm.base import BaseLLM, LLMConfig
from fluxibly.schema.response import (
    ContentItem,
    LLMMetadata,
    LLMOutput,
    LLMResponse,
    TokenUsage,
)
from fluxibly.schema.tools import ToolCall

try:
    from langchain_core.language_models import BaseChatModel
    from langchain_core.messages import (
        AIMessage,
        HumanMessage,
        SystemMessage,
        ToolMessage,
    )
except ImportError as e:
    raise ImportError(
        "LangChain packages are required for LangChainLLM. "
        "Install them with: pip install fluxibly[langchain]"
    ) from e


class LangChainLLM(BaseLLM):
    """LangChain-based LLM implementation.

    Wraps LangChain ChatModels for multi-provider support.
    Auto-detects provider from model name.
    """

    def __init__(self, config: LLMConfig) -> None:
        super().__init__(config)
        self.chat_model = self._create_chat_model()

    def _create_chat_model(self) -> BaseChatModel:
        """Create LangChain ChatModel based on model name."""
        model = self.config.model.lower()
        common_kwargs: dict[str, Any] = {
            "model": self.config.model,
        }
        if self.config.temperature is not None:
            common_kwargs["temperature"] = self.config.temperature
        if self.config.max_output_tokens:
            common_kwargs["max_tokens"] = self.config.max_output_tokens
        if self.config.api_key:
            common_kwargs["api_key"] = self.config.api_key

        if (
            "gpt" in model
            or "o1" in model
            or "o3" in model
            or model.startswith("openai/")
        ):
            from langchain_openai import ChatOpenAI

            return ChatOpenAI(**common_kwargs)

        if "claude" in model or model.startswith("anthropic/"):
            from langchain_anthropic import ChatAnthropic

            return ChatAnthropic(**common_kwargs)

        if "gemini" in model or model.startswith("google/"):
            from langchain_google_genai import ChatGoogleGenerativeAI

            return ChatGoogleGenerativeAI(**common_kwargs)

        # Default to OpenAI-compatible
        import logging as _logging

        _logging.getLogger(__name__).warning(
            "Unrecognized model '%s' — defaulting to ChatOpenAI. "
            "Set framework explicitly if this is unintended.",
            self.config.model,
        )
        from langchain_openai import ChatOpenAI

        if self.config.api_base:
            common_kwargs["base_url"] = self.config.api_base
        return ChatOpenAI(**common_kwargs)

    def prepare(self) -> dict[str, Any]:
        """Convert unified messages to LangChain message objects."""
        lc_messages = []
        for msg in self.messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")

            if role == "system":
                lc_messages.append(SystemMessage(content=content))
            elif role == "user":
                lc_messages.append(HumanMessage(content=content))
            elif role == "assistant":
                lc_messages.append(AIMessage(content=content))
            elif role == "tool":
                lc_messages.append(
                    ToolMessage(
                        content=content,
                        tool_call_id=msg.get("tool_call_id", ""),
                    )
                )

        prepared: dict[str, Any] = {"messages": lc_messages}

        # Bind tools if available
        if self.tools:
            formatted_tools = []
            for tool in self.tools:
                if tool.get("type") == "function":
                    formatted_tools.append(tool["function"])
            if formatted_tools:
                prepared["tools"] = formatted_tools

        return prepared

    def forward(self, **kwargs: Any) -> LLMResponse:
        """Execute LangChain inference."""
        prepared = self.prepare()

        self.logger.log_input(
            {
                "model": self.config.model,
                "message_count": len(prepared["messages"]),
            }
        )

        try:
            chat_model = self.chat_model
            if "tools" in prepared:
                chat_model = chat_model.bind_tools(prepared["tools"])

            response = chat_model.invoke(prepared["messages"])
            result = self._parse_output(response)
            self._apply_pricing(result)
            self.logger.log_output(result)
            return result

        except Exception as e:
            self.logger.log_error(e)
            raise

    # ── Response Parsing ─────────────────────────────────────────────

    def _parse_output(self, response: Any) -> LLMResponse:
        """Parse LangChain AIMessage to LLMResponse."""
        content_items: list[ContentItem] = []
        tool_calls: list[ToolCall] = []
        text_parts: list[str] = []

        # Extract text
        if isinstance(response.content, str):
            if response.content:
                text_parts.append(response.content)
                content_items.append(
                    ContentItem(type="text", text=response.content)
                )
        elif isinstance(response.content, list):
            for block in response.content:
                if isinstance(block, str):
                    text_parts.append(block)
                    content_items.append(ContentItem(type="text", text=block))
                elif isinstance(block, dict) and block.get("type") == "text":
                    text_parts.append(block["text"])
                    content_items.append(
                        ContentItem(type="text", text=block["text"])
                    )

        # Extract tool calls
        if hasattr(response, "tool_calls") and response.tool_calls:
            for tc in response.tool_calls:
                tool_call = ToolCall(
                    id=tc.get("id", ""),
                    name=tc["name"],
                    arguments=tc.get("args", {}),
                    type="function",
                )
                tool_calls.append(tool_call)
                content_items.append(
                    ContentItem(type="tool_call", tool_call=tool_call)
                )

        # Extract usage
        usage_metadata = getattr(response, "usage_metadata", {}) or {}

        return LLMResponse(
            output=LLMOutput(
                output_text="\n".join(text_parts),
                content=content_items,
                tool_calls=tool_calls,
            ),
            metadata=LLMMetadata(
                model=self.config.model,
                id=getattr(response, "id", "") or "",
                created_at="",
                status="completed",
                usage=TokenUsage(
                    input_tokens=usage_metadata.get("input_tokens", 0),
                    output_tokens=usage_metadata.get("output_tokens", 0),
                    total_tokens=usage_metadata.get("total_tokens", 0),
                ),
                stop_reason=self._determine_stop_reason(response),
                provider="langchain",
            ),
        )

    def _determine_stop_reason(self, response: Any) -> str:
        """Determine stop reason from LangChain response."""
        if hasattr(response, "tool_calls") and response.tool_calls:
            return "tool_calls"
        finish = getattr(response, "response_metadata", {}).get(
            "finish_reason", "stop"
        )
        return finish if isinstance(finish, str) else "stop"
