import functools
import inspect
import traceback
from collections.abc import Callable
from typing import TYPE_CHECKING, get_type_hints

import cloudpickle

from yera.agents.context import (
    AgentContext,
    no_active_agent_context,
)
from yera.agents.dataclasses import AgentInput, AgentMetadata
from yera.agents.exceptions import AgentRuntimeError, ErrorCause
from yera.agents.registry import get_agent_name, register_agent
from yera.agents.typing.coerce import (
    coerce_value,
)
from yera.agents.typing.utils import get_type_name
from yera.agents.typing.validate import assert_value_is_return_type, validate_signature
from yera.dsl.functions import sys_prompt
from yera.events import exit_event
from yera.events.exceptions import UnsupportedAwaitUserBlockError
from yera.events.models import ExitEventData
from yera.events.runtime import PyRuntimeExecutor, stream_handler
from yera.events.stream import (
    EventStream,
    failure_exit_event_already_pushed,
)
from yera.models.llm_atlas_proxy import llm

if TYPE_CHECKING:
    from yera import LLMContext


class AgentFunctionWrapper:
    def __init__(
        self,
        agent_function: Callable,
        name: str | None,
        description: str | None,
        model: "LLMContext | None",
        sys_prompt_str: str | None,
    ):
        self.model = model or llm.get_default()
        self.sys_prompt = sys_prompt_str

        self.agent_function = agent_function

        fn_name = agent_function.__name__
        module = agent_function.__module__
        identifier = f"{module}.{fn_name}"

        sig = inspect.signature(agent_function)
        type_hints = get_type_hints(agent_function)
        validate_signature(sig, type_hints, identifier)
        self.signature = sig
        self.type_hints = type_hints

        params = _get_agent_inputs(sig, type_hints)
        self.return_cls = _get_return_type(sig, type_hints)

        self.metadata = AgentMetadata(
            name=name or fn_name,
            description=description,
            module=module,
            docs=agent_function.__doc__,
            identifier=identifier,
            params=params,
            return_type=get_type_name(self.return_cls),
        )

        # Register this agent's metadata so the consumer can resolve the agent name
        # (including nested agents) from an event's agent_id.
        register_agent(self.metadata.identifier, self.metadata)

        functools.update_wrapper(self, agent_function)
        self.__signature__ = sig

    def _validate_and_coerce_args(self, args: tuple, kwargs: dict) -> dict:
        try:
            bound = self.signature.bind(*args, **kwargs)
        except TypeError as e:
            raise TypeError(
                f"Invalid arguments for {self.name} ({self.metadata.identifier}): {e}"
            ) from None

        bound.apply_defaults()

        coerced_args = {}

        for name, arg in bound.arguments.items():
            if arg is None:
                continue

            try:
                coerced_args[name] = coerce_value(arg, self.type_hints[name])
            except (TypeError, ValueError) as e:
                raise TypeError(
                    f"Argument '{name}' for agent '{self.metadata.identifier}': "
                    f"expected {get_type_name(self.type_hints[name])}, "
                    f"got {type(arg).__name__}. Error: {e}"
                ) from None

        return coerced_args

    def _handle_agent_exit(
        self,
        is_top_level: bool,
        ret_value=None,
        exc: BaseException | None = None,
    ):
        """Finish an agent run: either return the result (success) or handle the exception.

        Used from both the success path and the exception path of invoke().
        Returns the run result (including ret_value on success) or raises; never both.
        """
        if exc is None:
            if is_top_level:
                exit_event(
                    exit_code=0,
                    reason="Yera program completed successfully.",
                    return_value=ret_value,
                )
            return ret_value

        if isinstance(exc, UnsupportedAwaitUserBlockError):
            if is_top_level:
                return None
            raise exc

        if failure_exit_event_already_pushed():
            if is_top_level:
                return None
            raise exc

        error_cause = ErrorCause(error_type=exc.__class__.__name__)

        if is_top_level:
            exit_event(
                exit_code=1,
                reason=f"Error during agent execution: {exc!s}",
                return_value=traceback.format_exc(),
                error_cause=error_cause,
            )
            return None

        # First catcher (nested agent that failed): push then re-raise
        exit_event(
            exit_code=1,
            reason=f"Error during agent execution: {exc!s}",
            return_value=traceback.format_exc(),
            error_cause=error_cause,
        )
        raise exc

    def invoke(self, *args, **kwargs):
        with AgentContext(self.metadata) as ctx:
            model_entered = True

            try:
                c_args = self._validate_and_coerce_args(args, kwargs)

                if self.model:
                    self.model.__enter__()
                    model_entered = True

                    if self.sys_prompt:
                        sys_prompt(self.sys_prompt)

                ret_value = self.agent_function(**c_args)

                assert_value_is_return_type(ret_value, self.return_cls, self.metadata)

                return self._handle_agent_exit(
                    is_top_level=ctx.top_level(), ret_value=ret_value
                )
            except Exception as e:
                return self._handle_agent_exit(is_top_level=ctx.top_level(), exc=e)
            finally:
                if model_entered and self.model:
                    self.model.__exit__(None, None, None)

    @property
    def name(self) -> str:
        return self.metadata.name

    def __call__(self, *args, **kwargs):
        if len(kwargs) != 0:
            raise NotImplementedError("Agent kwargs are currently not supported.")

        EventStream.build()

        if no_active_agent_context():  # top of the agent stack, about to start
            with PyRuntimeExecutor(self, args, is_restricted_runtime=True):
                exit_evt = stream_handler(self.metadata)

            if not (
                exit_evt.block_type == "exit"
                and isinstance(exit_evt.data, ExitEventData)
            ):
                message = f"Expected run to terminate with an exit block, instead got {exit_evt.block_type!r}"
                raise AgentRuntimeError(
                    message=message,
                    exit_code=1,
                    reason=message,
                    agent_name=self.metadata.name,
                    agent_identifier=self.metadata.identifier,
                    error_cause=ErrorCause(
                        error_type="InvalidExitBlock",
                        block_type=exit_evt.block_type,
                    ),
                )

            if exit_evt.data.exit_code != 0:
                message = (
                    exit_evt.data.reason
                    if exit_evt.data.reason
                    else f"Agent exited with code {exit_evt.data.exit_code}"
                )
                exited_agent_id = exit_evt.agent_instance.agent_id
                exited_agent_name = get_agent_name(exited_agent_id)
                raise AgentRuntimeError(
                    message=message,
                    exit_code=exit_evt.data.exit_code,
                    reason=exit_evt.data.reason,
                    agent_name=exited_agent_name,
                    agent_identifier=exited_agent_id,
                    error_cause=exit_evt.data.error_cause,
                )

            return coerce_value(exit_evt.data.return_value, self.return_cls)
        return self.invoke(*args)

    def __reduce__(self):
        args = (
            self.metadata.name,
            self.metadata.description,
            self.model,
            self.sys_prompt,
        )
        return (
            _unpickle_agent_wrapper,
            (cloudpickle.dumps(self.agent_function), args),
        )


def _get_return_type(sig, type_hints):
    ret_type = type_hints.get("return", None)
    if ret_type is None and sig.return_annotation is not inspect.Signature.empty:
        ret_type = sig.return_annotation
    return ret_type


def _get_agent_inputs(sig, type_hints) -> list[AgentInput]:
    params = []
    for param_name, param in sig.parameters.items():
        if param_name in type_hints:
            param = param.replace(annotation=type_hints[param_name])
        params.append(AgentInput.from_parameter(param))
    return params


def _unpickle_agent_wrapper(pickled_func, args):
    agent_function = cloudpickle.loads(pickled_func)
    return AgentFunctionWrapper(agent_function, *args)


def agent(
    fn=None,
    *,
    name=None,
    description=None,
    model=None,
    sys_prompt=None,
) -> AgentFunctionWrapper:
    args = (name, description, model, sys_prompt)

    def decorator(func) -> AgentFunctionWrapper:
        return AgentFunctionWrapper(func, *args)

    if fn is not None and callable(fn):
        return decorator(fn)
    return decorator
