"""HttpPipelineContext - HTTP-based implementation of PipelineContext.

This context implementation uses the ApiClient to make HTTP calls
to the Alloy Runtime server. It runs inside Modal sandboxes.
"""

import re

from collections.abc import Awaitable
from typing import Any, 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.logging.protocol import PipelineLogger
from alloy_runtime_types.dtos.generation import (
    AgentSource,
    DirectModelSource,
    GenerateTextRequest,
    GenerationParameters,
    SchemaReference,
)
from alloy_runtime_types.dtos.managed_tools import ManagedToolRequest
from alloy_runtime_types.enums.provider import Provider

from alloy_runtime_sdk.pipeline.exceptions import ApprovalPendingException
from alloy_runtime_sdk.pipeline.types import (
    ApprovalResult,
    CostBreakdown,
    GenerateTextResult,
    GenerateTextTurnResult,
    TokenUsage,
)


_ID_COMPONENT_PATTERN = re.compile(r"^[a-zA-Z0-9_.-]+$")


class ConfigNamespace:
    """Converts a dict to a namespace object for dot-access.

    Allows users to access config like: config.synthesis.agents
    instead of: config["synthesis"]["agents"]
    """

    def __init__(self, data: dict[str, Any]):
        for key, value in data.items():
            if isinstance(value, dict):
                setattr(self, key, ConfigNamespace(cast(dict[str, Any], value)))
            else:
                setattr(self, key, value)

    def __repr__(self) -> str:
        return f"ConfigNamespace({vars(self)})"


class HttpPipelineContext:
    """HTTP-based implementation of PipelineContext.

    This context runs inside Modal sandboxes and uses the ApiClient
    to make HTTP calls back to the Alloy Runtime server.

    Attributes:
        config: Runtime configuration converted to dot-access namespace.
        execution_id: UUID of this pipeline execution.
    """

    def __init__(
        self,
        *,
        api_key: str,
        server_url: str,
        execution_id: UUID,
        config: dict[str, Any],
    ):
        self._api_key = api_key
        self._server_url = server_url
        self._execution_id = execution_id
        self._config = ConfigNamespace(config) if config else ConfigNamespace({})
        bound_logger = get_or_configure_logger("alloy_runtime_sdk.pipeline").bind(
            execution_id=str(execution_id)
        )
        self._logger = PipelineLogger(bound_logger)
        self._client = ApiClient(
            base_url=server_url, api_key=api_key, logger=self._logger
        )
        self._used_step_ids: set[str] = set()
        self._step_turn_counters: dict[str, int] = {}
        self._approval_counter: int = 0
        self._used_approval_ids: set[str] = set()

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

        Converted to namespace for dot-access: config.synthesis.agents
        """
        return self._config

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

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

    async def generate_text(
        self,
        *,
        step_id: str,
        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,
        chat_session_id: UUID | None = None,
    ) -> GenerateTextResult:
        """Generate text using an agent or direct model.

        Makes an HTTP request to the server's generate_text endpoint.
        """
        request = self._build_generate_text_request(
            step_id=step_id,
            agent=agent,
            model_source=model_source,
            user_message=user_message,
            user_message_template=user_message_template,
            user_message_variables=user_message_variables,
            system_instruction=system_instruction,
            system_instruction_template=system_instruction_template,
            system_instruction_variables=system_instruction_variables,
            output_schema=output_schema,
            managed_tools=managed_tools,
            include_tool_call_details=include_tool_call_details,
            tags=tags,
            chat_session_id=chat_session_id,
        )

        response = await self._client.generate_text(request)

        return GenerateTextResult(
            execution_id=response.execution_id,
            chat_session_id=response.chat_session_id,
            output_text=response.output_text,
            structured_output=response.structured_output,
            usage=TokenUsage(
                input_tokens=response.usage.input_tokens,
                output_tokens=response.usage.output_tokens,
                cache_creation_tokens=response.usage.cache_creation_tokens,
                cache_read_tokens=response.usage.cache_read_tokens,
            ),
            cost=CostBreakdown(
                input_cost_usd=response.cost.input_cost_usd,
                output_cost_usd=response.cost.output_cost_usd,
                cache_cost_usd=response.cost.cache_cost_usd,
                total_cost_usd=response.cost.total_cost_usd,
            ),
            provider_key=response.provider_key,
            provider_model_name=response.provider_model_name,
            finish_reason=response.finish_reason,
            metadata=response.metadata,
        )

    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:
        step_id = self._mint_step_turn_id(step_namespace)
        result = await self.generate_text(
            step_id=step_id,
            agent=agent,
            model_source=model_source,
            user_message=user_message,
            user_message_template=user_message_template,
            user_message_variables=user_message_variables,
            system_instruction=system_instruction,
            system_instruction_template=system_instruction_template,
            system_instruction_variables=system_instruction_variables,
            output_schema=output_schema,
            managed_tools=managed_tools,
            include_tool_call_details=include_tool_call_details,
            tags=tags,
            chat_session_id=chat_session_id,
        )
        return GenerateTextTurnResult(step_id=step_id, result=result)

    def _build_generate_text_request(
        self,
        *,
        step_id: str,
        agent: str | UUID | None,
        model_source: dict[str, Any] | None,
        user_message: str | None,
        user_message_template: str | UUID | None,
        user_message_variables: dict[str, Any] | None,
        system_instruction: str | None,
        system_instruction_template: str | UUID | None,
        system_instruction_variables: dict[str, Any] | None,
        output_schema: str | UUID | None,
        managed_tools: list[ManagedToolRequest] | None,
        include_tool_call_details: bool,
        tags: list[str] | None,
        chat_session_id: UUID | None,
    ) -> GenerateTextRequest:
        self._claim_step_id(step_id)

        request_data: dict[str, Any] = {
            "model_source": self._build_model_source(
                agent=agent, model_source=model_source
            ),
            "user_message": user_message,
            "user_message_template": user_message_template,
            "user_message_variables": user_message_variables,
            "system_instruction": system_instruction,
            "system_instruction_template": system_instruction_template,
            "system_instruction_variables": system_instruction_variables,
            "managed_tools": managed_tools,
            "output_schema": SchemaReference(schema_id=output_schema)
            if output_schema
            else None,
            "tags": tags or [],
            "stream": False,
            "external_id": f"pe:{self._execution_id}:step:{step_id}",
            "pipeline_execution_id": self._execution_id,
            "execution_purpose": "pipeline_step",
            "chat_session_id": chat_session_id,
            "include_tool_call_details": include_tool_call_details,
        }
        return GenerateTextRequest(**request_data)

    def _mint_step_turn_id(self, step_namespace: str) -> str:
        self._validate_id_component(value=step_namespace, label="step_namespace")
        turn_index = self._step_turn_counters.get(step_namespace, 0)
        self._step_turn_counters[step_namespace] = turn_index + 1
        return f"{step_namespace}.turn_{turn_index}"

    def _claim_step_id(self, step_id: str) -> None:
        if step_id in self._used_step_ids:
            raise ValueError(
                f"Duplicate step_id '{step_id}' in pipeline execution. "
                "Each generate_text() call must have a unique step_id."
            )

        self._validate_id_component(value=step_id, label="step_id")

        self._used_step_ids.add(step_id)

    def _validate_id_component(self, *, value: str, label: str) -> None:
        if not _ID_COMPONENT_PATTERN.match(value):
            raise ValueError(
                f"{label} '{value}' contains invalid characters. "
                "Only alphanumeric characters, hyphens, underscores, and dots are allowed."
            )

    def _build_model_source(
        self,
        *,
        agent: str | UUID | None,
        model_source: dict[str, Any] | None,
    ) -> AgentSource | DirectModelSource:
        if agent is not None:
            return AgentSource(agent=agent)

        if model_source is None:
            raise ValueError("Either agent or model_source must be provided")

        return DirectModelSource(
            provider_key=Provider(model_source["provider_key"]),
            provider_model_name=model_source["provider_model_name"],
            generation_params=GenerationParameters(
                **model_source.get("generation_params", {})
            )
            if model_source.get("generation_params")
            else None,
            request_headers=model_source.get("request_headers"),
        )

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

        Uses asyncio.gather with a semaphore to limit concurrency.
        Results are returned in the same order as input coroutines.
        """
        from alloy_runtime_sdk.concurrency import parallel as _parallel

        return await _parallel(coroutines, max_concurrency=max_concurrency)

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

        On first encounter, creates a pending checkpoint and raises
        ApprovalPendingException. On re-run, if the checkpoint has been
        decided, returns the ApprovalResult immediately.
        """
        if approval_id is None:
            approval_id = f"auto_{self._approval_counter}"
            self._approval_counter += 1
            while approval_id in self._used_approval_ids:
                approval_id = f"auto_{self._approval_counter}"
                self._approval_counter += 1
        elif approval_id in self._used_approval_ids:
            raise ValueError(
                f"Duplicate approval_id '{approval_id}' in pipeline execution"
            )

        self._validate_id_component(value=approval_id, label="approval_id")

        self._used_approval_ids.add(approval_id)
        approval_key = f"pe:{self._execution_id}:approval:{approval_id}"

        decision = await self._client.get_approval_checkpoint(
            self._execution_id, approval_key
        )
        if decision and decision.status != "pending":
            return ApprovalResult(
                approved=bool(decision.approved),
                decision_data=decision.decision_data or {},
                approved_by_user_id=decision.decided_by_user_id,
                decided_at=decision.decided_at,
            )

        if decision is None:
            await self._client.create_approval_checkpoint(
                self._execution_id, approval_key
            )

        raise ApprovalPendingException(approval_key=approval_key, prompt=prompt)

    async def close(self) -> None:
        """Close the HTTP client."""
        await self._client.close()
