#!/usr/bin/env python3
"""
ToRivers Automation Sandbox Entrypoint.

This script is the entry point for running automations in the Docker sandbox.
It loads the automation, executes it with the provided input, and writes
the results to the output file.
"""

from __future__ import annotations

import asyncio
import importlib.util
import json
import os
import sys
import time
import traceback
from pathlib import Path
from typing import Any


def load_automation(automation_path: str) -> Any:
    """
    Load the automation class from the given path.

    Args:
        automation_path: Path to the automation directory.

    Returns:
        The automation instance.

    Raises:
        ImportError: If automation cannot be loaded.
    """
    main_file = Path(automation_path) / "main.py"
    if not main_file.exists():
        raise ImportError(f"main.py not found in {automation_path}")

    # Add automation path to sys.path for imports
    sys.path.insert(0, automation_path)

    # Load the main module
    spec = importlib.util.spec_from_file_location("main", main_file)
    if spec is None or spec.loader is None:
        raise ImportError(f"Cannot load spec for {main_file}")

    module = importlib.util.module_from_spec(spec)
    sys.modules["main"] = module
    spec.loader.exec_module(module)

    # Get the automation class/instance
    if hasattr(module, "automation"):
        automation = module.automation
        # If it's a class, instantiate it
        if isinstance(automation, type):
            return automation()
        return automation
    else:
        raise ImportError("No 'automation' attribute found in main.py")


def setup_mocks() -> dict[str, Any]:
    """
    Set up mock credentials and services.

    Returns:
        Dictionary of mock services.
    """
    try:
        from torivers_sdk.testing.mocks import (
            MockCredentialProxy,
            MockHTTPClient,
            MockLLMClient,
            MockStorageClient,
        )

        return {
            "credentials": MockCredentialProxy([]),
            "llm": MockLLMClient(),
            "http": MockHTTPClient(),
            "storage": MockStorageClient(),
        }
    except ImportError:
        return {}


def write_output(output_file: str, data: dict[str, Any]) -> None:
    """
    Write output to the specified file.

    Args:
        output_file: Path to output file.
        data: Data to write.
    """
    try:
        with open(output_file, "w") as f:
            json.dump(data, f, indent=2, default=str)
    except Exception as e:
        print(f"Error writing output: {e}", file=sys.stderr)


async def run_automation(
    automation: Any,
    input_data: dict[str, Any],
    use_mocks: bool,
) -> dict[str, Any]:
    """
    Run the automation with the given input.

    Args:
        automation: The automation instance.
        input_data: Input data for the automation.
        use_mocks: Whether to use mock credentials.

    Returns:
        Dictionary with execution results.
    """
    start_time = time.time()
    tokens_used = 0

    try:
        # Set up mocks if requested
        mocks = setup_mocks() if use_mocks else {}

        # Validate input
        if hasattr(automation, "validate_input"):
            automation.validate_input(input_data)

        # Build the graph
        if hasattr(automation, "build_graph"):
            graph = automation.build_graph()
            compiled = graph.compile()
        else:
            raise RuntimeError("Automation does not have build_graph method")

        # Prepare initial state
        from torivers_sdk.context.execution import ExecutionContext
        from torivers_sdk.context.progress import ProgressReporter

        progress = ProgressReporter(execution_id="sandbox-execution")
        context = ExecutionContext(
            execution_id="sandbox-execution",
            progress=progress,
        )

        initial_state = {
            "execution_id": "sandbox-execution",
            "input_data": input_data,
            "output_data": {},
            "context": context,
            "credentials": mocks.get("credentials"),
            "current_step": "",
            "errors": [],
        }

        # Run lifecycle hooks
        if hasattr(automation, "on_startup_async"):
            await automation.on_startup_async()
        if hasattr(automation, "on_startup"):
            automation.on_startup()

        # Execute the graph
        final_state = await compiled.ainvoke(initial_state)

        # Run shutdown hooks
        if hasattr(automation, "on_shutdown"):
            automation.on_shutdown()
        if hasattr(automation, "on_shutdown_async"):
            await automation.on_shutdown_async()

        # Extract results
        output_data = final_state.get("output_data", {})
        errors = final_state.get("errors", [])

        return {
            "success": len(errors) == 0,
            "output_data": output_data,
            "errors": errors,
            "execution_time_seconds": time.time() - start_time,
            "tokens_used": tokens_used,
        }

    except Exception as e:
        traceback.print_exc()
        return {
            "success": False,
            "output_data": {},
            "error": str(e),
            "execution_time_seconds": time.time() - start_time,
            "tokens_used": tokens_used,
        }


def main() -> int:
    """
    Main entry point.

    Returns:
        Exit code (0 for success, non-zero for failure).
    """
    # Get configuration from environment
    automation_path = os.environ.get("AUTOMATION_PATH", "/automation")
    input_data_str = os.environ.get("INPUT_DATA", "{}")
    mock_credentials = os.environ.get("MOCK_CREDENTIALS", "true").lower() == "true"
    output_file = os.environ.get("OUTPUT_FILE", "/tmp/automation_output.json")

    print(f"[Sandbox] Loading automation from: {automation_path}")
    print(f"[Sandbox] Mock credentials: {mock_credentials}")

    # Parse input data
    try:
        input_data = json.loads(input_data_str)
    except json.JSONDecodeError as e:
        print(f"[Sandbox] Error parsing INPUT_DATA: {e}", file=sys.stderr)
        write_output(
            output_file,
            {
                "success": False,
                "output_data": {},
                "error": f"Invalid INPUT_DATA JSON: {e}",
            },
        )
        return 1

    # Load automation
    try:
        automation = load_automation(automation_path)
        print(f"[Sandbox] Loaded automation: {type(automation).__name__}")
    except Exception as e:
        print(f"[Sandbox] Error loading automation: {e}", file=sys.stderr)
        traceback.print_exc()
        write_output(
            output_file,
            {
                "success": False,
                "output_data": {},
                "error": f"Failed to load automation: {e}",
            },
        )
        return 1

    # Run automation
    print("[Sandbox] Starting automation execution...")
    result = asyncio.run(run_automation(automation, input_data, mock_credentials))

    # Write output
    write_output(output_file, result)
    print(f"[Sandbox] Execution complete. Success: {result.get('success', False)}")

    return 0 if result.get("success", False) else 1


if __name__ == "__main__":
    sys.exit(main())
