import argparse
import asyncio
import json
import sys
import traceback
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, cast
from uuid import UUID

from uuid_extension import uuid7

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.runtime import (
    PipelineRunnerError,
    run_pipeline_from_file,
)
from alloy_runtime_types.dtos.pipeline import (
    RegisterLocalPipelineExecutionRequest,
    UpdateLocalPipelineExecutionRequest,
)


@dataclass(slots=True)
class LocalPipelineRunResult:
    pipeline_execution_id: UUID
    payload: dict[str, Any] | None
    waiting_for_approval: bool
    approval_key: str | None
    prompt: str | None


async def _execute_local_pipeline(
    *,
    module_path: Path,
    server_url: str,
    api_key: str,
    config: dict[str, Any],
    pipeline_name: str | None,
    pipeline_execution_id: UUID,
) -> dict[str, Any]:
    return await run_pipeline_from_file(
        module_path=module_path,
        server_url=server_url,
        api_key=api_key,
        config=config,
        pipeline_name=pipeline_name,
        execution_id=pipeline_execution_id,
    )


async def _register_local_execution(
    *,
    pipeline_execution_id: UUID,
    server_url: str,
    api_key: str,
    module_path: Path,
    pipeline_name: str | None,
) -> UUID:
    logger = get_or_configure_logger("alloy_runtime_sdk.pipeline.local_runner").bind(
        pipeline_execution_id=str(pipeline_execution_id)
    )
    async with ApiClient(base_url=server_url, api_key=api_key, logger=logger) as client:
        response = await client.register_local_pipeline_execution(
            RegisterLocalPipelineExecutionRequest(
                pipeline_execution_id=pipeline_execution_id,
                module_path=str(module_path),
                pipeline_name=pipeline_name,
            )
        )
    return response.pipeline_execution_id


async def _update_local_execution(
    *,
    pipeline_execution_id: UUID,
    server_url: str,
    api_key: str,
    status: Literal["running", "waiting_for_approval", "completed", "failed"],
    result_data: dict[str, Any] | None = None,
    error_message: str | None = None,
    error_stack_trace: str | None = None,
) -> None:
    logger = get_or_configure_logger("alloy_runtime_sdk.pipeline.local_runner").bind(
        pipeline_execution_id=str(pipeline_execution_id)
    )
    async with ApiClient(base_url=server_url, api_key=api_key, logger=logger) as client:
        await client.update_local_pipeline_execution(
            pipeline_execution_id,
            UpdateLocalPipelineExecutionRequest(
                status=status,
                result_data=result_data,
                error_message=error_message,
                error_stack_trace=error_stack_trace,
            ),
        )


async def run_local_pipeline(
    *,
    module_path: str,
    server_url: str,
    api_key: str,
    config: dict[str, Any] | None = None,
    pipeline_name: str | None = None,
    pipeline_execution_id: UUID | None = None,
) -> LocalPipelineRunResult:
    execution_id = pipeline_execution_id or uuid7()
    resolved_module_path = Path(module_path)

    if not resolved_module_path.exists():
        raise PipelineRunnerError(
            f"Pipeline module '{resolved_module_path}' does not exist"
        )

    execution_id = await _register_local_execution(
        pipeline_execution_id=execution_id,
        server_url=server_url,
        api_key=api_key,
        module_path=resolved_module_path,
        pipeline_name=pipeline_name,
    )

    try:
        result = await _execute_local_pipeline(
            module_path=resolved_module_path,
            server_url=server_url,
            api_key=api_key,
            config=config or {},
            pipeline_name=pipeline_name,
            pipeline_execution_id=execution_id,
        )
    except ApprovalPendingException as exc:
        await _update_local_execution(
            pipeline_execution_id=execution_id,
            server_url=server_url,
            api_key=api_key,
            status="waiting_for_approval",
            result_data={"approval_key": exc.approval_key, "prompt": exc.prompt},
        )
        return LocalPipelineRunResult(
            pipeline_execution_id=execution_id,
            payload=None,
            waiting_for_approval=True,
            approval_key=exc.approval_key,
            prompt=exc.prompt,
        )
    except PipelineRunnerError as exc:
        await _update_local_execution(
            pipeline_execution_id=execution_id,
            server_url=server_url,
            api_key=api_key,
            status="failed",
            error_message=str(exc),
        )
        raise
    except Exception as exc:
        await _update_local_execution(
            pipeline_execution_id=execution_id,
            server_url=server_url,
            api_key=api_key,
            status="failed",
            error_message=str(exc),
            error_stack_trace=traceback.format_exc(),
        )
        raise

    await _update_local_execution(
        pipeline_execution_id=execution_id,
        server_url=server_url,
        api_key=api_key,
        status="completed",
        result_data=result,
    )

    return LocalPipelineRunResult(
        pipeline_execution_id=execution_id,
        payload=result,
        waiting_for_approval=False,
        approval_key=None,
        prompt=None,
    )


def _parse_config(raw_config: str | None) -> dict[str, Any]:
    if raw_config is None:
        return {}

    config_text = raw_config
    if raw_config.startswith("@"):
        config_path = Path(raw_config[1:])
        config_text = config_path.read_text(encoding="utf-8")

    parsed = cast(object, json.loads(config_text))
    if not isinstance(parsed, dict):
        raise PipelineRunnerError("Local runner config must decode to a JSON object")
    typed_parsed = cast(dict[str, Any], parsed)
    return {key: value for key, value in typed_parsed.items()}


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="python -m alloy_runtime_sdk.pipeline.local_runner"
    )
    parser.add_argument("--file", required=True, dest="module_path")
    parser.add_argument("--function-name", dest="pipeline_name")
    parser.add_argument("--config")
    parser.add_argument("--server-url", required=True)
    parser.add_argument("--api-key", required=True)
    parser.add_argument("--pipeline-execution-id", type=UUID)
    return parser


def main(argv: list[str] | None = None) -> None:
    parser = _build_parser()
    args = parser.parse_args(argv)

    try:
        result = asyncio.run(
            run_local_pipeline(
                module_path=args.module_path,
                pipeline_name=args.pipeline_name,
                config=_parse_config(args.config),
                server_url=args.server_url,
                api_key=args.api_key,
                pipeline_execution_id=args.pipeline_execution_id,
            )
        )
        print(
            json.dumps(
                {
                    "success": True,
                    "pipeline_execution_id": str(result.pipeline_execution_id),
                    "waiting_for_approval": result.waiting_for_approval,
                    "approval_key": result.approval_key,
                    "prompt": result.prompt,
                    "result": result.payload,
                }
            )
        )
    except PipelineRunnerError as exc:
        print(json.dumps({"success": False, "error": str(exc)}))
        sys.exit(1)
    except Exception as exc:
        print(json.dumps({"success": False, "error": str(exc)}))
        sys.exit(1)


if __name__ == "__main__":
    main()
