#!/usr/bin/env python3
#
# Copyright 2025 Google LLC
# SPDX-License-Identifier: Apache-2.0

"""Tests for the action module."""

import json
from typing import cast

import pytest

from genkit._core._action import (
    Action,
    ActionKind,
    ActionRunContext,
    DapQualifiedName,
    create_action_key,
    parse_action_key,
    parse_dap_qualified_name,
    parse_plugin_name_from_action_name,
)
from genkit._core._error import GenkitError


def test_action_enum_behaves_like_str() -> None:
    """Ensure the ActionType behaves like a string.

    This test verifies that the ActionType enum values can be compared
    directly with strings and that the correct variants are used.
    """
    assert ActionKind.CUSTOM == 'custom'
    assert ActionKind.EMBEDDER == 'embedder'
    assert ActionKind.EVALUATOR == 'evaluator'
    assert ActionKind.EXECUTABLE_PROMPT == 'executable-prompt'
    assert ActionKind.FLOW == 'flow'
    assert ActionKind.MODEL == 'model'
    assert ActionKind.PROMPT == 'prompt'
    assert ActionKind.TOOL == 'tool'
    assert ActionKind.UTIL == 'util'


def test_parse_action_key_valid() -> None:
    """Parse action key valid."""
    test_cases = [
        ('/prompt/my-prompt', (ActionKind.PROMPT, 'my-prompt')),
        ('/model/gpt-4', (ActionKind.MODEL, 'gpt-4')),
        (
            '/model/vertexai/gemini-1.0',
            (ActionKind.MODEL, 'vertexai/gemini-1.0'),
        ),
        ('/custom/test-action', (ActionKind.CUSTOM, 'test-action')),
        ('/flow/my-flow', (ActionKind.FLOW, 'my-flow')),
    ]

    for key, expected in test_cases:
        kind, name = parse_action_key(key)
        assert kind == expected[0]
        assert name == expected[1]


def test_parse_action_key_invalid_format() -> None:
    """Parse action key invalid format."""
    invalid_keys = [
        'invalid_key',  # Missing separator
        '/missing-kind',  # Missing kind
        'missing-name/',  # Missing name
        '',  # Empty string
        '/',  # Just separator
    ]

    for key in invalid_keys:
        with pytest.raises(ValueError, match='Invalid action key format'):
            parse_action_key(key)


def test_parse_dap_qualified_name() -> None:
    """Parse provider:innerKind/innerName segments."""
    assert parse_dap_qualified_name('my-dap:tool/echo') == DapQualifiedName('my-dap', 'tool', 'echo')
    assert parse_dap_qualified_name('plugin/foo:model/bar') is None
    assert parse_dap_qualified_name('plain-name') is None
    assert parse_dap_qualified_name('no-slash:toolonly') is None
    assert parse_dap_qualified_name(':tool/x') is None


def test_create_action_key() -> None:
    """Create action key."""
    assert create_action_key(ActionKind.CUSTOM, 'foo') == '/custom/foo'
    assert create_action_key(ActionKind.MODEL, 'foo') == '/model/foo'
    assert create_action_key(ActionKind.PROMPT, 'foo') == '/prompt/foo'
    assert create_action_key(ActionKind.TOOL, 'foo') == '/tool/foo'
    assert create_action_key(ActionKind.UTIL, 'foo') == '/util/foo'


def test_sync_action_rejected() -> None:
    """Sync functions are rejected - all actions must be async."""

    def sync_foo() -> str:
        return 'syncFoo'

    with pytest.raises(TypeError, match='Action handlers must be async functions'):
        Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=sync_foo)  # type: ignore[arg-type]


@pytest.mark.asyncio
async def test_define_async_action() -> None:
    """Define and run an async action."""

    async def async_foo() -> str:
        """An async action that returns 'asyncFoo'."""
        return 'asyncFoo'

    action = Action(name='asyncFoo', kind=ActionKind.CUSTOM, fn=async_foo)

    assert (await action.run()).response == 'asyncFoo'
    assert (await async_foo()) == 'asyncFoo'


@pytest.mark.asyncio
async def test_define_async_action_with_input() -> None:
    """Define and run an async action with input."""

    async def async_foo(input: str) -> str:
        """An async action that returns 'asyncFoo' with an input."""
        return f'asyncFoo {input}'

    action = Action(name='asyncFoo', kind=ActionKind.CUSTOM, fn=async_foo)

    assert (await action.run('foo')).response == 'asyncFoo foo'
    assert (await async_foo('foo')) == 'asyncFoo foo'


@pytest.mark.asyncio
async def test_define_async_action_with_input_and_context() -> None:
    """Define and run async action with input and context."""

    async def async_foo(input: str, ctx: ActionRunContext) -> str:
        """An async action that returns 'syncFoo' with an input and context."""
        return f'syncFoo {input} {ctx.context["foo"]}'

    action = Action(name='syncFoo', kind=ActionKind.CUSTOM, fn=async_foo)

    assert (await action.run('foo', context={'foo': 'bar'})).response == 'syncFoo foo bar'
    assert (await async_foo('foo', ActionRunContext(context={'foo': 'bar'}))) == 'syncFoo foo bar'


@pytest.mark.asyncio
async def test_streaming_action_with_callback() -> None:
    """Streaming action with on_chunk callback."""

    async def foo(
        input: str,
        ctx: ActionRunContext,
    ) -> int:
        ctx.send_chunk('1')
        ctx.send_chunk('2')
        return 3

    action = Action(name='foo', kind=ActionKind.CUSTOM, fn=foo)

    chunks: list[object] = []
    result = await action.run('foo', on_chunk=chunks.append)

    assert result.response == 3
    assert chunks == ['1', '2']


@pytest.mark.asyncio
async def test_streaming_action_with_stream_method() -> None:
    """Streaming action using the stream() method."""

    async def foo(
        input: str,
        ctx: ActionRunContext,
    ) -> int:
        ctx.send_chunk('1')
        ctx.send_chunk('2')
        return 3

    action = Action(name='foo', kind=ActionKind.CUSTOM, fn=foo)

    chunks: list[object] = []
    result = action.stream('foo')
    async for chunk in result.stream:
        chunks.append(chunk)

    assert await result.response == 3
    assert chunks == ['1', '2']


def test_parse_plugin_name_from_action_name() -> None:
    """Parse plugin name from the action name."""
    assert parse_plugin_name_from_action_name('foo') is None
    assert parse_plugin_name_from_action_name('foo/bar') == 'foo'
    assert parse_plugin_name_from_action_name('foo/bar/baz') == 'foo'


@pytest.mark.asyncio
async def test_propagates_context_via_contextvar() -> None:
    """Context is properly propagated via contextvar."""

    async def foo(_: str | None, ctx: ActionRunContext) -> str:
        return json.dumps(ctx.context)

    foo_action = cast(Action[str | None, str], Action(name='foo', kind=ActionKind.CUSTOM, fn=foo))

    async def bar() -> str:
        return (await foo_action.run()).response

    bar_action = cast(Action[None, str], Action(name='bar', kind=ActionKind.CUSTOM, fn=bar))

    async def baz() -> str:
        return (await bar_action.run()).response

    baz_action = cast(Action[None, str], Action(name='baz', kind=ActionKind.CUSTOM, fn=baz))

    first = baz_action.run(context={'foo': 'bar'})
    second = baz_action.run(context={'bar': 'baz'})

    assert (await second).response == '{"bar": "baz"}'
    assert (await first).response == '{"foo": "bar"}'


@pytest.mark.asyncio
async def test_action_raises_errors() -> None:
    """Action raises error with necessary metadata."""

    async def foo(_: str | None, ctx: ActionRunContext) -> None:
        raise Exception('oops')

    action = Action(name='fooAction', kind=ActionKind.CUSTOM, fn=foo)

    with pytest.raises(GenkitError, match=r'.*Error while running action fooAction.*') as e:
        await action.run()

    assert 'stack' in e.value.details
    assert 'trace_id' in e.value.details
    assert str(e.value.cause) == 'oops'


@pytest.mark.asyncio
async def test_run_raises_on_none_input_when_input_required() -> None:
    """run() raises GenkitError when input is None but the action requires it."""

    async def typed_fn(input: str) -> str:
        return f'got {input}'

    action = Action(name='typedAction', kind=ActionKind.CUSTOM, fn=typed_fn)

    with pytest.raises(GenkitError, match=r'.*requires input but none was provided.*'):
        await action.run(input=None)


@pytest.mark.asyncio
async def test_run_succeeds_with_valid_input() -> None:
    """run() succeeds when valid input is provided."""

    async def typed_fn(input: str) -> str:
        return f'got {input}'

    action = Action(name='typedAction', kind=ActionKind.CUSTOM, fn=typed_fn)

    result = await action.run(input='hello')
    assert result.response == 'got hello'


@pytest.mark.asyncio
async def test_run_no_input_type_allows_none() -> None:
    """run() allows None input when action has no input type."""

    async def no_input_fn() -> str:
        return 'no input needed'

    action = Action(name='noInputAction', kind=ActionKind.CUSTOM, fn=no_input_fn)

    result = await action.run(input=None)
    assert result.response == 'no input needed'


@pytest.mark.asyncio
async def test_run_defaulted_input_arg_allows_none() -> None:
    """A function with a Python default for its input should be callable with no input.

    Otherwise `await my_flow()` on `async def my_flow(name: str = 'world')`
    would surprise the caller with INVALID_ARGUMENT — typing accepts the
    call but the runtime would reject it.
    """

    async def greet(name: str = 'world') -> str:
        return f'hi {name}'

    action = Action(name='greet', kind=ActionKind.CUSTOM, fn=greet)

    assert (await action.run(input='Alice')).response == 'hi Alice'
    assert (await action.run(input=None)).response == 'hi world'
    assert (await action.run()).response == 'hi world'


@pytest.mark.asyncio
async def test_run_defaulted_input_arg_allows_none_with_ctx() -> None:
    """Same as above but for 2-arg (input + ctx) actions."""

    async def greet(name: str = 'world', ctx: ActionRunContext | None = None) -> str:
        return f'hi {name}'

    action = Action(name='greet_ctx', kind=ActionKind.CUSTOM, fn=greet)

    assert (await action.run(input='Bob')).response == 'hi Bob'
    assert (await action.run(input=None)).response == 'hi world'
    assert (await action.run()).response == 'hi world'
