"""Canonical LLM invocation path (capability routing + provider selection).

Callers may specify a capability name (``"fast"``, ``"reasoning"``, etc.)
or an explicit model name (``"openai/gpt-4o-mini"``).

This module resolves the capability, selects the provider, and returns
the normalized ``LLMGenerateResult`` from ``core.types.llm``.

ALL PYDANTIC — LLMMessage, LLMToolSpec, ToolChoice models.
No dicts in the public API.
"""

from __future__ import annotations

import time
from typing import TYPE_CHECKING, Literal
from uuid import uuid4

if TYPE_CHECKING:
    from core.common import JsonObject
    from core.pathways.context import Context
from core.llm.providers import (
    LLMProviderError,
    ModelCapability,
    get_model_for_capability,
    get_provider,
)
from core.types.llm import (
    LLMGenerateResult,
    LLMMessage,
    LLMStreamChunk,
    LLMToolSpec,
    ToolChoice,
)

_EVENT_BUFFER_MAX = 2000


def resolve_model(model: str) -> str:
    """Resolve a capability name to an explicit model name when applicable."""
    capability_names = {c.value for c in ModelCapability}
    m = (model or "").strip()
    if m in capability_names:
        return get_model_for_capability(m)
    return m


def _publish_llm_event(
    *,
    context: Context | None,
    event_type: str,
    call_id: str,
    ok: bool | None = None,
    duration_ms: float | None = None,
    error: str | None = None,
    data: JsonObject | None = None,
) -> None:
    """Best-effort publish of LLM lifecycle events into the event buffer."""
    if context is None:
        return

    try:
        from core.pathways.models.execution import PathwayEvent

        ev = PathwayEvent(
            event_type=event_type,
            pathway_id=context.pathway_id or None,
            node_id=context.run._current_node_id or None,
            error=error,
            data={
                "execution_id": context.execution_id or "",
                "call_id": call_id,
                "ok": ok,
                "duration_ms": duration_ms,
                **(data or {}),
            },
        )
        payload = ev.model_dump()
        ts = payload.get("timestamp")
        if ts is not None:
            payload["timestamp"] = ts.isoformat() if hasattr(ts, "isoformat") else ts

        # Buffer events on RunState for downstream consumers.
        context.run.event_buffer.append(payload)
        if len(context.run.event_buffer) > _EVENT_BUFFER_MAX:
            del context.run.event_buffer[: len(context.run.event_buffer) - _EVENT_BUFFER_MAX]
    except Exception:
        # Never break LLM calls due to observability failures.
        return


async def generate(
    *,
    prompt: str = "",
    messages: list[LLMMessage] | None = None,
    model: str = "reasoning",
    system: str | None = None,
    temperature: float = 0.7,
    max_tokens: int = 4096,
    tools: list[LLMToolSpec] | None = None,
    tool_choice: ToolChoice | None = None,
    images: list[str] | None = None,
    response_format: Literal["text", "json"] = "text",
    context: Context | None = None,
) -> LLMGenerateResult:
    """Generate via the routed provider and return a normalized result.

    All parameters and the return type are Pydantic models.
    Callers pass LLMMessage/LLMToolSpec models, get LLMGenerateResult back.
    """
    if not prompt and not messages:
        raise ValueError("prompt or messages required")

    resolved_model = resolve_model(model)
    provider = get_provider(resolved_model)

    call_id = f"llm_{uuid4().hex[:10]}"
    _publish_llm_event(
        context=context,
        event_type="llm_called",
        call_id=call_id,
        ok=None,
        duration_ms=None,
        data={
            "requested_model": model,
            "resolved_model": resolved_model,
            "provider": getattr(provider, "__class__", type(provider)).__name__,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "response_format": response_format,
            "has_prompt": bool(prompt),
            "prompt_chars": len(prompt or ""),
            "message_count": len(messages or []),
            "tool_schema_count": len(tools or []),
            "has_images": bool(images),
            "image_count": len(images or []),
        },
    )
    started_at = time.time()

    # Convert Pydantic models to dicts for the provider wire format.
    messages_wire = [m.model_dump(exclude_none=True) for m in messages] if messages else None
    tools_wire = [t.model_dump(exclude_none=True) for t in tools] if tools else None
    tool_choice_wire: str | dict[str, object] | None = None
    if isinstance(tool_choice, str):
        tool_choice_wire = tool_choice
    elif tool_choice is not None:
        tool_choice_wire = tool_choice.model_dump(exclude_none=True)

    try:
        from core.llm.retry import with_llm_retry

        _generate_with_retry = with_llm_retry(provider.generate)
        result = await _generate_with_retry(
            prompt=prompt,
            messages=messages_wire,
            model=resolved_model,
            system=system,
            temperature=temperature,
            max_tokens=max_tokens,
            tools=tools_wire,
            tool_choice=tool_choice_wire,
            images=images,
            response_format=response_format,
        )

        # Convert provider result to our canonical Pydantic model
        if isinstance(result, LLMGenerateResult):
            typed_result = result
        else:
            typed_result = LLMGenerateResult.model_validate(
                result
                if isinstance(result, dict)
                else result.__dict__
                if hasattr(result, "__dict__")
                else {
                    "content": str(result),
                    "model": resolved_model,
                    "tokens_used": 0,
                    "cost_usd": 0.0,
                },
            )

        dur = (time.time() - started_at) * 1000
        _publish_llm_event(
            context=context,
            event_type="llm_completed",
            call_id=call_id,
            ok=True,
            duration_ms=dur,
            data={
                "model": typed_result.model or resolved_model,
                "provider": typed_result.provider,
                "finish_reason": typed_result.finish_reason,
                "tokens_used": typed_result.tokens_used,
                "input_tokens": typed_result.input_tokens,
                "output_tokens": typed_result.output_tokens,
                "has_tool_calls": bool(typed_result.tool_calls),
                "tool_call_count": len(typed_result.tool_calls or []),
            },
        )
    except LLMProviderError:
        dur = (time.time() - started_at) * 1000
        _publish_llm_event(
            context=context,
            event_type="llm_failed",
            call_id=call_id,
            ok=False,
            duration_ms=dur,
            error="LLMProviderError",
            data={"requested_model": model, "resolved_model": resolved_model},
        )
        raise
    except Exception:
        dur = (time.time() - started_at) * 1000
        _publish_llm_event(
            context=context,
            event_type="llm_failed",
            call_id=call_id,
            ok=False,
            duration_ms=dur,
            error="unexpected",
            data={"requested_model": model, "resolved_model": resolved_model},
        )
        raise
    else:
        return typed_result


async def generate_stream(
    *,
    prompt: str = "",
    messages: list[LLMMessage] | None = None,
    model: str = "reasoning",
    system: str | None = None,
    temperature: float = 0.7,
    max_tokens: int = 4096,
    tools: list[LLMToolSpec] | None = None,
    tool_choice: ToolChoice | None = None,
    images: list[str] | None = None,
    response_format: Literal["text", "json"] = "text",
    context: Context | None = None,
):
    """Stream generation via the routed provider, yielding ``LLMStreamChunk``.

    Mirrors ``generate()`` but yields incremental chunks instead of a
    single result.  A final ``llm_completed`` event is still published to
    the event buffer after the stream finishes.

    Yields
    ------
    LLMStreamChunk
    """
    if not prompt and not messages:
        raise ValueError("prompt or messages required")

    resolved_model = resolve_model(model)
    provider = get_provider(resolved_model)

    call_id = f"llm_{uuid4().hex[:10]}"
    _publish_llm_event(
        context=context,
        event_type="llm_called",
        call_id=call_id,
        data={
            "requested_model": model,
            "resolved_model": resolved_model,
            "provider": getattr(provider, "__class__", type(provider)).__name__,
            "streaming": True,
        },
    )
    started_at = time.time()

    # Convert Pydantic models to dicts for the provider wire format.
    messages_wire = [m.model_dump(exclude_none=True) for m in messages] if messages else None
    tools_wire = [t.model_dump(exclude_none=True) for t in tools] if tools else None
    tool_choice_wire: str | dict[str, object] | None = None
    if isinstance(tool_choice, str):
        tool_choice_wire = tool_choice
    elif tool_choice is not None:
        tool_choice_wire = tool_choice.model_dump(exclude_none=True)

    total_content = ""
    chunk_count = 0
    last_finish_reason: str | None = None

    try:
        async for raw_chunk in provider.generate_stream(
            prompt=prompt,
            messages=messages_wire,
            model=resolved_model,
            system=system,
            temperature=temperature,
            max_tokens=max_tokens,
            tools=tools_wire,
            tool_choice=tool_choice_wire,
            images=images,
            response_format=response_format,
        ):
            chunk = (
                raw_chunk
                if isinstance(raw_chunk, LLMStreamChunk)
                else LLMStreamChunk.model_validate(raw_chunk)
            )
            total_content += chunk.delta
            chunk_count += 1
            if chunk.finish_reason:
                last_finish_reason = chunk.finish_reason
            yield chunk

        dur = (time.time() - started_at) * 1000
        _publish_llm_event(
            context=context,
            event_type="llm_completed",
            call_id=call_id,
            ok=True,
            duration_ms=dur,
            data={
                "model": resolved_model,
                "streaming": True,
                "chunks": chunk_count,
                "finish_reason": last_finish_reason,
                "content_length": len(total_content),
            },
        )
    except LLMProviderError:
        dur = (time.time() - started_at) * 1000
        _publish_llm_event(
            context=context,
            event_type="llm_failed",
            call_id=call_id,
            ok=False,
            duration_ms=dur,
            error="LLMProviderError",
            data={"requested_model": model, "resolved_model": resolved_model, "streaming": True},
        )
        raise
    except Exception:
        dur = (time.time() - started_at) * 1000
        _publish_llm_event(
            context=context,
            event_type="llm_failed",
            call_id=call_id,
            ok=False,
            duration_ms=dur,
            error="unexpected",
            data={"requested_model": model, "resolved_model": resolved_model, "streaming": True},
        )
        raise


__all__ = ["LLMProviderError", "generate", "generate_stream", "resolve_model"]
