"""Stateless pipeline functions for prompt execution.

Pipeline stages:
    prepare()       →  render + parse + expand  →  list[Message]
    execute()       →  single LLM call          →  raw response
    process()       →  response extraction      →  clean result
    run()           →  load + prepare + execute + process
    run_agent()     →  load + prepare + (execute → tool loop) + process
"""

from __future__ import annotations

import inspect
import json
from collections.abc import Callable
from typing import Any

from agentschema import PromptAgent

from ..renderers._common import _thread_nonces_local
from ..tracing.tracer import trace
from .discovery import (
    InvokerError,
    get_executor,
    get_parser,
    get_processor,
    get_renderer,
)
from .types import RICH_KINDS, ContentPart, Message, ThreadMarker

__all__ = [
    "validate_inputs",
    "prepare",
    "prepare_async",
    "execute",
    "execute_async",
    "process",
    "process_async",
    "run",
    "run_async",
    "run_agent",
    "run_agent_async",
    "headless",
    # Helpers (used by tests)
    "_get_rich_input_names",
    "_inject_thread_markers",
    "_expand_thread_markers",
    "_dict_to_message",
    "_dict_content_to_part",
]


# ---------------------------------------------------------------------------
# Input validation
# ---------------------------------------------------------------------------


def validate_inputs(
    agent: PromptAgent,
    inputs: dict[str, Any],
) -> dict[str, Any]:
    """Validate and fill defaults for inputs against ``agent.inputSchema``.

    Parameters
    ----------
    agent:
        The loaded PromptAgent.
    inputs:
        User-provided input values.

    Returns
    -------
    dict
        Validated inputs with defaults applied.

    Raises
    ------
    ValueError
        If required inputs are missing or unknown inputs are provided
        in strict mode.
    """
    if agent.inputSchema is None:
        return dict(inputs)

    schema = agent.inputSchema
    props = {p.name: p for p in schema.properties}
    result = dict(inputs)

    # Apply defaults for missing inputs
    for name, prop in props.items():
        if name not in result:
            if prop.default is not None:
                result[name] = prop.default
            elif prop.example is not None:
                result[name] = prop.example
            elif prop.required:
                raise ValueError(f"Required input '{name}' not provided and has no default value.")

    # Strict mode: reject unknown keys
    if schema.strict:
        unknown = set(result.keys()) - set(props.keys())
        if unknown:
            raise ValueError(
                f"Unknown input(s) in strict mode: {', '.join(sorted(unknown))}. "
                f"Declared inputs: {', '.join(sorted(props.keys()))}"
            )

    return result


# ---------------------------------------------------------------------------
# Thread-marker helpers
# ---------------------------------------------------------------------------


def _get_rich_input_names(agent: PromptAgent) -> dict[str, str]:
    """Return {property_name: kind} for all rich-kind inputs."""
    if agent.inputSchema is None:
        return {}
    return {p.name: p.kind for p in agent.inputSchema.properties if p.kind in RICH_KINDS}


def _inject_thread_markers(
    messages: list[Message],
    thread_nonces: dict[str, str],
) -> list[Message | ThreadMarker]:
    """Scan parsed messages for nonce markers emitted by the renderer.

    When a thread-kind input appears as ``{{thread_var}}`` in the template,
    the renderer substitutes a nonce marker string. After parsing, that
    marker ends up inside a ``Message``'s ``TextPart`` content. This
    function finds those markers, splits the message, and inserts a
    ``ThreadMarker`` at the correct position.

    Parameters
    ----------
    messages:
        Parsed message list (no ThreadMarkers yet).
    thread_nonces:
        Mapping from nonce marker strings to input property names,
        provided by the renderer.

    Returns
    -------
    list[Message | ThreadMarker]
        Messages with ``ThreadMarker`` objects injected at nonce positions.
    """
    from .types import TextPart

    result: list[Message | ThreadMarker] = []

    for msg in messages:
        text = msg.text
        # Check if any nonce marker appears in this message's text
        found_marker = None
        for marker, name in thread_nonces.items():
            if marker in text:
                found_marker = (marker, name)
                break

        if found_marker is None:
            result.append(msg)
            continue

        marker, name = found_marker

        # Split the message text at the nonce marker
        before, _, after = text.partition(marker)
        before = before.strip()
        after = after.strip()

        if before:
            result.append(
                Message(
                    role=msg.role,
                    parts=[TextPart(value=before)],
                    metadata=dict(msg.metadata),
                )
            )

        result.append(ThreadMarker(name=name))

        if after:
            result.append(
                Message(
                    role=msg.role,
                    parts=[TextPart(value=after)],
                    metadata=dict(msg.metadata),
                )
            )

    return result


def _expand_thread_markers(
    messages: list[Message | ThreadMarker],
    inputs: dict[str, Any],
    rich_inputs: dict[str, str],
) -> list[Message]:
    """Replace ThreadMarker entries with actual messages from inputs.

    If no markers exist but thread-kind inputs are provided, append
    thread messages at the end.
    """
    expanded: list[Message] = []
    marker_found = False

    for item in messages:
        if isinstance(item, ThreadMarker):
            marker_found = True
            thread_data = inputs.get(item.name, [])
            if isinstance(thread_data, list):
                for msg in thread_data:
                    if isinstance(msg, Message):
                        expanded.append(msg)
                    elif isinstance(msg, dict):
                        expanded.append(_dict_to_message(msg))
            # else: skip non-list thread values
        else:
            expanded.append(item)

    # If no markers but thread inputs exist, append at end
    if not marker_found:
        for name, kind in rich_inputs.items():
            if kind == "thread" and name in inputs:
                thread_data = inputs[name]
                if isinstance(thread_data, list):
                    for msg in thread_data:
                        if isinstance(msg, Message):
                            expanded.append(msg)
                        elif isinstance(msg, dict):
                            expanded.append(_dict_to_message(msg))

    return expanded


def _dict_to_message(d: dict[str, Any]) -> Message:
    """Convert a plain dict to a Message, preserving structure."""
    from .types import TextPart

    role = d.get("role", "user")
    content = d.get("content", "")

    parts: list[ContentPart] = []
    if isinstance(content, str):
        parts = [TextPart(value=content)]
    elif isinstance(content, list):
        # Already structured content parts — preserve as-is
        for item in content:
            if isinstance(item, ContentPart):
                parts.append(item)
            elif isinstance(item, dict):
                parts.append(_dict_content_to_part(item))
            else:
                parts.append(TextPart(value=str(item)))
    else:
        parts = [TextPart(value=str(content))]

    # Collect metadata (everything except role and content)
    metadata = {k: v for k, v in d.items() if k not in ("role", "content", "parts")}

    return Message(role=role, parts=parts, metadata=metadata)


def _dict_content_to_part(d: dict[str, Any]) -> ContentPart:
    """Convert a dict content item to a ContentPart."""
    from .types import AudioPart, FilePart, ImagePart, TextPart

    kind = d.get("kind", d.get("type", "text"))
    if kind == "text":
        return TextPart(value=d.get("value", d.get("text", "")))
    elif kind == "image" or kind == "image_url":
        source = d.get("source", "")
        if not source and "image_url" in d:
            url_data = d["image_url"]
            source = url_data.get("url", "") if isinstance(url_data, dict) else str(url_data)
        return ImagePart(
            source=source,
            detail=d.get("detail"),
            media_type=d.get("mediaType", d.get("media_type")),
        )
    elif kind == "file":
        return FilePart(
            source=d.get("source", ""),
            media_type=d.get("mediaType", d.get("media_type")),
        )
    elif kind == "audio":
        return AudioPart(
            source=d.get("source", ""),
            media_type=d.get("mediaType", d.get("media_type")),
        )
    else:
        return TextPart(value=str(d))


# ---------------------------------------------------------------------------
# Pipeline: prepare() — shared config resolution
# ---------------------------------------------------------------------------


def _resolve_prepare_config(
    agent: PromptAgent,
) -> tuple[str, str, bool]:
    """Extract format kind, parser kind, and strict flag from an agent's template config."""
    format_kind = "jinja2"
    parser_kind = "prompty"
    is_strict = True

    if agent.template is not None:
        if agent.template.format is not None and agent.template.format.kind not in (
            "",
            "*",
        ):
            format_kind = agent.template.format.kind
        if agent.template.parser is not None and agent.template.parser.kind not in (
            "",
            "*",
        ):
            parser_kind = agent.template.parser.kind
        if agent.template.format is not None and agent.template.format.strict is not None:
            is_strict = agent.template.format.strict

    return format_kind, parser_kind, is_strict


def _finalize_messages(
    messages: list[Message],
    thread_nonces: dict[str, str],
    inputs: dict[str, Any],
    rich_inputs: dict[str, str],
) -> list[Message]:
    """Inject thread markers and expand them with actual conversation messages."""
    expanded: list[Message | ThreadMarker] = list(messages)
    if thread_nonces:
        expanded = _inject_thread_markers(messages, thread_nonces)
    return _expand_thread_markers(expanded, inputs, rich_inputs)


# ---------------------------------------------------------------------------
# Pipeline: prepare()
# ---------------------------------------------------------------------------


@trace
def prepare(
    agent: PromptAgent,
    inputs: dict[str, Any] | None = None,
) -> list[Message]:
    """Render, parse, and expand a prompt into a message array.

    Pipeline:
        1. Validate inputs against ``agent.inputSchema``
        2. Discover parser; if ``Format.strict``, call ``pre_render()``
        3. Discover renderer; call ``render()``
        4. Call ``parser.parse()``
        5. Expand thread markers with structured messages from inputs

    Parameters
    ----------
    agent:
        A loaded ``PromptAgent``.
    inputs:
        Input values for template rendering.

    Returns
    -------
    list[Message]
        Model-agnostic message array ready for execution.
    """
    inputs = validate_inputs(agent, inputs or {})
    rich_inputs = _get_rich_input_names(agent)
    format_kind, parser_kind, is_strict = _resolve_prepare_config(agent)
    template = agent.instructions or ""

    parser = get_parser(parser_kind)

    parse_context: dict[str, Any] = {}
    if is_strict and hasattr(parser, "pre_render"):
        template, parse_context = parser.pre_render(template)  # type: ignore[union-attr]

    renderer = get_renderer(format_kind)
    rendered = renderer.render(agent, template, inputs)

    thread_nonces: dict[str, str] = getattr(_thread_nonces_local, "nonces", {})
    _thread_nonces_local.nonces = {}

    messages = parser.parse(agent, rendered, **parse_context)
    return _finalize_messages(messages, thread_nonces, inputs, rich_inputs)


@trace
async def prepare_async(
    agent: PromptAgent,
    inputs: dict[str, Any] | None = None,
) -> list[Message]:
    """Async variant of :func:`prepare`."""
    inputs = validate_inputs(agent, inputs or {})
    rich_inputs = _get_rich_input_names(agent)
    format_kind, parser_kind, is_strict = _resolve_prepare_config(agent)
    template = agent.instructions or ""

    parser = get_parser(parser_kind)

    parse_context: dict[str, Any] = {}
    if is_strict and hasattr(parser, "pre_render"):
        template, parse_context = parser.pre_render(template)  # type: ignore[union-attr]

    renderer = get_renderer(format_kind)
    rendered = await renderer.render_async(agent, template, inputs)

    thread_nonces: dict[str, str] = getattr(_thread_nonces_local, "nonces", {})
    _thread_nonces_local.nonces = {}

    messages = await parser.parse_async(agent, rendered, **parse_context)
    return _finalize_messages(messages, thread_nonces, inputs, rich_inputs)


# ---------------------------------------------------------------------------
# Pipeline: execute()
# ---------------------------------------------------------------------------


@trace
def execute(
    agent: PromptAgent,
    messages: list[Message],
) -> Any:
    """Execute a single LLM call with the given messages.

    The executor is discovered via entry point using
    ``agent.model.provider``. It maps abstract messages to the
    provider's wire format internally.

    Parameters
    ----------
    agent:
        A loaded ``PromptAgent`` with model configuration.
    messages:
        Abstract message array from :func:`prepare`.

    Returns
    -------
    Any
        Raw LLM response (provider-specific).
    """
    provider = agent.model.provider or ""
    if not provider:
        raise InvokerError("prompty.executors", "(no provider set)")
    executor = get_executor(provider)
    return executor.execute(agent, messages)


@trace
async def execute_async(
    agent: PromptAgent,
    messages: list[Message],
) -> Any:
    """Async variant of :func:`execute`."""
    provider = agent.model.provider or ""
    if not provider:
        raise InvokerError("prompty.executors", "(no provider set)")
    executor = get_executor(provider)
    return await executor.execute_async(agent, messages)


# ---------------------------------------------------------------------------
# Pipeline: process()
# ---------------------------------------------------------------------------


@trace
def process(
    agent: PromptAgent,
    response: Any,
) -> Any:
    """Extract a clean result from a raw LLM response.

    Parameters
    ----------
    agent:
        The ``PromptAgent`` used for the call.
    response:
        Raw response from :func:`execute`.

    Returns
    -------
    Any
        Clean result (``str``, tool calls, parsed JSON, etc.).
    """
    provider = agent.model.provider or ""
    if not provider:
        raise InvokerError("prompty.processors", "(no provider set)")
    processor = get_processor(provider)
    return processor.process(agent, response)


@trace
async def process_async(
    agent: PromptAgent,
    response: Any,
) -> Any:
    """Async variant of :func:`process`."""
    provider = agent.model.provider or ""
    if not provider:
        raise InvokerError("prompty.processors", "(no provider set)")
    processor = get_processor(provider)
    return await processor.process_async(agent, response)


# ---------------------------------------------------------------------------
# Pipeline: run()
# ---------------------------------------------------------------------------


@trace
def run(
    prompt: str | PromptAgent,
    inputs: dict[str, Any] | None = None,
    *,
    raw: bool = False,
) -> Any:
    """Full pipeline: load → prepare → execute → process.

    Parameters
    ----------
    prompt:
        Path to a ``.prompty`` file, or a pre-loaded ``PromptAgent``.
    inputs:
        Input values for template rendering.
    raw:
        If ``True``, skip processing and return the raw LLM response.

    Returns
    -------
    Any
        The processed result (or raw response if ``raw=True``).
    """
    from .loader import load

    agent = load(prompt) if isinstance(prompt, str) else prompt
    messages = prepare(agent, inputs)
    response = execute(agent, messages)
    if raw:
        return response
    return process(agent, response)


@trace
async def run_async(
    prompt: str | PromptAgent,
    inputs: dict[str, Any] | None = None,
    *,
    raw: bool = False,
) -> Any:
    """Async variant of :func:`run`."""
    from .loader import load_async

    if isinstance(prompt, str):
        agent = await load_async(prompt)
    else:
        agent = prompt
    messages = await prepare_async(agent, inputs)
    response = await execute_async(agent, messages)
    if raw:
        return response
    return await process_async(agent, response)


# ---------------------------------------------------------------------------
# Pipeline: run_agent() — agent loop with tool execution
# ---------------------------------------------------------------------------

_DEFAULT_MAX_ITERATIONS = 10


def _execute_tool(
    fn: Callable[..., Any],
    fn_name: str,
    arguments_json: str,
) -> str:
    """Execute a tool function with JSON arguments, handling errors gracefully.

    Returns the tool result as a string. On JSON parse error or tool
    function exception, returns an error message so the model can self-correct.
    """
    try:
        fn_args = json.loads(arguments_json)
    except json.JSONDecodeError as e:
        return f"Error: invalid JSON arguments for '{fn_name}': {e}"

    try:
        result = fn(**fn_args)
    except Exception as e:
        return f"Error calling '{fn_name}': {type(e).__name__}: {e}"

    return str(result)


async def _execute_tool_async(
    fn: Callable[..., Any],
    fn_name: str,
    arguments_json: str,
) -> str:
    """Async variant of :func:`_execute_tool`."""
    try:
        fn_args = json.loads(arguments_json)
    except json.JSONDecodeError as e:
        return f"Error: invalid JSON arguments for '{fn_name}': {e}"

    try:
        if inspect.iscoroutinefunction(fn):
            result = await fn(**fn_args)
        else:
            result = fn(**fn_args)
    except Exception as e:
        return f"Error calling '{fn_name}': {type(e).__name__}: {e}"

    return str(result)


def _has_tool_calls(response: Any) -> bool:
    """Check if an LLM response contains tool calls."""
    return (
        hasattr(response, "choices")
        and response.choices
        and response.choices[0].finish_reason == "tool_calls"
        and response.choices[0].message.tool_calls
    )


def _build_tool_result_messages(
    response: Any,
    tools: dict[str, Callable[..., Any]],
) -> tuple[list[Message], bool]:
    """Execute tool calls from the response and build result messages.

    Returns (messages_to_append, had_missing_tool). The messages include
    the assistant's tool-call message and each tool result.
    """
    from .types import TextPart

    tool_calls = response.choices[0].message.tool_calls
    result_messages: list[Message] = []

    # Assistant message with tool_calls metadata
    result_messages.append(
        Message(
            role="assistant",
            parts=[],
            metadata={"tool_calls": [tc.model_dump() for tc in tool_calls]},
        )
    )

    for tc in tool_calls:
        fn_name = tc.function.name
        fn = tools.get(fn_name)
        if fn is None:
            tool_result = (
                f"Error: tool function '{fn_name}' not registered. Available: {', '.join(sorted(tools)) or '(none)'}"
            )
        elif inspect.iscoroutinefunction(fn):
            tool_result = f"Error: async tool '{fn_name}' cannot be called in sync mode"
        else:
            tool_result = _execute_tool(fn, fn_name, tc.function.arguments)

        result_messages.append(
            Message(
                role="tool",
                parts=[TextPart(value=tool_result)],
                metadata={"tool_call_id": tc.id, "name": fn_name},
            )
        )

    return result_messages, False


async def _build_tool_result_messages_async(
    response: Any,
    tools: dict[str, Callable[..., Any]],
) -> list[Message]:
    """Async variant of :func:`_build_tool_result_messages`."""
    from .types import TextPart

    tool_calls = response.choices[0].message.tool_calls
    result_messages: list[Message] = []

    result_messages.append(
        Message(
            role="assistant",
            parts=[],
            metadata={"tool_calls": [tc.model_dump() for tc in tool_calls]},
        )
    )

    for tc in tool_calls:
        fn_name = tc.function.name
        fn = tools.get(fn_name)
        if fn is None:
            tool_result = (
                f"Error: tool function '{fn_name}' not registered. Available: {', '.join(sorted(tools)) or '(none)'}"
            )
        else:
            tool_result = await _execute_tool_async(fn, fn_name, tc.function.arguments)

        result_messages.append(
            Message(
                role="tool",
                parts=[TextPart(value=tool_result)],
                metadata={"tool_call_id": tc.id, "name": fn_name},
            )
        )

    return result_messages


@trace
def run_agent(
    prompt: str | PromptAgent,
    inputs: dict[str, Any] | None = None,
    *,
    tools: dict[str, Callable[..., Any]] | None = None,
    max_iterations: int = _DEFAULT_MAX_ITERATIONS,
    raw: bool = False,
) -> Any:
    """Run a prompt with automatic tool-call execution loop.

    Similar to :func:`run`, but when the LLM returns tool calls, the
    specified tool functions are executed and their results are sent back
    to the model. This repeats until the model returns a normal response
    or *max_iterations* is reached.

    Parameters
    ----------
    prompt:
        Path to a ``.prompty`` file, or a pre-loaded ``PromptAgent``.
    inputs:
        Input values for template rendering.
    tools:
        Mapping of tool name → callable. Tool names must match the
        ``name`` field in the ``.prompty`` frontmatter ``tools:`` list.
    max_iterations:
        Maximum number of tool-call loop iterations (default 10).
    raw:
        If ``True``, skip processing and return the raw final response.

    Returns
    -------
    Any
        The processed result from the final LLM response.

    Raises
    ------
    ValueError
        If *max_iterations* is exceeded.
    """
    from ..tracing.tracer import Tracer
    from .loader import load

    agent = load(prompt) if isinstance(prompt, str) else prompt
    tools = tools or {}
    messages = prepare(agent, inputs)

    with Tracer.start("AgentLoop") as t:
        t("type", "agent")
        t("tools", list(tools.keys()))

        response = execute(agent, messages)
        iteration = 0

        while _has_tool_calls(response):
            iteration += 1
            if iteration > max_iterations:
                raise ValueError(
                    f"Agent loop exceeded max_iterations ({max_iterations}). "
                    f"The model kept requesting tool calls. Increase max_iterations or check your tools."
                )

            tool_messages, _ = _build_tool_result_messages(response, tools)
            messages.extend(tool_messages)
            response = execute(agent, messages)

        t("iterations", iteration)
        t("result", response)

    if raw:
        return response
    return process(agent, response)


@trace
async def run_agent_async(
    prompt: str | PromptAgent,
    inputs: dict[str, Any] | None = None,
    *,
    tools: dict[str, Callable[..., Any]] | None = None,
    max_iterations: int = _DEFAULT_MAX_ITERATIONS,
    raw: bool = False,
) -> Any:
    """Async variant of :func:`run_agent`."""
    from ..tracing.tracer import Tracer
    from .loader import load_async

    if isinstance(prompt, str):
        agent = await load_async(prompt)
    else:
        agent = prompt
    tools = tools or {}
    messages = await prepare_async(agent, inputs)

    with Tracer.start("AgentLoopAsync") as t:
        t("type", "agent")
        t("tools", list(tools.keys()))

        response = await execute_async(agent, messages)
        iteration = 0

        while _has_tool_calls(response):
            iteration += 1
            if iteration > max_iterations:
                raise ValueError(
                    f"Agent loop exceeded max_iterations ({max_iterations}). "
                    f"The model kept requesting tool calls. Increase max_iterations or check your tools."
                )

            tool_messages = await _build_tool_result_messages_async(response, tools)
            messages.extend(tool_messages)
            response = await execute_async(agent, messages)

        t("iterations", iteration)
        t("result", response)

    if raw:
        return response
    return await process_async(agent, response)


# ---------------------------------------------------------------------------


def headless(
    api: str = "chat",
    content: str | list | dict = "",
    *,
    model: str = "",
    provider: str = "openai",
    connection: dict[str, Any] | None = None,
    options: dict[str, Any] | None = None,
) -> PromptAgent:
    """Create a ``PromptAgent`` programmatically without a ``.prompty`` file.

    Useful for embedding calls, one-off completions, or cases where a file
    isn't needed.  The returned agent can be passed to :func:`execute` and
    :func:`process`.

    Parameters
    ----------
    api:
        The API type: ``"chat"``, ``"embedding"``, ``"image"``, ``"agent"``.
    content:
        Content to attach — for embeddings this is the input text/list,
        for images this is the prompt string. Stored in
        ``agent.metadata["content"]``.
    model:
        Model identifier (e.g. ``"gpt-4"``, ``"text-embedding-ada-002"``).
    provider:
        Provider name (``"openai"`` or ``"azure"``).
    connection:
        Connection config dict (``kind``, ``apiKey``, ``endpoint``, etc.).
    options:
        Model options dict (``temperature``, ``maxOutputTokens``, etc.).

    Returns
    -------
    PromptAgent
        A fully typed agent ready for :func:`execute` / :func:`process`.

    Examples
    --------
    >>> from prompty import headless, execute, process
    >>> agent = headless(
    ...     api="embedding",
    ...     model="text-embedding-ada-002",
    ...     provider="azure",
    ...     connection={
    ...         "kind": "key",
    ...         "endpoint": "https://my.openai.azure.com",
    ...         "apiKey": "sk-...",
    ...     },
    ...     content="hello world",
    ... )
    >>> response = execute(agent, agent.metadata["content"])
    >>> result = process(agent, response)
    """
    from agentschema import AgentDefinition

    data: dict[str, Any] = {
        "kind": "prompt",
        "name": "headless",
        "model": {
            "id": model,
            "provider": provider,
            "apiType": api,
        },
        "metadata": {
            "content": content,
        },
    }

    if connection:
        data["model"]["connection"] = connection
    if options:
        data["model"]["options"] = options

    agent = AgentDefinition.load(data)
    if not isinstance(agent, PromptAgent):
        raise TypeError(f"Expected PromptAgent, got {type(agent).__name__}")
    return agent
