import sys
import threading
from unittest.mock import patch
from typing import Callable

from pydantic import BaseModel, Field

from versionhq.agent.model import Agent, LLM
from versionhq.task.model import Task, ResponseField, TaskOutput, TaskExecutionType
from versionhq.tool.model import Tool, ToolSet
from versionhq.tool.decorator import tool

sys.setrecursionlimit(2097152)
threading.stack_size(134217728)



def test_async_execute_task():
    task = Task(description="Return string: 'test'", type=TaskExecutionType.ASYNC)

    with patch.object(Agent, "execute_task", return_value="test") as execute:
        res = task.execute()
        assert res.raw == "test"
        execute.assert_called_once_with(task=task, context=None, task_tools=list())


def test_sync_execute_with_task_context():
    """
    Use case = One agent handling multiple tasks sequentially using context set in the main task.
    """
    sub_task = Task(
        description="return the output following the given prompt.",
        response_fields=[
            ResponseField(title="subtask_result", data_type=str, required=True),
        ]
    )
    main_task = Task(
        description="return the output following the given prompt.",
        response_fields=[
            ResponseField(title="test1", data_type=int, required=True),
            ResponseField(title="test2", data_type=str, required=True),
        ],
        context=[sub_task,]
    )
    res = main_task.execute()

    assert isinstance(res, TaskOutput)
    assert res.task_id is main_task.id
    assert res.raw is not None
    assert isinstance(res.raw, str)
    assert res.json_dict is not None
    assert isinstance(res.json_dict, dict)
    assert res.pydantic is None
    assert sub_task.output is not None
    assert sub_task.output.json_dict is not None
    assert "subtask_result" in main_task.prompt()


def test_sync_execute_task_with_prompt_context():
    """
    Use case:
    - One agent handling multiple tasks sequentially using context set in the main task.
    - On top of that, the agent receives context when they execute the task.
    """
    class Outcome(BaseModel):
        test1: int = Field(default=None)
        test2: str = Field(default=None)

    sub_task = Task(
        description="return the output following the given prompt.",
        response_fields=[ResponseField(title="result", data_type=str, required=True),]
    )
    main_task = Task(
        description="return the output following the given prompt.",
        pydantic_output=Outcome,
        response_fields=[
            ResponseField(title="test1", data_type=int, required=True),
            ResponseField(title="test2", data_type=str, required=True),
        ],
        context=[sub_task]
    )
    res = main_task.execute(context="plan a Black Friday campaign.")

    assert isinstance(res, TaskOutput) and res.task_id is main_task.id
    assert res.raw and isinstance(res.raw, str)
    assert res.json_dict and isinstance(res.json_dict, dict)
    assert res.pydantic.test1 == res.json_dict["test1"] and res.pydantic.test2 == res.json_dict["test2"]
    assert sub_task.output is not None
    assert "result" in main_task.prompt()
    assert main_task.prompt_context == "plan a Black Friday campaign."
    assert "plan a Black Friday campaign." in main_task.prompt()


def test_callback():
    """
    See if the callback function is executed well with kwargs.
    """

    def callback_func(condition: str, test1: str):
        # task_id = str(id) if id else None
        return f"Result: {test1}, condition added: {condition}"

    task = Task(
        description="return the output following the given prompt.",
        response_fields=[
            ResponseField(title="test1", data_type=str, required=True),
        ],
        callback=callback_func,
        callback_kwargs=dict(condition="demo for pytest")
    )
    res = task.execute()

    assert res and isinstance(res, TaskOutput)
    assert res.task_id is task.id
    assert "demo for pytest" in res.callback_output


def test_delegate():
    agent = Agent(role="demo agent 6", goal="My amazing goals", maxit=1, max_tokens=3000)
    task = Task(
        description="return the output following the given prompt.",
        response_fields=[
            ResponseField(title="test1", data_type=str, required=True),
        ],
        allow_delegation=True
    )
    task.execute()

    assert task.output is not None
    assert "vhq-Delegated-Agent" in task.processed_agents
    assert task.delegations != 0


# def test_conditional_task():
#     task = Task(
#         description="erturn the output following the given prompt.",
#         response_fields=[ResponseField(title="test1", data_type=str, required=True),],
#     )
#     res = task.execute_sync(agent=base_agent)

#     conditional_task = ConditionalTask(
#         description="return the output following the given prompt.",
#         response_fields=[ResponseField(title="test1", data_type=str, required=True),],
#         condition=lambda x: bool("zzz" in task.output.raw)
#     )
#     should_execute = conditional_task.should_execute(context=res)

#     assert res.raw is not None
#     assert should_execute is False

#     conditional_res = conditional_task._handle_conditional_task(task_outputs=[res,], task_index=1, was_replayed=False)
#     if not should_execute:
#         assert conditional_res is None
#     else:
#         assert conditional_res.task_id is conditional_task.id


def test_store_task_log():
    task = Task(
        description="return the output following the given prompt.",
        response_fields=[ResponseField(title="task_1", data_type=str, required=True),],
    )
    assert task._task_output_handler.load() is not None


def test_task_with_agent_tools():
    simple_tool = Tool(name="simple tool", func=lambda x: "simple func")
    agent = Agent(role="demo", goal="execute tools", tools=[simple_tool,], maxit=1, max_tokens=3000)
    task = Task(description="execute tool", can_use_agent_tools=True, tool_res_as_final=True)
    res = task.execute(agent=agent)
    assert res.tool_output == "simple func"

    def empty_func():
        return "empty func"

    func_tool = Tool(name="func tool", func=empty_func)
    agent.tools = [func_tool]
    res = task.execute(agent=agent)
    assert res.tool_output == "empty func"


    def demo_func(message: str) -> str:
        return message + "_demo"

    class CustomTool(Tool):
        name: str = "custom tool"

    custom_tool = CustomTool(func=demo_func)
    agent.tools = [custom_tool]
    res = task.execute(agent=agent)
    assert "_demo" in res.tool_output


def test_task_with_tools():

    def random_func(message: str) -> str:
        return message + "_demo"

    tool = Tool(name="tool", func=random_func)
    tool_set = ToolSet(tool=tool, kwargs=dict(message="empty func"))
    task = Task(description="execute the given tools", tools=[tool_set,], tool_res_as_final=True)
    res = task.execute()
    assert res.tool_output == "empty func_demo"

    class CustomTool(Tool):
        name: str = "custom tool"
        func: Callable = None

    custom_tool = CustomTool(func=random_func)
    task.tools = [custom_tool]
    res = task.execute()
    assert "_demo" in res.tool_output

    task.tools = [custom_tool]
    res = task.execute()
    assert res.tool_output is not None



def test_task_without_response_format():
    task = Task(description="return a simple output with any random values.")
    res = task.execute()

    assert res and isinstance(res, TaskOutput)
    assert res.json_dict and isinstance(res.json_dict, dict)
    assert res.pydantic is None


def test_build_agent_without_developer_prompt():
    agent = Agent(
        role="analyst",
        goal="analyze the company's website and retrieve the product overview",
        backstory="You are competitive analysts who have abundand knowledge in marketing, product management.",
        use_developer_prompt=False
    )
    task = Task(description="return a simple output with any random values.")
    res = task.execute(agent=agent)

    assert res and isinstance(res, TaskOutput)
    assert res.json_dict and isinstance(res.json_dict, dict)
    assert res.pydantic is None


def test_callback_with_custom_output():
    class CustomOutput(BaseModel):
        test1: str
        test2: list[str]

    def dummy_func(message: str, test1: str, test2: list[str]) -> str:
        return f"""{message}: {test1}, {", ".join(test2)}"""

    task = Task(
        description="Amazing task",
        pydantic_output=CustomOutput,
        callback=dummy_func,
        callback_kwargs=dict(message="Hi! Here is the result: ")
    )
    res = task.execute(context="amazing context to consider.")

    assert res.task_id == task.id
    assert res.pydantic.test1 and res.pydantic.test2
    assert "Hi! Here is the result: " in res.callback_output and res.pydantic.test1 in res.callback_output and ", ".join(res.pydantic.test2) in res.callback_output


def test_task_with_agent_callback():
    import litellm

    def dummy_func(*args, **kwargs) -> str:
        return "Demo func"

    agent = Agent(role="demo", goal="amazing project goal", maxit=1, max_tokens=3000, callbacks=[dummy_func,])
    task = Task(description="Amazing task")
    res = task.execute(agent=agent)

    assert res.raw and res.task_id == task.id
    assert litellm.callbacks == [dummy_func]


def test_rpm():
    agent = Agent(role="demo", goal="use the given tools", max_tokens=3000, max_rpm=3)
    assert agent._rpm_controller and agent._rpm_controller.max_rpm == 3

    a = "hello"
    tool = Tool(func=lambda x: a + x)
    tool_set = ToolSet(tool=tool, kwargs={ "x": "_demo" })
    task = Task(description="Summarize overview of the given tool in sentence, then execute the tool.", tools=[tool_set,])
    res = task.execute(agent=agent)

    assert "hello_demo" in res.raw if res.raw else res
    assert agent._rpm_controller._current_rpm < agent._rpm_controller.max_rpm


def test_maxit():
    @tool
    def demo() -> str:
        """Get the final answer but don't give it yet, just re-use this
        tool non-stop."""
        return "demo"

    agent = Agent(role="demo", goal="amazing demo", maxit=2)
    task = Task(description="Summarize overview of the given tool in sentences.", tools=[demo,])

    with patch.object(LLM, "call", wraps=agent.llm.call) as mock:
        task.execute(agent=agent)
        assert mock.call_count <= 2


def test_evaluation():
    """
    See if the output will be evaluated accurately - when the task was given eval criteria
    """
    from versionhq.task.evaluate import Evaluation, EvaluationItem
    from versionhq.agent.inhouse_agents import vhq_task_evaluator

    task = Task(
        description="Research a topic to teach a kid aged 6 about math.",
        should_evaluate=True,
        eval_criteria=["Uniquness of the topic researched", "Fit to the target audience",],
    )
    res = task.execute()

    assert res.evaluation and isinstance(res.evaluation, Evaluation)
    assert [isinstance(item, EvaluationItem) and item.criteria in task.eval_criteria for item in res.evaluation.items]
    assert res.evaluation.latency and res.evaluation.tokens and res.evaluation.eval_by == vhq_task_evaluator
    assert res.evaluation.aggregate_score is not None and res.evaluation.suggestion_summary
