"""Pipeline runner for Modal sandbox execution.

This module is the entry point for running pipelines inside Modal sandboxes.
It reads environment variables for configuration and outputs results as JSON.

Usage from Modal:
    python -m alloy_runtime_sdk.pipeline.runner

Environment Variables:
    ALLOY_RUNTIME_API_KEY: API key for authenticating with the server
    ALLOY_RUNTIME_SERVER_URL: Base URL of the Alloy Runtime server
    ALLOY_RUNTIME_EXECUTION_ID: UUID of this pipeline execution
    ALLOY_RUNTIME_MODULE_PATH: File path to the Python source code
    ALLOY_RUNTIME_PIPELINE_NAME: Name of the @pipeline function to run (optional)
    ALLOY_RUNTIME_CONFIG: JSON-encoded runtime configuration
"""

import asyncio
import contextlib
import io
import json
import os
import sys
import traceback
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator, Literal, TextIO, cast
from uuid import UUID

from alloy_runtime_sdk.api_client.client import ApiClient
from alloy_runtime_sdk.logging.config import get_or_configure_logger
from alloy_runtime_sdk.pipeline.exceptions import ApprovalPendingException
from alloy_runtime_sdk.pipeline.http_context import HttpPipelineContext
from alloy_runtime_sdk.pipeline.runtime import (
    PipelineRunnerError,
    get_pipeline_function,
    load_pipeline_module,
    run_pipeline as _run_pipeline,
    run_pipeline_from_file,
)

HostedPipelineExecutionStatus = Literal[
    "running", "waiting_for_approval", "completed", "failed"
]

__all__ = [
    "PipelineRunnerError",
    "get_pipeline_function",
    "load_pipeline_module",
    "run_pipeline",
    "main",
]


class _BufferedTee:
    def __init__(self, sink: TextIO, buffer: io.StringIO) -> None:
        self._sink = sink
        self._buffer = buffer

    def write(self, s: str) -> int:
        self._buffer.write(s)
        return self._sink.write(s)

    def flush(self) -> None:
        self._buffer.flush()
        self._sink.flush()


@dataclass
class _CapturedPipelineOutput:
    stdout: io.StringIO
    stderr: io.StringIO


@dataclass
class _HostedExecutionOutcome:
    status: HostedPipelineExecutionStatus
    result_data: dict[str, Any] | None
    error_message: str | None
    error_stack_trace: str | None
    response_payload: dict[str, Any]
    exit_code: int


@contextlib.contextmanager
def _capture_pipeline_output() -> Iterator[_CapturedPipelineOutput]:
    stdout_buffer = io.StringIO()
    stderr_buffer = io.StringIO()
    live_stderr = cast(TextIO, sys.stderr)
    stdout_tee = _BufferedTee(live_stderr, stdout_buffer)
    stderr_tee = _BufferedTee(live_stderr, stderr_buffer)
    with (
        contextlib.redirect_stdout(cast(TextIO, stdout_tee)),
        contextlib.redirect_stderr(cast(TextIO, stderr_tee)),
    ):
        yield _CapturedPipelineOutput(stdout=stdout_buffer, stderr=stderr_buffer)


async def _notify_hosted_execution_update(
    *,
    finalization_token: str,
    server_url: str,
    execution_id: UUID,
    status: HostedPipelineExecutionStatus,
    result_data: dict[str, Any] | None = None,
    error_message: str | None = None,
    error_stack_trace: str | None = None,
    execution_stdout: str | None = None,
    execution_stderr: str | None = None,
    orchestration_state: str | None = None,
) -> None:
    logger = get_or_configure_logger("alloy_runtime_sdk.pipeline.runner").bind(
        execution_id=str(execution_id)
    )
    async with ApiClient(
        base_url=server_url, api_key=finalization_token, logger=logger
    ) as client:
        await client.finalize_hosted_pipeline_execution(
            execution_id,
            status=status,
            result_data=result_data,
            error_message=error_message,
            error_stack_trace=error_stack_trace,
            execution_stdout=execution_stdout,
            execution_stderr=execution_stderr,
            orchestration_state=orchestration_state,
        )


async def _emit_hosted_heartbeats(
    *,
    finalization_token: str,
    server_url: str,
    execution_id: UUID,
    interval_seconds: int,
    stop_event: asyncio.Event,
) -> None:
    first_update = True
    while not stop_event.is_set():
        try:
            await _notify_hosted_execution_update(
                finalization_token=finalization_token,
                server_url=server_url,
                execution_id=execution_id,
                status="running",
                orchestration_state="worker_started" if first_update else "running",
            )
        except Exception as exc:
            print(f"Hosted execution heartbeat failed: {exc}", file=sys.stderr)
        first_update = False

        try:
            await asyncio.wait_for(stop_event.wait(), timeout=interval_seconds)
        except TimeoutError:
            continue


async def run_pipeline(
    api_key: str,
    server_url: str,
    execution_id: UUID,
    module_content: str,
    pipeline_name: str | None,
    config: dict[str, Any],
) -> dict[str, Any]:
    """Run a pipeline with the given configuration.

    Args:
        api_key: API key for server authentication.
        server_url: Base URL of the Alloy Runtime server.
        execution_id: UUID of this pipeline execution.
        module_content: Python source code of the pipeline.
        pipeline_name: Name of the @pipeline function to run.
        config: Runtime configuration dict.

    Returns:
        The pipeline result as a dict.

    Raises:
        PipelineRunnerError: If execution fails.
        ApprovalPendingException: If pipeline pauses for approval.
    """
    return await _run_pipeline(
        api_key=api_key,
        server_url=server_url,
        execution_id=execution_id,
        module_content=module_content,
        pipeline_name=pipeline_name,
        config=config,
        context_factory=HttpPipelineContext,
    )


async def _run_pipeline_with_callbacks(
    *,
    api_key: str,
    finalization_token: str | None,
    server_url: str,
    execution_id: UUID,
    module_path: Path,
    pipeline_name: str | None,
    config: dict[str, Any],
    heartbeat_interval_seconds: int,
) -> _HostedExecutionOutcome:
    stop_event = asyncio.Event()
    heartbeat_task: asyncio.Task[None] | None = None
    if finalization_token:
        heartbeat_task = asyncio.create_task(
            _emit_hosted_heartbeats(
                finalization_token=finalization_token,
                server_url=server_url,
                execution_id=execution_id,
                interval_seconds=max(heartbeat_interval_seconds, 1),
                stop_event=stop_event,
            )
        )

    try:
        result = await run_pipeline_from_file(
            module_path=module_path,
            api_key=api_key,
            server_url=server_url,
            execution_id=execution_id,
            pipeline_name=pipeline_name,
            config=config,
        )
        return _HostedExecutionOutcome(
            status="completed",
            result_data=result,
            error_message=None,
            error_stack_trace=None,
            response_payload={"success": True, "result": result},
            exit_code=0,
        )
    except ApprovalPendingException as exc:
        return _HostedExecutionOutcome(
            status="waiting_for_approval",
            result_data={"approval_key": exc.approval_key, "prompt": exc.prompt},
            error_message=None,
            error_stack_trace=None,
            response_payload={
                "success": True,
                "waiting_for_approval": True,
                "approval_key": exc.approval_key,
                "prompt": exc.prompt,
            },
            exit_code=0,
        )
    except PipelineRunnerError as exc:
        return _HostedExecutionOutcome(
            status="failed",
            result_data=None,
            error_message=str(exc),
            error_stack_trace=None,
            response_payload={"success": False, "error": str(exc)},
            exit_code=1,
        )
    except Exception as exc:
        return _HostedExecutionOutcome(
            status="failed",
            result_data=None,
            error_message=str(exc),
            error_stack_trace=traceback.format_exc(),
            response_payload={"success": False, "error": str(exc)},
            exit_code=1,
        )
    finally:
        if heartbeat_task is not None:
            stop_event.set()
            with contextlib.suppress(Exception):
                await heartbeat_task


def main() -> None:
    """Main entry point for Modal sandbox execution.

    Reads configuration from environment variables, runs the pipeline,
    and outputs the result as JSON to stdout.
    """
    api_key = os.environ.get("ALLOY_RUNTIME_API_KEY")
    finalization_token = os.environ.get("ALLOY_RUNTIME_FINALIZATION_TOKEN")
    if not api_key:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": "ALLOY_RUNTIME_API_KEY environment variable is required",
                }
            )
        )
        sys.exit(1)

    server_url = os.environ.get("ALLOY_RUNTIME_SERVER_URL")
    if not server_url:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": "ALLOY_RUNTIME_SERVER_URL environment variable is required",
                }
            )
        )
        sys.exit(1)

    execution_id_str = os.environ.get("ALLOY_RUNTIME_EXECUTION_ID")
    if not execution_id_str:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": "ALLOY_RUNTIME_EXECUTION_ID environment variable is required",
                }
            )
        )
        sys.exit(1)

    try:
        execution_id = UUID(execution_id_str)
    except ValueError:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": f"Invalid execution ID: {execution_id_str}",
                }
            )
        )
        sys.exit(1)

    module_path_str = os.environ.get("ALLOY_RUNTIME_MODULE_PATH")
    if not module_path_str:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": "ALLOY_RUNTIME_MODULE_PATH environment variable is required",
                }
            )
        )
        sys.exit(1)

    pipeline_name = os.environ.get("ALLOY_RUNTIME_PIPELINE_NAME")
    heartbeat_interval_seconds = int(
        os.environ.get("ALLOY_RUNTIME_HEARTBEAT_INTERVAL_SECONDS", "15")
    )

    config_json = os.environ.get("ALLOY_RUNTIME_CONFIG", "{}")
    try:
        config = json.loads(config_json)
    except json.JSONDecodeError as exc:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": f"Invalid config JSON: {exc}",
                }
            )
        )
        sys.exit(1)

    with _capture_pipeline_output() as captured_output:
        outcome = asyncio.run(
            _run_pipeline_with_callbacks(
                api_key=api_key,
                finalization_token=finalization_token,
                server_url=server_url,
                execution_id=execution_id,
                module_path=Path(module_path_str),
                pipeline_name=pipeline_name,
                config=config,
                heartbeat_interval_seconds=heartbeat_interval_seconds,
            )
        )

    captured_stdout = captured_output.stdout.getvalue()
    captured_stderr = captured_output.stderr.getvalue()
    if finalization_token:
        try:
            asyncio.run(
                _notify_hosted_execution_update(
                    finalization_token=finalization_token,
                    server_url=server_url,
                    execution_id=execution_id,
                    status=outcome.status,
                    result_data=outcome.result_data,
                    error_message=outcome.error_message,
                    error_stack_trace=outcome.error_stack_trace,
                    execution_stdout=captured_stdout,
                    execution_stderr=captured_stderr,
                    orchestration_state=outcome.status,
                )
            )
        except Exception as callback_error:
            print(f"Hosted execution update failed: {callback_error}", file=sys.stderr)

    print(json.dumps(outcome.response_payload))
    if outcome.exit_code != 0:
        sys.exit(outcome.exit_code)


if __name__ == "__main__":
    main()
