from pathlib import Path
from typing import Any, Protocol, cast
from uuid import UUID

from alloy_runtime_sdk.pipeline.decorators import (
    clear_pipeline_registry,
    get_pipeline_registry,
)
from alloy_runtime_sdk.pipeline.http_context import HttpPipelineContext


class PipelineRunnerError(Exception):
    pass


class PipelineContextFactory(Protocol):
    def __call__(
        self,
        *,
        api_key: str,
        server_url: str,
        execution_id: UUID,
        config: dict[str, Any],
    ) -> HttpPipelineContext: ...


def load_pipeline_module(
    module_content: str, *, filename: str = "<pipeline_module>"
) -> None:
    clear_pipeline_registry()

    namespace: dict[str, Any] = {
        "__name__": "__pipeline_module__",
        "__file__": filename,
        "__builtins__": __builtins__,
    }

    try:
        exec(compile(module_content, filename, "exec"), namespace)
    except SyntaxError as exc:
        raise PipelineRunnerError(f"Syntax error in pipeline module: {exc}") from exc
    except Exception as exc:
        raise PipelineRunnerError(f"Failed to load pipeline module: {exc}") from exc

    registry = get_pipeline_registry()
    if not registry:
        raise PipelineRunnerError(
            "No @pipeline decorated functions found in module. "
            "Make sure your function is decorated with @pipeline."
        )


def get_pipeline_function(pipeline_name: str | None):
    registry = get_pipeline_registry()

    if pipeline_name:
        if pipeline_name not in registry:
            available = list(registry.keys())
            raise PipelineRunnerError(
                f"Pipeline '{pipeline_name}' not found. Available: {available}"
            )
        return registry[pipeline_name]

    if len(registry) == 1:
        return list(registry.values())[0]

    available = list(registry.keys())
    raise PipelineRunnerError(
        f"Multiple pipelines found: {available}. "
        f"Specify which registered pipeline name to run with pipeline_name."
    )


def serialize_pipeline_result(result: Any) -> dict[str, Any]:
    if result is None:
        return {"result": None}
    if isinstance(result, dict):
        return cast(dict[str, Any], result)
    if isinstance(result, list):
        return {"results": result}
    if hasattr(result, "model_dump"):
        dumped = result.model_dump(mode="json")
        if isinstance(dumped, dict):
            return cast(dict[str, Any], dumped)
        return {"result": dumped}
    return {"result": str(result)}


async def run_loaded_pipeline(
    *,
    api_key: str,
    server_url: str,
    execution_id: UUID,
    pipeline_func: Any,
    config: dict[str, Any],
    context_factory: PipelineContextFactory = HttpPipelineContext,
) -> dict[str, Any]:
    ctx = context_factory(
        api_key=api_key,
        server_url=server_url,
        execution_id=execution_id,
        config=config,
    )
    pipeline_name = getattr(
        pipeline_func,
        "_pipeline_name",
        getattr(pipeline_func, "__name__", "pipeline"),
    )

    try:
        ctx.logger.info(
            "pipeline_run_started",
            pipeline_name=pipeline_name,
            config_keys=sorted(config.keys()),
        )
        result = await pipeline_func(ctx)
        serialized_result = serialize_pipeline_result(result)
        ctx.logger.info(
            "pipeline_run_completed",
            pipeline_name=pipeline_name,
            result_keys=sorted(serialized_result.keys()),
        )
        return serialized_result
    except Exception as exc:
        ctx.logger.error(
            "pipeline_run_failed",
            pipeline_name=pipeline_name,
            error=str(exc),
            error_type=type(exc).__name__,
        )
        raise
    finally:
        await ctx.close()


async def run_pipeline(
    api_key: str,
    server_url: str,
    execution_id: UUID,
    module_content: str,
    pipeline_name: str | None,
    config: dict[str, Any],
    *,
    filename: str = "<pipeline_module>",
    context_factory: PipelineContextFactory = HttpPipelineContext,
) -> dict[str, Any]:
    load_pipeline_module(module_content, filename=filename)
    pipeline_func = get_pipeline_function(pipeline_name)
    return await run_loaded_pipeline(
        api_key=api_key,
        server_url=server_url,
        execution_id=execution_id,
        pipeline_func=pipeline_func,
        config=config,
        context_factory=context_factory,
    )


async def run_pipeline_from_file(
    *,
    module_path: Path,
    api_key: str,
    server_url: str,
    execution_id: UUID,
    pipeline_name: str | None,
    config: dict[str, Any],
    context_factory: PipelineContextFactory = HttpPipelineContext,
) -> dict[str, Any]:
    if not module_path.exists():
        raise PipelineRunnerError(f"Pipeline module '{module_path}' does not exist")
    if not module_path.is_file():
        raise PipelineRunnerError(f"Pipeline module '{module_path}' is not a file")

    module_content = module_path.read_text(encoding="utf-8")
    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,
        filename=str(module_path),
        context_factory=context_factory,
    )
