"""Session caller — LangChain-based LLM invocation for session conversations.

Implements the :class:`LangChainSessionCaller` that wraps a LangChain chat
model with full session history context, plus token usage extraction and
cost estimation helpers.  Used by :class:`SessionWorkflow` in
:mod:`cleveragents.application.services.session_workflow`.

Forgejo: #5784
"""

from __future__ import annotations

from collections.abc import Iterator
from typing import Any

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.messages.tool import ToolMessage

from cleveragents.domain.models.core.session import MessageRole
from cleveragents.tool.actor_runtime import LLMResponse, LLMToolCall

# Token cost estimation constants (USD per 1K tokens).
# These are rough defaults; real cost comes from the LLM response metadata.
_COST_PER_1K_INPUT = 0.001
_COST_PER_1K_OUTPUT = 0.002


# ---------------------------------------------------------------------------
# Token extraction and cost estimation helpers
# ---------------------------------------------------------------------------


def extract_content(response: Any) -> str:
    """Extract plain text content from a LangChain response object."""
    raw_content = getattr(response, "content", None)
    if raw_content is None:
        raw_content = getattr(response, "text", None)
    if raw_content is None:
        raw_content = str(response)
    if isinstance(raw_content, list):
        parts: list[str] = []
        for chunk in raw_content:
            if isinstance(chunk, dict):
                parts.append(str(chunk.get("text") or chunk.get("content", "")))
            else:
                parts.append(str(chunk))
        return " ".join(parts)
    return str(raw_content)


def extract_token_usage(response: Any) -> tuple[int, int]:
    """Extract (input_tokens, output_tokens) from a LangChain response.

    Attempts ``response_metadata`` and ``usage_metadata`` attributes.
    Returns (0, 0) when no usage information is present.
    """
    # LangChain ≥ 0.2: response_metadata may carry provider-specific usage
    meta = getattr(response, "response_metadata", None) or {}
    usage = meta.get("usage") or meta.get("token_usage") or {}
    if usage:
        try:
            inp = int(usage.get("input_tokens") or usage.get("prompt_tokens") or 0)
            out = int(usage.get("output_tokens") or usage.get("completion_tokens") or 0)
            return inp, out
        except (TypeError, ValueError):
            pass

    # LangChain ≥ 0.3: usage_metadata attribute
    usage_meta = getattr(response, "usage_metadata", None) or {}
    if usage_meta:
        try:
            inp = int(usage_meta.get("input_tokens") or 0)
            out = int(usage_meta.get("output_tokens") or 0)
            return inp, out
        except (TypeError, ValueError):
            pass

    return 0, 0


def estimate_cost(input_tokens: int, output_tokens: int) -> float:
    """Estimate LLM cost in USD from token counts."""
    raw = input_tokens * _COST_PER_1K_INPUT + output_tokens * _COST_PER_1K_OUTPUT
    return raw / 1000.0


def history_to_langchain_messages(
    messages: list[Any],
) -> list[Any]:
    """Convert session message history to LangChain message objects.

    Maps:
    - ``MessageRole.SYSTEM``    → ``SystemMessage``
    - ``MessageRole.USER``      → ``HumanMessage``
    - ``MessageRole.ASSISTANT`` → ``AIMessage``
    - ``MessageRole.TOOL``      → ``ToolMessage``

    Returns an ordered list of LangChain message objects, empty if
    no history is available.
    """
    lc_messages: list[Any] = []
    for msg in messages:
        role = msg.role
        content = msg.content
        if role == MessageRole.SYSTEM:
            lc_messages.append(SystemMessage(content=content))
        elif role == MessageRole.USER:
            lc_messages.append(HumanMessage(content=content))
        elif role == MessageRole.ASSISTANT:
            lc_messages.append(AIMessage(content=content))
        elif role == MessageRole.TOOL:
            tool_id = msg.tool_call_id or "unknown"
            lc_messages.append(ToolMessage(content=content, tool_call_id=tool_id))
        else:
            # Unknown role — default to human message to preserve context
            lc_messages.append(HumanMessage(content=content))
    return lc_messages


# ---------------------------------------------------------------------------
# LangChainSessionCaller
# ---------------------------------------------------------------------------


class LangChainSessionCaller:
    """LLMCaller implementation for session conversations.

    Wraps a LangChain chat model with full session history context.
    Implements the :class:`~cleveragents.tool.actor_runtime.LLMCaller`
    protocol so it can be used with :class:`ToolCallingRuntime`.

    The caller is stateful within a single ``tell()`` invocation: it
    accumulates the LangChain message thread across tool-call loop
    iterations.  Create a fresh instance for each ``tell()`` call.

    Args:
        llm: A LangChain ``BaseLanguageModel`` (chat model).
        history: Pre-converted list of LangChain message objects
            representing the session history *before* the current
            user message.
        session_system_prompt: Optional system prompt prepended to
            the message list when no system message exists in history.
    """

    def __init__(
        self,
        llm: Any,
        history: list[Any],
        session_system_prompt: str | None = (
            "You are the CleverAgents orchestrator. You interpret the user's "
            "natural-language requests and help them manage their software "
            "development workflows. Be concise, accurate, and helpful."
        ),
    ) -> None:
        self._llm = llm
        self._last_response: Any = None
        # Build initial accumulated messages: optional system + history
        self._accumulated: list[Any] = []

        # Prepend system prompt only when history doesn't already have one
        has_system = any(isinstance(m, SystemMessage) for m in history)
        if session_system_prompt and not has_system:
            self._accumulated.append(SystemMessage(content=session_system_prompt))
        self._accumulated.extend(history)
        self._first_call = True
        # Accumulated real token usage across all invoke() calls (M1).
        self._total_input_tokens = 0
        self._total_output_tokens = 0

    def invoke(
        self,
        prompt: str,
        tool_schemas: list[dict[str, Any]],
        tool_results: list[dict[str, Any]] | None = None,
        actor_config: dict[str, Any] | None = None,
    ) -> Any:
        """Invoke the LLM with the accumulated conversation context.

        On the first call (``tool_results is None``):
            Appends the user message and invokes the LLM.

        On subsequent calls (``tool_results`` present):
            Appends the previous AI response (if available) and tool
            results, then invokes the LLM again.

        Args:
            prompt: The original user prompt (used only on the first call).
            tool_schemas: Provider tool schemas (unused — no tool binding).
            tool_results: Tool execution results from the previous iteration.
            actor_config: Actor configuration dict (currently unused).

        Returns:
            :class:`~cleveragents.tool.actor_runtime.LLMResponse` with the
            LLM's text response.
        """
        if self._first_call:
            # First iteration: append the user's prompt and invoke
            self._accumulated.append(HumanMessage(content=prompt))
            self._first_call = False
        elif tool_results is not None:
            # Subsequent iteration: add previous AI response + tool results
            if self._last_response is not None:
                self._accumulated.append(self._last_response)
            for tr in tool_results:
                tool_result_content = (
                    str(tr.get("output", {}))
                    if tr.get("success")
                    else tr.get("error", "error")
                )
                self._accumulated.append(
                    ToolMessage(
                        content=tool_result_content,
                        tool_call_id=(
                            tr.get("call_id") or tr.get("tool_name", "unknown")
                        ),
                    )
                )

        # Invoke the LLM (no tool binding — session tell is conversational)
        response = self._llm.invoke(self._accumulated)
        self._last_response = response

        # Extract real token usage from the LLM response metadata (M1).
        inp, out = extract_token_usage(response)
        self._total_input_tokens += inp
        self._total_output_tokens += out

        content = extract_content(response)

        # Extract any tool calls from the response
        tool_calls: list[LLMToolCall] = []
        raw_tool_calls = getattr(response, "tool_calls", None) or []
        for tc in raw_tool_calls:
            if isinstance(tc, dict):
                name = tc.get("name") or tc.get("id", "")
                args = tc.get("args") or tc.get("arguments") or {}
                call_id = tc.get("id", "")
                if name:
                    tool_calls.append(
                        LLMToolCall(name=name, arguments=args, call_id=call_id)
                    )

        return LLMResponse(content=content, tool_calls=tool_calls)

    def get_token_usage(self) -> tuple[int, int]:
        """Return accumulated ``(input_tokens, output_tokens)`` across all calls.

        Tokens are extracted from real LLM response metadata via
        :func:`extract_token_usage`.  Returns ``(0, 0)`` when the LLM
        provides no usage metadata (stub / minimal mode).
        """
        return self._total_input_tokens, self._total_output_tokens


# ---------------------------------------------------------------------------
# Minimal stub LLM response helpers (used by SessionWorkflow._make_stub_llm)
# ---------------------------------------------------------------------------


def estimate_tokens(prompt: str, response: str) -> tuple[int, int]:
    """Estimate token counts using a simple word-based heuristic.

    Returns ``(input_tokens, output_tokens)`` based on the
    conventional approximation of ~4 characters per token.
    """
    input_tokens = max(1, len(prompt) // 4)
    output_tokens = max(1, len(response) // 4)
    return input_tokens, output_tokens


class _StubResp:
    """Minimal LLM response object for stub mode."""

    content = "(no LLM configured)"

    def __init__(self) -> None:
        self.tool_calls: list[Any] = []
        self.response_metadata: dict[str, Any] = {}


class _StubChunk:
    """Minimal streaming chunk for stub mode."""

    content = "(no LLM configured)"


class _MinimalStubLLM:
    """Minimal LLM stub for environments without langchain_community."""

    def invoke(self, messages: Any, **kwargs: Any) -> _StubResp:  # pragma: no cover
        """Return a stub response."""
        return _StubResp()

    def stream(
        self, messages: Any, **kwargs: Any
    ) -> Iterator[_StubChunk]:  # pragma: no cover
        """Return a single stub chunk."""
        yield _StubChunk()


__all__ = [
    "LangChainSessionCaller",
    "estimate_cost",
    "extract_content",
    "extract_token_usage",
    "history_to_langchain_messages",
]
