"""Credential resolution stage for resolving credential references."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from agent_flows.exceptions import CredentialError
from agent_flows.utils.logging import get_logger

if TYPE_CHECKING:
    from agent_flows.core.flow_engine.pipeline import ExecutionState
    from agent_flows.models.credentials import CredentialBundle
    from agent_flows.services.realtimex import CredentialResolver

logger = get_logger(__name__)


class CredentialResolutionStage:
    """Resolves credential references in step configuration.

    Runs: Pre-execution only
    Modifies: Context resolved_credentials

    Credentials are defined in step config as:
    {
        "credentials": {
            "alias_name": {"id": "credential_id"}
        }
    }

    Resolved credentials are added to context for executor use.
    """

    def __init__(self, credential_resolver: CredentialResolver | None = None) -> None:
        """Initialize credential resolution stage.

        Args:
            credential_resolver: Optional resolver for fetching credentials from RTX.
        """
        self._credential_resolver = credential_resolver

    async def process(self, state: ExecutionState) -> ExecutionState:
        """Resolve credential references in step configuration.

        Args:
            state: Current execution state

        Returns:
            New state with resolved credentials in context
        """
        # Only run before execution
        if state.is_post_execution:
            return state

        # Skip if no credential manager
        if not self._credential_resolver:
            return state

        # Resolve credentials
        resolved = await self._resolve_credentials(state.step.config)

        # Skip if no credentials to resolve
        if not resolved:
            return state

        # Update context with resolved credentials
        new_context = state.context.with_credentials(resolved)
        return state.with_context(new_context)

    async def _resolve_credentials(self, config: dict[str, Any]) -> dict[str, CredentialBundle]:
        """Resolve credential references defined in the step configuration.

        Args:
            config: Step configuration

        Returns:
            Dictionary mapping alias to resolved credential bundle
        """
        raw_credentials = config.get("credentials")
        if not isinstance(raw_credentials, dict):
            return {}

        resolved: dict[str, CredentialBundle] = {}
        for alias, reference in raw_credentials.items():
            if not isinstance(reference, dict):
                logger.warning(
                    "Skipping credential reference with invalid structure",
                    extra={
                        "alias": alias,
                        "reference_type": type(reference).__name__,
                    },
                )
                continue

            credential_id = reference.get("id")
            if not credential_id:
                raise CredentialError(
                    f"Credential reference missing id for alias '{alias}'",
                    details={"alias": alias},
                )

            bundle = await self._credential_resolver.get(str(credential_id))
            resolved[str(alias)] = bundle

        return resolved
