from __future__ import annotations

from collections.abc import Awaitable, Callable
from typing import assert_type

import pytest
from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_ai.messages import ImageUrl, UserContent
from pydantic_ai.run import AgentRunResult
from pydantic_ai.tools import RunContext

from keble_helpers.ai import (
    ImagePromptChecker,
    arun_with_image_fallback,
    typed_tool,
    typed_tool_plain,
)


class _TypedOutput(BaseModel):
    value: str


class _TypedStubAgent:
    def __init__(self) -> None:
        self.prompts: list[str | list[UserContent]] = []

    async def run(
        self,
        user_prompt: str | list[UserContent] | None = None,
        *,
        deps: int | None = None,
        **kwargs: object,
    ) -> AgentRunResult[_TypedOutput]:
        _ = (deps, kwargs)
        normalized_prompt: str | list[UserContent]
        if isinstance(user_prompt, str):
            normalized_prompt = user_prompt
        elif user_prompt is None:
            normalized_prompt = ""
        else:
            normalized_prompt = list(user_prompt)
        self.prompts.append(normalized_prompt)
        return AgentRunResult(output=_TypedOutput(value="ok"))


@pytest.mark.asyncio
async def test_arun_with_image_fallback_preserves_agent_run_result_type() -> None:
    agent = _TypedStubAgent()

    result = await arun_with_image_fallback(
        agent=agent,
        prompt=["input", ImageUrl(url="https://img.example.com/a.jpg")],
        image_prompt_checker=ImagePromptChecker(enabled=False),
    )

    typed_result = assert_type(result, AgentRunResult[_TypedOutput])
    output = assert_type(typed_result.output, _TypedOutput)
    assert output.value == "ok"
    assert agent.prompts == [["input", ImageUrl(url="https://img.example.com/a.jpg")]]


@pytest.mark.asyncio
async def test_checker_method_preserves_agent_run_result_type() -> None:
    agent = _TypedStubAgent()
    checker = ImagePromptChecker(enabled=False)

    result = await checker.arun_with_image_fallback(
        agent=agent,
        prompt="input text",
        deps=3,
    )

    typed_result = assert_type(result, AgentRunResult[_TypedOutput])
    output = assert_type(typed_result.output, _TypedOutput)
    assert output.value == "ok"
    assert agent.prompts == ["input text"]


def test_typed_tool_registers_context_tool_without_losing_callable_type() -> None:
    agent: Agent[int, str] = Agent("test", deps_type=int, output_type=str)

    @typed_tool(agent, require_parameter_descriptions=True)
    async def repeat(ctx: RunContext[int], word: str) -> str:
        """Repeat a word with deps.

        Args:
            word: Word to repeat.
        """

        return f"{ctx.deps}:{word}"

    typed_repeat = assert_type(
        repeat,
        Callable[[RunContext[int], str], Awaitable[str]],
    )
    assert typed_repeat is repeat
    assert "repeat" in agent._function_toolset.tools
    assert agent._function_toolset.tools["repeat"].function is repeat


def test_typed_tool_plain_registers_plain_tool_without_losing_callable_type() -> None:
    agent: Agent[None, str] = Agent("test", output_type=str)

    def slugify(name: str) -> str:
        return name.strip().lower().replace(" ", "-")

    registered = typed_tool_plain(
        agent,
        slugify,
        name="slugify",
        description="Normalize a label into a slug.",
    )

    typed_registered = assert_type(registered, Callable[[str], str])
    assert typed_registered is slugify
    assert typed_registered("Hello There") == "hello-there"
    assert "slugify" in agent._function_toolset.tools
    assert agent._function_toolset.tools["slugify"].function is slugify
