"""PipelineContext protocol - defines the primitives users can call.

The context implementation provides these methods for pipeline code.
Users write pipelines using these methods, with config values from runtime.
"""

from collections.abc import Awaitable
from typing import Any, Protocol
from uuid import UUID

from alloy_runtime_sdk.logging.protocol import PipelineLogger
from alloy_runtime_sdk.pipeline.types import (
    ApprovalResult,
    GenerateTextResult,
    GenerateTextTurnResult,
)
from alloy_runtime_types.dtos.managed_tools import ManagedToolRequest


class PipelineContext(Protocol):
    """Runtime context for pipeline execution.

    This protocol defines the primitives (atoms) that users compose
    to build pipelines. The implementation (HttpPipelineContext) executes
    these operations via HTTP calls to the server.

    Attributes:
        config: Runtime configuration passed at execution time.
                Users access this like template variables: config.synthesis.agents
        execution_id: Unique identifier for this pipeline execution.

    Example:
        @pipeline
        async def my_pipeline(ctx: PipelineContext):
            config = ctx.config

            # Use primitives with config values
            result = await ctx.generate_text(
                step_id="main",
                agent=config.agent,
                user_message_template=config.template,
            )

            return result
    """

    @property
    def config(self) -> Any:
        """Runtime configuration passed at execution.

        This is the config dict from the execute request, converted
        to a namespace object for dot-access: config.synthesis.agents

        Users define what config their pipeline expects, similar to
        template variables.
        """
        ...

    @property
    def execution_id(self) -> UUID:
        """Unique identifier for this pipeline execution."""
        ...

    @property
    def logger(self) -> PipelineLogger: ...

    async def generate_text(
        self,
        *,
        step_id: str,
        # Model source - one of agent OR model_source required
        agent: str | UUID | None = None,
        model_source: dict[str, Any] | None = None,
        # User message - one required
        user_message: str | None = None,
        user_message_template: str | UUID | None = None,
        user_message_variables: dict[str, Any] | None = None,
        # System instruction - optional
        system_instruction: str | None = None,
        system_instruction_template: str | UUID | None = None,
        system_instruction_variables: dict[str, Any] | None = None,
        # Output schema - optional
        output_schema: str | UUID | None = None,
        managed_tools: list[ManagedToolRequest] | None = None,
        include_tool_call_details: bool = False,
        # Tags - optional
        tags: list[str] | None = None,
        chat_session_id: UUID | None = None,
    ) -> GenerateTextResult:
        """Generate text using an agent or direct model.

        This is the pipeline-specific AI generation primitive.
        `chat_session_id` continues an existing conversation, while `step_id`
        remains the caller-owned immutable identifier for this pipeline step.

        Args:
            step_id: Stable identifier for this step. Used to generate
                     deterministic idempotency keys for pipeline replay.
                     Every pipeline-origin generate_text call must provide it.
            agent: Agent identifier (name or UUID). Mutually exclusive with model_source.
            model_source: Direct model config dict with provider_key and provider_model_name.
                         Mutually exclusive with agent.
            user_message: Direct user message text.
            user_message_template: Template identifier for user message.
            user_message_variables: Variables for the user message template.
            system_instruction: Direct system instruction text.
            system_instruction_template: Template identifier for system instruction.
            system_instruction_variables: Variables for the system instruction template.
            output_schema: Schema identifier for structured output.
            managed_tools: Request-level managed tools for direct model generation.
            include_tool_call_details: Whether to include tool call details in response metadata.
            tags: Tags to apply to generated content.
            chat_session_id: Existing chat session to continue for a follow-up turn.

        Returns:
            GenerateTextResult with output_text, chat_session_id,
            structured_output, usage, and cost.

        Example:
            # Using an agent
            result = await ctx.generate_text(
                step_id="research",
                agent="researcher",
                user_message_template="research-task",
                user_message_variables={"topic": config.topic},
            )

            # Using direct model
            result = await ctx.generate_text(
                step_id="hello",
                model_source={"provider_key": "openai", "provider_model_name": "gpt-4o"},
                user_message="Hello!",
            )
        """
        ...

    async def generate_text_turn(
        self,
        *,
        step_namespace: str,
        chat_session_id: UUID,
        agent: str | UUID | None = None,
        model_source: dict[str, Any] | None = None,
        user_message: str | None = None,
        user_message_template: str | UUID | None = None,
        user_message_variables: dict[str, Any] | None = None,
        system_instruction: str | None = None,
        system_instruction_template: str | UUID | None = None,
        system_instruction_variables: dict[str, Any] | None = None,
        output_schema: str | UUID | None = None,
        managed_tools: list[ManagedToolRequest] | None = None,
        include_tool_call_details: bool = False,
        tags: list[str] | None = None,
    ) -> GenerateTextTurnResult: ...

    async def parallel(
        self,
        coroutines: list[Awaitable[Any]],
        *,
        max_concurrency: int = 10,
    ) -> list[Any]:
        """Execute multiple operations concurrently.

        Runs the provided coroutines in parallel with a semaphore limiting
        concurrent execution. Results are returned in the same order as
        the input coroutines.

        Args:
            coroutines: List of awaitable operations (e.g., generate_text calls).
            max_concurrency: Maximum number of concurrent operations.

        Returns:
            List of results in the same order as input coroutines.

        Example:
            # Run multiple agents in parallel
            results = await ctx.parallel([
                ctx.generate_text(
                    step_id=f"agent-{index}",
                    agent=agent,
                    user_message_template="task",
                )
                for index, agent in enumerate(config.agents)
            ])

            # Cartesian product: test each prompt against each model
            results = await ctx.parallel([
                ctx.generate_text(
                    step_id=f"prompt-{prompt_index}-model-{model_index}",
                    model_source=model,
                    system_instruction=prompt.output_text,
                    user_message=config.test_input,
                )
                for prompt_index, prompt in enumerate(synthesis_results)
                for model_index, model in enumerate(config.test_models)
            ])
        """
        ...

    async def wait_for_approval(
        self,
        prompt: str | None = None,
        *,
        approval_id: str | None = None,
    ) -> ApprovalResult:
        """Pause execution until human approval is received.

        Creates a durable approval checkpoint. On first encounter, the pipeline
        pauses. When a human decides via the API and the pipeline is re-run,
        decided checkpoints return immediately.

        Args:
            prompt: Message displayed to the approver.
            approval_id: Stable identifier for this approval point. If omitted,
                        auto-generated from an incrementing counter. Explicit IDs
                        are recommended for pipelines with dynamic control flow.

        Returns:
            ApprovalResult with the decision and any feedback.

        Raises:
            ApprovalPendingException: Internally raised to pause execution.
                                     Caught by runtime to checkpoint state.
        """
        ...
