"""Tests for staso.decorators (@agent, @tool, @trace)."""
from __future__ import annotations

import asyncio
import json

import pytest

import staso as st


class TestAgentDecorator:
    def test_sync_agent(self, capturing):
        @st.agent(name="my-agent")
        def run(msg: str) -> str:
            return f"echo: {msg}"

        result = run("hello")
        assert result == "echo: hello"
        assert len(capturing) == 1
        assert capturing[0]["name"] == "my-agent"
        assert capturing[0]["kind"] == "agent"
        assert capturing[0]["status"] == "ok"

    def test_async_agent(self, capturing):
        @st.agent(name="async-agent")
        async def run(msg: str) -> str:
            return f"echo: {msg}"

        result = asyncio.run(run("hello"))
        assert result == "echo: hello"
        assert len(capturing) == 1
        assert capturing[0]["name"] == "async-agent"
        assert capturing[0]["kind"] == "agent"

    def test_agent_default_name(self, capturing):
        @st.agent()
        def my_function():
            return 42

        my_function()
        assert capturing[0]["name"] == "my_function"

    def test_agent_captures_input(self, capturing):
        @st.agent(name="test-agent")
        def run():
            return "ok"

        run()
        input_val = json.loads(capturing[0]["input"])
        assert input_val["agent_name"] == "test-agent"

    def test_agent_error_propagates(self, capturing):
        @st.agent(name="failing")
        def run():
            raise ValueError("boom")

        with pytest.raises(ValueError, match="boom"):
            run()

        assert len(capturing) == 1
        assert capturing[0]["status"] == "error"
        assert capturing[0]["error_message"] == "boom"

    def test_agent_async_error_propagates(self, capturing):
        @st.agent(name="failing-async")
        async def run():
            raise RuntimeError("async boom")

        with pytest.raises(RuntimeError, match="async boom"):
            asyncio.run(run())

        assert capturing[0]["status"] == "error"

    def test_agent_with_tags(self, capturing):
        @st.agent(name="tagged", tags=["tier-1", "support"])
        def run():
            return "ok"

        run()
        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["tags"] == ["tier-1", "support"]

    def test_agent_preserves_function_name(self):
        @st.agent(name="custom")
        def original_func():
            pass

        assert original_func.__name__ == "original_func"


class TestToolDecorator:
    def test_sync_tool(self, capturing):
        @st.tool(name="search")
        def search(query: str, limit: int = 5) -> str:
            return f"results for {query}"

        result = search("hello", limit=10)
        assert result == "results for hello"
        assert len(capturing) == 1
        assert capturing[0]["name"] == "search"
        assert capturing[0]["kind"] == "tool"
        assert capturing[0]["status"] == "ok"

    def test_tool_captures_args(self, capturing):
        @st.tool(name="search")
        def search(query: str, limit: int = 5) -> str:
            return "ok"

        search("hello", limit=10)
        input_val = json.loads(capturing[0]["input"])
        assert input_val["query"] == "hello"
        assert input_val["limit"] == 10

    def test_tool_captures_output(self, capturing):
        @st.tool(name="search")
        def search(query: str) -> str:
            return "found it"

        search("test")
        assert json.loads(capturing[0]["output"])["result"] == "found it"

    def test_async_tool(self, capturing):
        @st.tool(name="async-search")
        async def search(query: str) -> str:
            return f"async: {query}"

        result = asyncio.run(search("hello"))
        assert result == "async: hello"
        assert capturing[0]["name"] == "async-search"
        assert capturing[0]["kind"] == "tool"

    def test_tool_error(self, capturing):
        @st.tool(name="broken")
        def broken():
            raise ConnectionError("timeout")

        with pytest.raises(ConnectionError, match="timeout"):
            broken()

        assert capturing[0]["status"] == "error"
        assert capturing[0]["error_message"] == "timeout"

    def test_tool_default_name(self, capturing):
        @st.tool()
        def my_tool_func():
            return 1

        my_tool_func()
        assert capturing[0]["name"] == "my_tool_func"

    def test_tool_with_tags(self, capturing):
        @st.tool(name="tagged", tags=["retrieval"])
        def search():
            return "ok"

        search()
        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["tags"] == ["retrieval"]

    def test_tool_positional_and_keyword_args(self, capturing):
        @st.tool(name="multi")
        def multi(a: int, b: str, c: float = 1.0) -> str:
            return "ok"

        multi(1, "two", c=3.0)
        input_val = json.loads(capturing[0]["input"])
        assert input_val["a"] == 1
        assert input_val["b"] == "two"
        assert input_val["c"] == 3.0


class TestTraceDecorator:
    def test_sync_trace(self, capturing):
        @st.trace(name="classify", kind="CHAIN")
        def classify(msg: str) -> str:
            return "intent"

        result = classify("hello")
        assert result == "intent"
        assert capturing[0]["name"] == "classify"
        assert capturing[0]["kind"] == "chain"

    def test_trace_default_kind_is_chain(self, capturing):
        @st.trace(name="step")
        def step():
            return "ok"

        step()
        assert capturing[0]["kind"] == "chain"

    def test_trace_retriever_kind(self, capturing):
        @st.trace(name="fetch", kind="RETRIEVER")
        def fetch():
            return "data"

        fetch()
        assert capturing[0]["kind"] == "retriever"

    def test_trace_llm_kind(self, capturing):
        @st.trace(name="llm", kind="LLM")
        def call_llm():
            return "response"

        call_llm()
        assert capturing[0]["kind"] == "llm"

    def test_async_trace(self, capturing):
        @st.trace(name="async-chain", kind="CHAIN")
        async def process():
            return "done"

        result = asyncio.run(process())
        assert result == "done"
        assert capturing[0]["name"] == "async-chain"

    def test_trace_captures_io(self, capturing):
        """@trace captures inputs/outputs."""
        @st.trace(name="step")
        def step(x: int) -> int:
            return x * 2

        step(5)
        input_val = json.loads(capturing[0]["input"])
        assert input_val["x"] == 5
        output_val = json.loads(capturing[0]["output"])
        assert output_val["result"] == 10

    def test_trace_error(self, capturing):
        @st.trace(name="failing")
        def failing():
            raise TypeError("wrong type")

        with pytest.raises(TypeError):
            failing()

        assert capturing[0]["status"] == "error"

    def test_trace_with_tags(self, capturing):
        @st.trace(name="tagged", kind="CHAIN", tags=["nlp"])
        def classify():
            return "ok"

        classify()
        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["tags"] == ["nlp"]


class TestNesting:
    def test_tool_inside_agent_shares_trace_id(self, capturing):
        @st.tool(name="inner-tool")
        def my_tool():
            return "done"

        @st.agent(name="outer-agent")
        def my_agent():
            return my_tool()

        my_agent()

        assert len(capturing) == 2
        assert capturing[0]["trace_id"] == capturing[1]["trace_id"]

    def test_child_has_parent_span_id(self, capturing):
        @st.tool(name="child")
        def child():
            return "ok"

        @st.agent(name="parent")
        def parent():
            return child()

        parent()

        # child ends first (captured[0]), parent ends second (captured[1])
        child_span = capturing[0]
        parent_span = capturing[1]
        assert child_span["parent_span_id"] == parent_span["span_id"]

    def test_deep_nesting(self, capturing):
        @st.trace(name="level3", kind="CHAIN")
        def level3():
            return "deep"

        @st.tool(name="level2")
        def level2():
            return level3()

        @st.agent(name="level1")
        def level1():
            return level2()

        level1()

        assert len(capturing) == 3
        trace_ids = {s["trace_id"] for s in capturing}
        assert len(trace_ids) == 1


class TestDecoratorMetadata:
    def test_agent_metadata(self, capturing):
        @st.agent(name="meta-agent", metadata={"version": "1.0", "env": "test"})
        def run():
            return "ok"

        run()
        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["version"] == "1.0"
        assert attrs["env"] == "test"

    def test_tool_metadata(self, capturing):
        @st.tool(name="meta-tool", metadata={"source": "db"})
        def search():
            return "found"

        search()
        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["source"] == "db"

    def test_trace_metadata(self, capturing):
        @st.trace(name="meta-trace", metadata={"step": 1})
        def step():
            return "done"

        step()
        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["step"] == 1

    def test_metadata_with_tags(self, capturing):
        @st.agent(name="both", tags=["prod"], metadata={"v": "2"})
        def run():
            return "ok"

        run()
        attrs = json.loads(capturing[0]["attributes"])
        assert attrs["tags"] == ["prod"]
        assert attrs["v"] == "2"


class TestIOCaptureToggle:
    def test_agent_no_input(self, capturing):
        @st.agent(name="no-in", capture_input=False)
        def run(x):
            return x

        run(42)
        assert capturing[0]["input"] == ""

    def test_agent_no_output(self, capturing):
        @st.agent(name="no-out", capture_output=False)
        def run():
            return "secret"

        run()
        output = capturing[0]["output"]
        assert output == "" or json.loads(output) == {}

    def test_tool_no_io(self, capturing):
        @st.tool(name="silent", capture_input=False, capture_output=False)
        def search(q):
            return "result"

        search("query")
        assert capturing[0]["input"] == ""
        output = capturing[0]["output"]
        assert output == "" or json.loads(output) == {}

    def test_trace_no_io(self, capturing):
        @st.trace(name="quiet", capture_input=False, capture_output=False)
        def step(x):
            return x * 2

        step(5)
        assert capturing[0]["input"] == ""
        output = capturing[0]["output"]
        assert output == "" or json.loads(output) == {}

    def test_default_captures_io(self, capturing):
        @st.tool(name="loud")
        def search(q):
            return "found"

        search("test")
        assert json.loads(capturing[0]["input"])["q"] == "test"
        assert json.loads(capturing[0]["output"])["result"] == "found"
