"""
Integration tests for lifecycle hooks with LocalSandbox.
"""

from typing import Annotated

import pytest
from langgraph.graph import StateGraph

from torivers_sdk.automation.base import Automation
from torivers_sdk.automation.metadata import AutomationMetadata
from torivers_sdk.automation.state import BaseState, last_value
from torivers_sdk.testing.sandbox import LocalSandbox


class MyState(BaseState):
    """Test state for sandbox tests."""

    result: Annotated[str, last_value]


@pytest.mark.asyncio
class TestSandboxAsyncHooks:
    """Test suite for async hooks integration with LocalSandbox."""

    async def test_async_hooks_called_in_sandbox(self) -> None:
        """Test that async hooks are called during sandbox execution."""
        call_order = []

        class MyAutomation(Automation[MyState]):
            def metadata(self) -> AutomationMetadata:
                return AutomationMetadata(
                    name="test-auto",
                    title="Test Automation",
                    version="1.0.0",
                    description="Test automation with async hooks",
                )

            def get_state_class(self) -> type[MyState]:
                return MyState

            def build_graph(self) -> StateGraph:
                graph = StateGraph(MyState)
                graph.add_node("process", self.process)
                graph.set_entry_point("process")
                return graph

            def process(self, state: MyState) -> dict:
                call_order.append("process")
                return {
                    "result": "done",
                    "output_data": {
                        "_version": 1,
                        "_blocks": [
                            {
                                "type": "text",
                                "label": "Result",
                                "content": "done",
                                "format": "plain",
                            }
                        ],
                    },
                }

            async def on_startup_async(self) -> None:
                call_order.append("async_startup")

            def on_startup(self) -> None:
                call_order.append("sync_startup")

            def on_shutdown(self) -> None:
                call_order.append("sync_shutdown")

            async def on_shutdown_async(self) -> None:
                call_order.append("async_shutdown")

        automation = MyAutomation()
        sandbox = LocalSandbox()

        result = await sandbox.run(automation, {"test": "data"})

        # Verify execution was successful
        assert result.success
        assert result.output_data["_blocks"][0]["content"] == "done"

        # Verify hook call order
        assert call_order == [
            "async_startup",
            "sync_startup",
            "process",
            "sync_shutdown",
            "async_shutdown",
        ]

    async def test_async_hooks_with_state_modification(self) -> None:
        """Test async hooks that perform async operations."""
        import asyncio

        class MyAutomation(Automation[MyState]):
            def metadata(self) -> AutomationMetadata:
                return AutomationMetadata(
                    name="test-auto",
                    title="Test",
                    version="1.0.0",
                    description="Test with async operations",
                )

            def get_state_class(self) -> type[MyState]:
                return MyState

            def build_graph(self) -> StateGraph:
                graph = StateGraph(MyState)
                graph.add_node("process", self.process)
                graph.set_entry_point("process")
                return graph

            def process(self, state: MyState) -> dict:
                return {"result": "done"}

            async def on_startup_async(self) -> None:
                # Simulate async initialization (e.g., DB connection)
                await asyncio.sleep(0.01)
                self.initialized = True

            async def on_shutdown_async(self) -> None:
                # Simulate async cleanup (e.g., close connections)
                await asyncio.sleep(0.01)
                self.cleaned_up = True

        automation = MyAutomation()
        sandbox = LocalSandbox()

        result = await sandbox.run(automation, {"test": "data"})

        assert result.success
        assert hasattr(automation, "initialized")
        assert automation.initialized is True
        assert hasattr(automation, "cleaned_up")
        assert automation.cleaned_up is True

    async def test_only_async_hooks(self) -> None:
        """Test automation with only async hooks (no sync hooks)."""
        startup_called = False
        shutdown_called = False

        class MyAutomation(Automation[MyState]):
            def metadata(self) -> AutomationMetadata:
                return AutomationMetadata(
                    name="test-auto",
                    title="Test",
                    version="1.0.0",
                    description="Test with only async hooks",
                )

            def get_state_class(self) -> type[MyState]:
                return MyState

            def build_graph(self) -> StateGraph:
                graph = StateGraph(MyState)
                graph.add_node("process", self.process)
                graph.set_entry_point("process")
                return graph

            def process(self, state: MyState) -> dict:
                return {"result": "done"}

            async def on_startup_async(self) -> None:
                nonlocal startup_called
                startup_called = True

            async def on_shutdown_async(self) -> None:
                nonlocal shutdown_called
                shutdown_called = True

        automation = MyAutomation()
        sandbox = LocalSandbox()

        result = await sandbox.run(automation, {"test": "data"})

        assert result.success
        assert startup_called
        assert shutdown_called

    async def test_only_sync_hooks(self) -> None:
        """Test backward compatibility with only sync hooks."""
        startup_called = False
        shutdown_called = False

        class MyAutomation(Automation[MyState]):
            def metadata(self) -> AutomationMetadata:
                return AutomationMetadata(
                    name="test-auto",
                    title="Test",
                    version="1.0.0",
                    description="Test with only sync hooks",
                )

            def get_state_class(self) -> type[MyState]:
                return MyState

            def build_graph(self) -> StateGraph:
                graph = StateGraph(MyState)
                graph.add_node("process", self.process)
                graph.set_entry_point("process")
                return graph

            def process(self, state: MyState) -> dict:
                return {"result": "done"}

            def on_startup(self) -> None:
                nonlocal startup_called
                startup_called = True

            def on_shutdown(self) -> None:
                nonlocal shutdown_called
                shutdown_called = True

        automation = MyAutomation()
        sandbox = LocalSandbox()

        result = await sandbox.run(automation, {"test": "data"})

        assert result.success
        assert startup_called
        assert shutdown_called
