"""Tests for the @python_step decorator."""

from pydantic import BaseModel

from areteos import InputMapping, Step, StepType, python_step
from areteos.step_decorator import PythonStepWrapper

# --- Test functions (defined at module level so inspect.getsource works) ---


@python_step
def simple_step(input: dict) -> dict:
    return {"result": input.get("value", 0) * 2}


class AnalyzeInput(BaseModel):
    text: str
    max_words: int = 50


class AnalyzeOutput(BaseModel):
    word_count: int
    truncated: str


@python_step
def typed_step(input: AnalyzeInput) -> AnalyzeOutput:
    words = input.text.split()[: input.max_words]
    return AnalyzeOutput(word_count=len(words), truncated=" ".join(words))


@python_step(requires=["pandas>=2.0", "numpy>=1.26"])
def step_with_deps(input: dict) -> dict:
    import pandas as pd

    df = pd.DataFrame(input["data"])
    return {"mean": float(df["value"].mean())}


@python_step(requires=["requests>=2.31"], python_version="3.12")
def step_custom_python(input: dict) -> dict:
    return {"ok": True}


# --- Tests ---


class TestDecoratorBasics:
    def test_returns_wrapper(self):
        assert isinstance(simple_step, PythonStepWrapper)

    def test_preserves_name(self):
        assert simple_step.__name__ == "simple_step"

    def test_callable_directly(self):
        result = simple_step({"value": 21})
        assert result == {"result": 42}

    def test_typed_callable_directly(self):
        result = typed_step(AnalyzeInput(text="hello world foo", max_words=2))
        assert result.word_count == 2
        assert result.truncated == "hello world"


class TestScriptGeneration:
    def test_simple_script_has_main(self):
        script = simple_step.script
        assert "def _main():" in script
        assert "simple_step(_input)" in script
        assert "json.dumps" in script

    def test_simple_script_has_input_loading(self):
        script = simple_step.script
        assert "sys.argv[1]" in script
        assert "json.load" in script

    def test_simple_script_no_pep723(self):
        script = simple_step.script
        assert "# /// script" not in script

    def test_typed_script_has_model_definitions(self):
        script = typed_step.script
        assert "class AnalyzeInput(BaseModel):" in script
        assert "class AnalyzeOutput(BaseModel):" in script
        assert "from pydantic import BaseModel" in script

    def test_typed_script_validates_input(self):
        script = typed_step.script
        assert "AnalyzeInput.model_validate(data)" in script

    def test_typed_script_dumps_output(self):
        script = typed_step.script
        assert "model_dump" in script

    def test_typed_script_includes_pydantic_dep(self):
        # Typed steps auto-add pydantic to PEP 723 requirements
        script = typed_step.script
        assert "# /// script" in script
        assert "pydantic" in script


class TestPEP723:
    def test_requirements_generate_header(self):
        script = step_with_deps.script
        assert "# /// script" in script
        assert '"pandas>=2.0"' in script
        assert '"numpy>=1.26"' in script
        assert "# ///" in script

    def test_custom_python_version(self):
        script = step_custom_python.script
        assert '>="3.12"' in script or ">=3.12" in script

    def test_requires_python_default(self):
        script = step_with_deps.script
        assert ">=3.11" in script


class TestAsStep:
    def test_produces_step(self):
        step = simple_step.as_step("my-step", "My Step")
        assert isinstance(step, Step)
        assert step.id == "my-step"
        assert step.name == "My Step"
        assert step.type == StepType.TOOL
        assert step.config["tool"] == "python"
        assert "script_content" in step.config

    def test_step_config_contains_script(self):
        step = simple_step.as_step("test", "Test")
        assert "simple_step" in step.config["script_content"]
        assert "_main()" in step.config["script_content"]

    def test_step_with_options(self):
        step = simple_step.as_step(
            "test",
            "Test",
            timeout_ms=10_000,
            max_retries=3,
            input_mappings=[
                InputMapping(source_step_id="prev", source_field="out", target_field="value")
            ],
        )
        assert step.timeout_ms == 10_000
        assert step.max_retries == 3
        assert len(step.input_mappings) == 1

    def test_typed_step_preserves_types(self):
        step = typed_step.as_step("typed", "Typed Step")
        # The script should contain the model classes
        assert "AnalyzeInput" in step.config["script_content"]
        assert "AnalyzeOutput" in step.config["script_content"]


class TestTypeIntrospection:
    def test_input_type_for_typed(self):
        assert typed_step.input_type is AnalyzeInput

    def test_output_type_for_typed(self):
        assert typed_step.output_type is AnalyzeOutput

    def test_input_fields_for_typed(self):
        assert "text" in typed_step.input_fields
        assert "max_words" in typed_step.input_fields

    def test_output_fields_for_typed(self):
        assert "word_count" in typed_step.output_fields
        assert "truncated" in typed_step.output_fields

    def test_input_type_for_untyped(self):
        assert simple_step.input_type is None

    def test_output_type_for_untyped(self):
        assert simple_step.output_type is None

    def test_input_fields_for_untyped(self):
        assert simple_step.input_fields == []
