import asyncio
from collections.abc import Coroutine
import inspect
from typing import Any

from pydantic import BaseModel, Field
from typing_extensions import TypeIs

from crewai.agent import Agent
from crewai.lite_agent_output import LiteAgentOutput
from crewai.llms.base_llm import BaseLLM
from crewai.tasks.task_output import TaskOutput


def _is_coroutine(
    obj: LiteAgentOutput | Coroutine[Any, Any, LiteAgentOutput],
) -> TypeIs[Coroutine[Any, Any, LiteAgentOutput]]:
    """Check if obj is a coroutine for type narrowing."""
    return inspect.iscoroutine(obj)


class LLMGuardrailResult(BaseModel):
    valid: bool = Field(
        description="Whether the task output complies with the guardrail"
    )
    feedback: str | None = Field(
        description="A feedback about the task output if it is not valid",
        default=None,
    )


class LLMGuardrail:
    """It validates the output of another task using an LLM.

    This class is used to validate the output from a Task based on specified criteria.
    It uses an LLM to validate the output and provides a feedback if the output is not valid.

    Args:
        description (str): The description of the validation criteria.
        llm (LLM, optional): The language model to use for code generation.
    """

    def __init__(
        self,
        description: str,
        llm: BaseLLM,
    ):
        self.description = description

        self.llm: BaseLLM = llm

    def _validate_output(self, task_output: TaskOutput) -> LiteAgentOutput:
        agent = Agent(
            role="Guardrail Agent",
            goal="Validate the output of the task",
            backstory="You are a expert at validating the output of a task. By providing effective feedback if the output is not valid.",
            llm=self.llm,
        )

        query = f"""
        Ensure the following task result complies with the given guardrail.

        Task result:
        {task_output.raw}

        Guardrail:
        {self.description}

        Your task:
        - Confirm if the Task result complies with the guardrail.
        - If not, provide clear feedback explaining what is wrong (e.g., by how much it violates the rule, or what specific part fails).
        - Focus only on identifying issues — do not propose corrections.
        - If the Task result complies with the guardrail, saying that is valid
        """

        kickoff_result = agent.kickoff(query, response_format=LLMGuardrailResult)
        if _is_coroutine(kickoff_result):
            return asyncio.run(kickoff_result)
        return kickoff_result

    def __call__(self, task_output: TaskOutput) -> tuple[bool, Any]:
        """Validates the output of a task based on specified criteria.

        Args:
            task_output (TaskOutput): The output to be validated.

        Returns:
            Tuple[bool, Any]: A tuple containing:
                - bool: True if validation passed, False otherwise
                - Any: The validation result or error message
        """

        try:
            result = self._validate_output(task_output)
            if not isinstance(result.pydantic, LLMGuardrailResult):
                raise ValueError("The guardrail result is not a valid pydantic model")

            if result.pydantic.valid:
                return True, task_output.raw
            return False, result.pydantic.feedback
        except Exception as e:
            return False, f"Error while validating the task output: {e!s}"
