"""
Local sandbox for testing automations.

This module provides a local sandbox environment that simulates
the ToRivers runtime for development and testing purposes.
"""

import asyncio
from dataclasses import dataclass
from typing import Any, TypeVar

from torivers_sdk.automation.base import Automation
from torivers_sdk.automation.output import JsonDataBlock, TextBlock, build_output_data
from torivers_sdk.automation.state import BaseState
from torivers_sdk.context.credentials import CredentialProxy
from torivers_sdk.context.execution import ExecutionContext
from torivers_sdk.context.mocks import MockCredentialProxy
from torivers_sdk.context.progress import ProgressEntry, ProgressReporter

S = TypeVar("S", bound=BaseState)


@dataclass
class ExecutionResult:
    """Result of a sandbox execution."""

    success: bool
    output_data: dict[str, Any]
    errors: list[dict[str, Any]]
    progress_entries: list[ProgressEntry]
    execution_time_seconds: float
    final_state: dict[str, Any]

    def get_output(self, key: str, default: Any = None) -> Any:
        """Get a value from json_data output blocks."""
        blocks = self.output_data.get("_blocks")
        if not isinstance(blocks, list):
            return default

        for block in blocks:
            if not isinstance(block, dict):
                continue
            if block.get("type") != "json_data":
                continue

            data = block.get("data")
            if isinstance(data, dict) and key in data:
                return data[key]

        return default


@dataclass
class SandboxConfig:
    """Configuration for the local sandbox."""

    execution_id: str = "test-execution-001"
    timeout_seconds: float = 300.0
    max_steps: int = 100


class LocalSandbox:
    """
    Local sandbox for testing automations.

    This sandbox provides a simulated runtime environment that allows
    developers to test their automations locally before submission.

    Example:
        from torivers_sdk.testing import LocalSandbox, MockCredentialProxy

        # Create sandbox
        sandbox = LocalSandbox()

        # Configure mock credentials
        credentials = MockCredentialProxy(["gmail", "sheets"])
        sandbox.set_credentials(credentials)

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

        # Check results
        assert result.success
        assert result.output_data.get("_blocks")
    """

    def __init__(self, config: SandboxConfig | None = None) -> None:
        """
        Initialize the local sandbox.

        Args:
            config: Optional sandbox configuration
        """
        self._config = config or SandboxConfig()
        self._credentials: CredentialProxy | None = None
        self._progress_reporter: ProgressReporter | None = None

    def set_credentials(self, credentials: CredentialProxy) -> None:
        """
        Set the credential proxy for the sandbox.

        Args:
            credentials: Credential proxy (usually MockCredentialProxy)
        """
        self._credentials = credentials

    async def run(
        self,
        automation: Automation[S],
        input_data: dict[str, Any],
    ) -> ExecutionResult:
        """
        Run an automation in the sandbox.

        Args:
            automation: The automation instance to run
            input_data: Input data for the automation

        Returns:
            Execution result with output and diagnostics
        """
        import time

        start_time = time.time()

        # Create progress reporter
        self._progress_reporter = ProgressReporter(
            execution_id=self._config.execution_id
        )

        # Create execution context
        context = ExecutionContext(
            execution_id=self._config.execution_id,
            progress=self._progress_reporter,
        )

        # Set up credentials
        allowed_services = (
            automation.get_required_credentials()
            + automation.get_optional_credentials()
        )
        credentials = self._credentials or MockCredentialProxy(allowed_services)

        # Validate input
        try:
            automation.validate_input(input_data)
        except Exception as e:
            return ExecutionResult(
                success=False,
                output_data={},
                errors=[{"type": "validation_error", "message": str(e)}],
                progress_entries=self._progress_reporter.entries,
                execution_time_seconds=time.time() - start_time,
                final_state={},
            )

        # Build and compile the graph
        try:
            graph = automation.build_graph()
            compiled = graph.compile()
        except Exception as e:
            return ExecutionResult(
                success=False,
                output_data={},
                errors=[{"type": "graph_error", "message": str(e)}],
                progress_entries=self._progress_reporter.entries,
                execution_time_seconds=time.time() - start_time,
                final_state={},
            )

        # Prepare initial state
        initial_state: dict[str, Any] = {
            "execution_id": self._config.execution_id,
            "input_data": input_data,
            "output_data": {},
            "context": context,
            "credentials": credentials,
            "current_step": "",
            "errors": [],
        }

        # Run the graph
        try:
            # Call async startup hook first, then sync hook
            await automation.on_startup_async()
            automation.on_startup()

            final_state = await asyncio.wait_for(
                self._run_graph(compiled, initial_state),
                timeout=self._config.timeout_seconds,
            )

            # Call sync shutdown hook first, then async hook
            automation.on_shutdown()
            await automation.on_shutdown_async()

            return ExecutionResult(
                success=len(final_state.get("errors", [])) == 0,
                output_data=self._normalize_output_data(final_state.get("output_data")),
                errors=final_state.get("errors", []),
                progress_entries=self._progress_reporter.entries,
                execution_time_seconds=time.time() - start_time,
                final_state=final_state,
            )

        except asyncio.TimeoutError:
            return ExecutionResult(
                success=False,
                output_data={},
                errors=[{"type": "timeout", "message": "Execution timed out"}],
                progress_entries=self._progress_reporter.entries,
                execution_time_seconds=time.time() - start_time,
                final_state={},
            )

        except Exception as e:
            automation.on_error(e)
            return ExecutionResult(
                success=False,
                output_data={},
                errors=[{"type": "execution_error", "message": str(e)}],
                progress_entries=self._progress_reporter.entries,
                execution_time_seconds=time.time() - start_time,
                final_state={},
            )

    async def _run_graph(
        self,
        compiled: Any,
        initial_state: dict[str, Any],
    ) -> dict[str, Any]:
        """
        Execute the compiled graph.

        Args:
            compiled: Compiled LangGraph
            initial_state: Initial state dictionary

        Returns:
            Final state after execution
        """
        # Run the graph asynchronously
        result = await compiled.ainvoke(initial_state)
        return result

    def _normalize_output_data(self, output_data: Any) -> dict[str, Any]:
        """Normalize sandbox output to the `_blocks` envelope."""
        if isinstance(output_data, dict):
            blocks = output_data.get("_blocks")
            if isinstance(blocks, list) and blocks:
                return output_data
            if output_data:
                return build_output_data(
                    JsonDataBlock(
                        label="Execution Output",
                        data=output_data,
                        display_hint="key_value",
                    )
                )

        return build_output_data(
            TextBlock(
                label="Execution Output",
                content="Automation completed successfully.",
                format="plain",
            )
        )
