"""Session workflow — orchestrator LLM invocation for ``agents session tell``.

Replaces the M3 stub with real LLM actor invocation.  The user prompt is
sent to the session's bound orchestrator actor via the A2A ``message/send`` /
``message/stream`` operations routed through :class:`A2aLocalFacade`.

Call chain in local mode::

    tell CLI  ──▸  A2aLocalFacade.dispatch("message/send")
              ──▸  SessionWorkflow.tell()
              ──▸  LangChainSessionCaller.invoke()
              ──▸  LangChain chat model (real LLM)

Forgejo: #5784
"""

from __future__ import annotations

import time
from collections.abc import Callable, Iterator
from typing import TYPE_CHECKING, Any

import structlog
from langchain_core.messages import HumanMessage, SystemMessage
from pydantic import BaseModel, ConfigDict, Field

try:
    from langchain_community.llms import FakeListLLM
except ImportError:  # pragma: no cover — langchain_community optional
    FakeListLLM = None

from cleveragents.application.services.prompt_sanitizer import PromptSanitizer
from cleveragents.application.services.session_caller import (
    LangChainSessionCaller,
    _MinimalStubLLM,
    estimate_cost,
    estimate_tokens,
    extract_content,
    history_to_langchain_messages,
)
from cleveragents.application.services.strategy_resolution import _parse_actor_name
from cleveragents.domain.models.core.session import (
    MessageRole,
    SessionActorNotConfiguredError,
    SessionMessage,
)
from cleveragents.tool.actor_runtime import (
    LLMCaller,
    ToolCallingRuntime,
    ToolCallRunResult,
)
from cleveragents.tool.registry import ToolRegistry
from cleveragents.tool.runner import ToolRunner

if TYPE_CHECKING:
    from cleveragents.domain.models.core.session import SessionService
    from cleveragents.providers.registry import ProviderRegistry

logger = structlog.get_logger(__name__)
_sanitizer = PromptSanitizer()

# System prompt for the orchestrator actor in session context.
_SESSION_SYSTEM_PROMPT = (
    "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."
)


# ---------------------------------------------------------------------------
# TellResult
# ---------------------------------------------------------------------------


class TellResult(BaseModel):
    """Result of a :meth:`SessionWorkflow.tell` invocation.

    Attributes:
        session_id: Session ULID.
        user_message: Original user prompt.
        assistant_message: Orchestrator response text.
        input_tokens: Input token count.
        output_tokens: Output token count.
        cost: Estimated cost in USD.
        duration_ms: Wall-clock LLM invocation duration (ms).
        tool_calls_count: Number of tool calls made.
    """

    session_id: str = Field(..., description="ULID of the session")
    user_message: str = Field(..., description="Original user prompt")
    assistant_message: str = Field(..., description="Orchestrator response text")
    input_tokens: int = Field(0, ge=0, description="Input token count")
    output_tokens: int = Field(0, ge=0, description="Output token count")
    cost: float = Field(0.0, ge=0.0, description="Estimated cost in USD")
    duration_ms: float = Field(0.0, ge=0.0, description="LLM invocation duration (ms)")
    tool_calls_count: int = Field(0, ge=0, description="Number of tool calls made")

    model_config = ConfigDict(
        str_strip_whitespace=True,
        validate_assignment=True,
    )

    def usage_dict(self) -> dict[str, Any]:
        """Return a JSON-serializable usage summary dict."""
        return {
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "cost_usd": round(self.cost, 6),
            "duration_ms": round(self.duration_ms, 1),
            "tool_calls": self.tool_calls_count,
        }


# ---------------------------------------------------------------------------
# SessionWorkflow
# ---------------------------------------------------------------------------


class SessionWorkflow:
    """Orchestrates ``agents session tell`` — real LLM actor invocation.

    Resolves the session's actor (or ``--actor`` override), invokes the LLM
    via :class:`ToolCallingRuntime`, persists the response, and records
    token usage via :class:`SessionService`.

    Args:
        session_service: Session persistence service.
        provider_registry: Optional provider registry for LLM resolution.
            When ``None``, falls back to a stub response.
        tool_registry: Optional ``ToolRegistry`` (defaults to empty).
        llm_factory: Optional ``(actor_name: str) -> Any`` test injection point.
            Tests inject stub LLMs through this factory.
        actor_resolver: Optional ``(actor_name: str) -> str | None`` for
            namespace/name → provider/model resolution.
        actor_options_resolver: Optional ``(actor_name: str) -> dict | None``
            for resolving actor-level options (openai_api_base, etc.).
        max_iterations: Max tool-call loop iterations (default 25).
    """

    def __init__(
        self,
        session_service: SessionService,
        provider_registry: ProviderRegistry | None = None,
        tool_registry: ToolRegistry | None = None,
        llm_factory: Callable[[str], Any] | None = None,
        actor_resolver: Callable[[str], str | None] | None = None,
        actor_options_resolver: Callable[[str], dict[str, Any] | None] | None = None,
        max_iterations: int = 25,
    ) -> None:
        self._session_service = session_service
        self._provider_registry = provider_registry
        self._tool_registry = tool_registry or ToolRegistry()
        self._llm_factory = llm_factory
        self._actor_resolver = actor_resolver
        self._actor_options_resolver = actor_options_resolver
        self._max_iterations = max_iterations
        self._logger = logger.bind(component="session_workflow")
        # Populated by tell_stream() so callers can read real usage metrics
        # after the stream is exhausted (addresses M3, M4).
        self.last_stream_usage: dict[str, Any] = {}

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def tell(
        self,
        session_id: str,
        prompt: str,
        actor_override: str | None = None,
    ) -> TellResult:
        """Send a user prompt to the session's orchestrator actor (non-streaming).

        Appends the user message, resolves the actor, invokes the LLM,
        persists the assistant response, and records token usage.

        Raises:
            SessionNotFoundError: If the session does not exist.
            SessionActorNotConfiguredError: If no actor is configured.
            ValueError: If prompt is empty.
        """
        if not prompt or not prompt.strip():
            raise ValueError("prompt must not be empty")

        session = self._session_service.get(session_id)
        actor_name = actor_override or session.actor_name
        if not actor_name:
            raise SessionActorNotConfiguredError(
                "Session has no actor configured. "
                "Create a session with --actor <actor> or pass --actor."
            )

        # Load history before appending the user message so the slice in
        # _invoke_actor reads a stable snapshot (addresses m11).
        # TODO(#5784-M9): If _invoke_actor() raises below, the session retains
        # an orphaned user message with no paired assistant response.  Implement
        # compensation logic that removes the orphaned message on failure, or
        # persist the user message transactionally with the assistant response.
        # Sanitize the prompt before passing to the LLM (addresses M7).
        sanitized_prompt = _sanitizer.sanitize_user_input(prompt).sanitized

        self._session_service.append_message(
            session_id=session_id,
            role=MessageRole.USER,
            content=prompt,
        )

        start = time.monotonic()
        run_result, real_input_tokens, real_output_tokens = self._invoke_actor(
            session_id=session_id,
            prompt=sanitized_prompt,
            actor_name=actor_name,
        )
        duration_ms = (time.monotonic() - start) * 1000.0

        assistant_content = run_result.content or ""

        # Persist assistant response
        self._session_service.append_message(
            session_id=session_id,
            role=MessageRole.ASSISTANT,
            content=assistant_content if assistant_content else "(no response)",
        )

        # Use real token counts from LLM response metadata when available,
        # falling back to the heuristic estimate otherwise (M1).
        if real_input_tokens > 0 or real_output_tokens > 0:
            input_tokens = real_input_tokens
            output_tokens = real_output_tokens
        else:
            input_tokens, output_tokens = estimate_tokens(
                sanitized_prompt, assistant_content
            )
        cost = estimate_cost(input_tokens, output_tokens)
        self._session_service.update_token_usage(
            session_id=session_id,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost=cost,
        )

        self._logger.info(
            "session.tell.complete",
            session_id=session_id,
            actor=actor_name,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            tool_calls=len(run_result.tool_call_history),
            duration_ms=round(duration_ms, 1),
        )

        return TellResult(
            session_id=session_id,
            user_message=prompt,
            assistant_message=assistant_content,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost=cost,
            duration_ms=round(duration_ms, 1),
            tool_calls_count=len(run_result.tool_call_history),
        )

    def tell_stream(
        self,
        session_id: str,
        prompt: str,
        actor_override: str | None = None,
    ) -> Iterator[str]:
        """Stream the assistant's response token-by-token.

        Loads session history before appending the user message (addresses C1).
        After the stream is exhausted, the full response is persisted, token
        usage is recorded, and ``last_stream_usage`` is populated (M3, M4).

        Yields:
            Successive text tokens from the LLM.

        Raises:
            SessionNotFoundError: If the session does not exist.
            SessionActorNotConfiguredError: If no actor is configured.
            ValueError: If prompt is empty.
        """
        if not prompt or not prompt.strip():
            raise ValueError("prompt must not be empty")

        # Reset stale usage data before any operation that could raise (M4).
        self.last_stream_usage = {}

        session = self._session_service.get(session_id)
        actor_name = actor_override or session.actor_name
        if not actor_name:
            raise SessionActorNotConfiguredError(
                "Session has no actor configured. "
                "Create a session with --actor <actor> or pass --actor."
            )

        # Load history BEFORE appending the user message (addresses C1).
        # Build the LC message list from existing history + the new prompt
        # so the LLM does not receive a duplicate HumanMessage.
        existing_history = self._session_service.get_messages(session_id)

        # Sanitize the prompt before passing to the LLM (addresses M7).
        sanitized_prompt = _sanitizer.sanitize_user_input(prompt).sanitized
        lc_messages = self._build_lc_messages_from_history(
            existing_history, sanitized_prompt
        )

        # Append user message to the session (persisted separately from
        # the LC message list used for the LLM call).
        self._session_service.append_message(
            session_id=session_id,
            role=MessageRole.USER,
            content=prompt,
        )

        llm = self._resolve_llm(actor_name)

        start = time.monotonic()
        collected_chunks: list[str] = []

        # try/finally guarantees persistence even when the consumer
        # breaks early.  The stream is explicitly closed in the finally
        # block to release resources (M8).
        stream = llm.stream(lc_messages)
        try:
            for chunk in stream:
                token = extract_content(chunk)
                if token:
                    collected_chunks.append(token)
                    yield token
        finally:
            if hasattr(stream, "close"):
                stream.close()
            duration_ms = (time.monotonic() - start) * 1000.0
            full_content = "".join(collected_chunks)

            # Persist and record usage
            self._session_service.append_message(
                session_id=session_id,
                role=MessageRole.ASSISTANT,
                content=full_content if full_content else "(no response)",
            )
            input_tokens, output_tokens = estimate_tokens(
                sanitized_prompt, full_content
            )
            cost = estimate_cost(input_tokens, output_tokens)
            self._session_service.update_token_usage(
                session_id=session_id,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cost=cost,
            )

            # Populate last_stream_usage so CLI can read real metrics (M3, M4)
            self.last_stream_usage = {
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "cost_usd": cost,
                "duration_ms": round(duration_ms, 1),
                "tool_calls": 0,
            }

            self._logger.info(
                "session.tell_stream.complete",
                session_id=session_id,
                actor=actor_name,
                duration_ms=round(duration_ms, 1),
            )

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _resolve_llm(self, actor_name: str) -> Any:
        """Create a LangChain LLM for the given actor name.

        Delegates to ``llm_factory`` (test injection), ``ProviderRegistry``,
        or falls back to a stub when neither is available.

        Pre-resolves namespace/name actor references (e.g. ``local/my-actor``)
        via ``actor_resolver`` before passing the name to the provider
        registry (addresses #11254).
        """
        # Test injection point — skips provider registry when llm_factory is set.
        # Must be checked before _provider_registry fallback so tests can inject
        # stubs even when the registry is None (the default).
        if self._llm_factory is not None:
            result = self._llm_factory(actor_name)
            if result is not None:
                return result
            # Factory returned None; fall through to provider registry

        if self._provider_registry is None:
            return self._make_stub_llm()

        # Pre-resolve namespace/name references (e.g. local/my-strategist)
        # to provider/model format via the injected actor_resolver.
        # When no resolver is available or resolution fails, fall through
        # to _parse_actor_name which treats the input as namespace/name
        # with the default provider (avoiding a ValueError crash).
        resolved: str | None = None
        if self._actor_resolver is not None:
            try:
                resolved = self._actor_resolver(actor_name)
            except Exception:
                resolved = None
        effective_name = resolved if resolved is not None else actor_name

        provider_type, model_id = _parse_actor_name(effective_name)
        actor_kwargs: dict[str, Any] = {}
        if self._actor_options_resolver is not None:
            try:
                options = self._actor_options_resolver(actor_name)
                if options:
                    from cleveragents.actor.config import build_llm_kwargs_from_options

                    actor_kwargs = build_llm_kwargs_from_options(
                        options, logger=self._logger
                    )
            except Exception:
                pass
        return self._provider_registry.create_llm(
            provider_type=provider_type,
            model_id=model_id,
            **actor_kwargs,
        )

    @staticmethod
    def _make_stub_llm() -> Any:
        """Return a deterministic stub LLM for environments without API keys."""
        if FakeListLLM is not None:
            return FakeListLLM(responses=["(no LLM configured)"] * 100)
        return _MinimalStubLLM()

    def _build_lc_messages_from_history(
        self,
        history_messages: list[SessionMessage],
        prompt: str,
    ) -> list[Any]:
        """Build LangChain messages from pre-loaded history + new prompt.

        Used by ``tell_stream()`` to avoid the duplicate-user-message bug (C1).
        The caller controls when history is sampled.
        """
        lc_history = history_to_langchain_messages(history_messages)

        messages: list[Any] = []
        if not any(isinstance(m, SystemMessage) for m in lc_history):
            messages.append(SystemMessage(content=_SESSION_SYSTEM_PROMPT))
        messages.extend(lc_history)
        messages.append(HumanMessage(content=prompt))
        return messages

    def _invoke_actor(
        self,
        session_id: str,
        prompt: str,
        actor_name: str,
    ) -> tuple[ToolCallRunResult, int, int]:
        """Run the tool-calling loop against the resolved LLM.

        Returns ``(run_result, input_tokens, output_tokens)`` with real token
        counts from LLM response metadata.  Callers fall back to
        ``estimate_tokens()`` when both counts are zero.
        """
        llm = self._resolve_llm(actor_name)

        # Load history — only SessionNotFoundError is handled;
        # infrastructure errors propagate (addresses M5).
        # TODO(#5784-m7): Session message history (assistant/tool roles) is fed to
        # the LLM without sanitization.  Apply PromptSanitizer to historical
        # messages to prevent prompt injection from prior assistant responses.
        history_msgs = self._session_service.get_messages(session_id)

        # Exclude the user message we just appended (last message)
        # so it is not duplicated when LangChainSessionCaller appends it
        lc_history = history_to_langchain_messages(history_msgs[:-1])

        caller: LLMCaller = LangChainSessionCaller(
            llm=llm,
            history=lc_history,
        )

        runner = ToolRunner(registry=self._tool_registry)
        runtime = ToolCallingRuntime(
            registry=self._tool_registry,
            runner=runner,
            llm_caller=caller,
            max_iterations=self._max_iterations,
        )

        run_result = runtime.run_tool_loop(prompt=prompt)
        # Read real token usage accumulated by LangChainSessionCaller (M1).
        real_input, real_output = caller.get_token_usage()
        return run_result, real_input, real_output


__all__ = [
    "LangChainSessionCaller",
    "SessionWorkflow",
    "TellResult",
]
