"""Built-in Node Handlers -- Pre-built nodes for Graph Orchestrator v18.

Provides reusable handlers:
  - llm_node:       LLM call with prompt template
  - tool_node:      MCP tool call
  - transform_node: State transformation (map/filter/merge)
  - condition_node: Condition check (router)
  - delay_node:     Wait time
  - log_node:       Logging/debug
  - accumulate_node: Collects results from parallel branches
  - gate_node:      Checks if condition is met (gating)
  - retry_wrapper:  Wrapper for retry logic

All handlers have signature: async (GraphState) -> GraphState
Router handlers return: async (GraphState) -> str
"""

from __future__ import annotations

import asyncio
import time
from typing import TYPE_CHECKING, Any

from cognithor.utils.logging import get_logger

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from cognithor.graph.types import GraphState

log = get_logger(__name__)


# ── LLM Node ────────────────────────────────────────────────────


def llm_node(
    prompt_template: str,
    *,
    input_key: str = "messages",
    output_key: str = "response",
    model: str = "",
    llm_handler: Callable | None = None,
) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Creates an LLM call handler.

    Args:
        prompt_template: Template with {variable} placeholders
        input_key: State key for input
        output_key: State key for output
        model: LLM model (optional)
        llm_handler: Custom LLM call function
    """

    async def handler(state: GraphState) -> GraphState:
        # Render template
        prompt = prompt_template
        for key in state:
            placeholder = "{" + key + "}"
            if placeholder in prompt:
                prompt = prompt.replace(placeholder, str(state[key]))

        if llm_handler:
            response = await llm_handler(prompt, model=model)
        else:
            # Fallback: prompt as response (for tests without LLM)
            response = f"[LLM would process: {prompt[:200]}]"

        state[output_key] = response
        return state

    return handler


# ── Tool Node ────────────────────────────────────────────────────


def tool_node(
    tool_name: str,
    *,
    params_key: str = "tool_params",
    result_key: str = "tool_result",
    tool_executor: Callable | None = None,
) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Creates an MCP tool call handler.

    Args:
        tool_name: Name of the tool
        params_key: State key for tool parameters
        result_key: State key for tool result
        tool_executor: Custom tool executor
    """

    async def handler(state: GraphState) -> GraphState:
        params = state.get(params_key, {})

        if tool_executor:
            result = await tool_executor(tool_name, params)
        else:
            result = {"tool": tool_name, "params": params, "status": "simulated"}

        state[result_key] = result
        return state

    return handler


# ── Transform Node ───────────────────────────────────────────────


def transform_node(
    transform_fn: Callable[[dict[str, Any]], dict[str, Any]],
) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Creates a state transformation handler.

    Args:
        transform_fn: Sync function that transforms dict->dict
    """

    async def handler(state: GraphState) -> GraphState:
        data = state.to_dict()
        transformed = transform_fn(data)
        state.update(transformed)
        return state

    return handler


# ── Condition Node (Router) ──────────────────────────────────────


def condition_node(
    condition_fn: Callable[[dict[str, Any]], str],
) -> Callable[[GraphState], Awaitable[str]]:
    """Creates a condition router.

    Args:
        condition_fn: Function that returns dict->str (edge name)
    """

    async def handler(state: GraphState) -> str:
        data = state.to_dict()
        return condition_fn(data)

    return handler


def threshold_router(
    key: str,
    threshold: float,
    above: str = "above",
    below: str = "below",
) -> Callable[[GraphState], Awaitable[str]]:
    """Router that decides based on a threshold value."""

    async def handler(state: GraphState) -> str:
        value = state.get(key, 0)
        try:
            return above if float(value) >= threshold else below
        except (TypeError, ValueError):
            return below

    return handler


def key_router(
    key: str,
    mapping: dict[str, str] | None = None,
    *,
    default: str = "__default__",
) -> Callable[[GraphState], Awaitable[str]]:
    """Router that routes based on a state value.

    Args:
        key: State key whose value is checked
        mapping: Optional mapping {value: edge_name}
        default: Fallback when no match
    """

    async def handler(state: GraphState) -> str:
        value = str(state.get(key, ""))
        if mapping:
            return mapping.get(value, default)
        return value or default

    return handler


# ── Delay Node ───────────────────────────────────────────────────


def delay_node(seconds: float) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Waits for a specified time."""

    async def handler(state: GraphState) -> GraphState:
        await asyncio.sleep(seconds)
        return state

    return handler


# ── Log Node ─────────────────────────────────────────────────────


def log_node(
    message: str = "",
    *,
    log_keys: list[str] | None = None,
    state_key: str = "__log__",
) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Logs state information."""

    async def handler(state: GraphState) -> GraphState:
        info: dict[str, Any] = {"timestamp": time.strftime("%H:%M:%S", time.gmtime())}
        if message:
            info["message"] = message
        if log_keys:
            for k in log_keys:
                info[k] = state.get(k, "N/A")
        else:
            info["state_keys"] = list(state.keys())

        # Build log history in state
        history = state.get(state_key, [])
        if not isinstance(history, list):
            history = []
        history.append(info)
        state[state_key] = history

        log.debug("graph_log_node", **info)
        return state

    return handler


# ── Accumulate Node ──────────────────────────────────────────────


def accumulate_node(
    source_keys: list[str],
    target_key: str = "accumulated",
) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Collects values from multiple state keys into a list."""

    async def handler(state: GraphState) -> GraphState:
        accumulated = []
        for key in source_keys:
            value = state.get(key)
            if value is not None:
                accumulated.append({"key": key, "value": value})
        state[target_key] = accumulated
        return state

    return handler


# ── Gate Node ────────────────────────────────────────────────────


def gate_node(
    check_fn: Callable[[dict[str, Any]], bool],
    *,
    error_message: str = "Gate check failed",
) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Checks a condition and raises exception if not met."""

    async def handler(state: GraphState) -> GraphState:
        data = state.to_dict()
        if not check_fn(data):
            raise ValueError(error_message)
        return state

    return handler


# ── Counter Node ─────────────────────────────────────────────────


def counter_node(
    key: str = "iteration",
    increment: int = 1,
) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Increments a counter in the state."""

    async def handler(state: GraphState) -> GraphState:
        current = state.get(key, 0)
        state[key] = current + increment
        return state

    return handler


# ── Set Value Node ───────────────────────────────────────────────


def set_value_node(**values: Any) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Sets fixed values in the state."""

    async def handler(state: GraphState) -> GraphState:
        for k, v in values.items():
            state[k] = v
        return state

    return handler


# ── Merge Node ───────────────────────────────────────────────────


def merge_node(
    merge_fn: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
) -> Callable[[GraphState], Awaitable[GraphState]]:
    """Merge point for parallel branches.

    Can optionally apply a merge function.
    """

    async def handler(state: GraphState) -> GraphState:
        if merge_fn:
            merged = merge_fn(state.to_dict())
            state.update(merged)
        return state

    return handler
