"""
Base provider abstractions for foundry-mcp.

This module provides the core provider contracts adapted from the Foundry CLI,
enabling pluggable LLM backends for CLI operations. The abstractions support
capability negotiation, request/response normalization, and lifecycle hooks.

Design principles:
- Frozen dataclasses for immutability
- Enum-based capabilities for type-safe routing
- Status codes aligned with existing ProviderStatus patterns
- Error hierarchy for granular exception handling
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Sequence, Set

if TYPE_CHECKING:
    from foundry_mcp.core.responses import ErrorType


class ProviderCapability(Enum):
    """
    Feature flags a provider can expose to routing heuristics.

    These capabilities enable callers to select providers based on
    required features (vision, streaming, etc.) and allow registries
    to route requests to appropriate backends.

    Values:
        TEXT: Basic text generation capability
        VISION: Image/vision input processing
        FUNCTION_CALLING: Tool/function invocation support
        STREAMING: Incremental response streaming
        THINKING: Extended reasoning/chain-of-thought support
    """

    TEXT = "text_generation"
    VISION = "vision"
    FUNCTION_CALLING = "function_calling"
    STREAMING = "streaming"
    THINKING = "thinking"


class ProviderStatus(Enum):
    """
    Normalized execution outcomes emitted by providers.

    These status codes provide a consistent interface for handling provider
    responses across different backends, enabling uniform error handling
    and retry logic.

    Values:
        SUCCESS: Operation completed successfully
        TIMEOUT: Operation exceeded time limit (retryable)
        NOT_FOUND: Provider/resource not available (not retryable)
        INVALID_OUTPUT: Provider returned malformed response (not retryable)
        ERROR: Generic error during execution (retryable)
        CANCELED: Operation was explicitly canceled (not retryable)
    """

    SUCCESS = "success"
    TIMEOUT = "timeout"
    NOT_FOUND = "not_found"
    INVALID_OUTPUT = "invalid_output"
    ERROR = "error"
    CANCELED = "canceled"

    def is_retryable(self) -> bool:
        """
        Check if this status represents a transient failure that may succeed on retry.

        Retryable statuses:
            - TIMEOUT: May succeed with more time or if temporary resource contention resolves
            - ERROR: Generic errors may be transient (network issues, rate limits, etc.)

        Non-retryable statuses:
            - SUCCESS: Already succeeded
            - NOT_FOUND: Provider/tool not available (won't change on retry)
            - INVALID_OUTPUT: Provider responded but output was malformed
            - CANCELED: Explicitly canceled (shouldn't retry)

        Returns:
            True if the status is retryable, False otherwise
        """
        return self in (ProviderStatus.TIMEOUT, ProviderStatus.ERROR)

    def to_error_type(self) -> "ErrorType":
        """
        Map provider status to foundry-mcp ErrorType for response envelopes.

        This enables consistent error categorization across MCP and CLI surfaces,
        allowing callers to handle provider errors using the standard error taxonomy.

        Mapping:
            - SUCCESS: raises ValueError (not an error state)
            - TIMEOUT: UNAVAILABLE (503 analog, retryable)
            - NOT_FOUND: NOT_FOUND (404 analog, not retryable)
            - INVALID_OUTPUT: VALIDATION (400 analog, not retryable)
            - ERROR: INTERNAL (500 analog, retryable)
            - CANCELED: INTERNAL (operation aborted)

        Returns:
            ErrorType enum value corresponding to this status

        Raises:
            ValueError: If called on SUCCESS status (not an error)
        """
        from foundry_mcp.core.responses import ErrorType

        if self == ProviderStatus.SUCCESS:
            raise ValueError("SUCCESS status cannot be mapped to an error type")

        mapping = {
            ProviderStatus.TIMEOUT: ErrorType.UNAVAILABLE,
            ProviderStatus.NOT_FOUND: ErrorType.NOT_FOUND,
            ProviderStatus.INVALID_OUTPUT: ErrorType.VALIDATION,
            ProviderStatus.ERROR: ErrorType.INTERNAL,
            ProviderStatus.CANCELED: ErrorType.INTERNAL,
        }
        return mapping[self]


@dataclass(frozen=True)
class ProviderRequest:
    """
    Normalized request payload for provider execution.

    This dataclass encapsulates all parameters needed to make a generation
    request to any provider backend. Fields follow common LLM API conventions
    to ensure portability across different providers.

    Attributes:
        prompt: The user's input prompt/message
        system_prompt: Optional system/instruction prompt
        model: Model identifier (provider-specific, e.g., "pro", "flash")
        timeout: Request timeout in seconds (None = provider default)
        temperature: Sampling temperature (0.0-2.0, None = provider default)
        max_tokens: Maximum output tokens (None = provider default)
        metadata: Arbitrary request metadata (tracing IDs, feature flags, etc.)
        stream: Whether to request streaming response
        attachments: File paths or URIs for multimodal inputs
    """

    prompt: str
    system_prompt: Optional[str] = None
    model: Optional[str] = None
    timeout: Optional[float] = None
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    stream: bool = False
    attachments: Sequence[str] = field(default_factory=list)


@dataclass(frozen=True)
class TokenUsage:
    """
    Token accounting information reported by providers.

    Tracks input, output, and cached tokens for cost estimation
    and usage monitoring.

    Attributes:
        input_tokens: Tokens consumed by the prompt
        output_tokens: Tokens generated in the response
        cached_input_tokens: Tokens served from cache (if supported)
        total_tokens: Sum of all token counts
    """

    input_tokens: int = 0
    output_tokens: int = 0
    cached_input_tokens: int = 0
    total_tokens: int = 0


@dataclass(frozen=True)
class ProviderResult:
    """
    Normalized provider response.

    This dataclass encapsulates all data returned from a provider execution,
    providing a consistent interface regardless of the backend used.

    Attributes:
        content: Final text output (aggregated if streaming was used)
        provider_id: Canonical provider identifier (e.g., "gemini", "codex")
        model_used: Fully-qualified model identifier (e.g., "gemini:pro")
        status: ProviderStatus describing execution outcome
        tokens: Token usage data (if reported by provider)
        duration_ms: Execution duration in milliseconds
        stderr: Captured stderr/log output for debugging
        raw_payload: Provider-specific metadata (traces, debug info, etc.)
    """

    content: str
    provider_id: str
    model_used: str
    status: ProviderStatus
    tokens: TokenUsage = field(default_factory=TokenUsage)
    duration_ms: Optional[float] = None
    stderr: Optional[str] = None
    raw_payload: Dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class ModelDescriptor:
    """
    Describes a model supported by a provider.

    Attributes:
        id: Provider-specific model identifier (e.g., "pro", "flash")
        display_name: Human-friendly name for UIs/logs
        capabilities: Feature flags supported by this model
        routing_hints: Optional metadata for routing (cost, latency, etc.)
    """

    id: str
    display_name: Optional[str] = None
    capabilities: Set[ProviderCapability] = field(default_factory=set)
    routing_hints: Dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class ProviderMetadata:
    """
    Provider-level metadata shared with registries and consumers.

    This dataclass describes a provider's capabilities, supported models,
    and configuration, enabling informed routing decisions.

    Attributes:
        provider_id: Canonical provider identifier (e.g., "gemini", "codex")
        display_name: Human-friendly provider name
        models: Supported model descriptors
        default_model: Model ID used when no override supplied
        capabilities: Aggregate capabilities across all models
        security_flags: Provider-specific sandbox/safety configuration
        extra: Arbitrary metadata (version info, auth requirements, etc.)
    """

    provider_id: str
    display_name: Optional[str] = None
    models: Sequence[ModelDescriptor] = field(default_factory=list)
    default_model: Optional[str] = None
    capabilities: Set[ProviderCapability] = field(default_factory=set)
    security_flags: Dict[str, Any] = field(default_factory=dict)
    extra: Dict[str, Any] = field(default_factory=dict)


# =============================================================================
# Error Hierarchy
# =============================================================================


class ProviderError(RuntimeError):
    """Base exception for provider orchestration errors."""

    def __init__(self, message: str, *, provider: Optional[str] = None):
        self.provider = provider
        super().__init__(message)


class ProviderUnavailableError(ProviderError):
    """Raised when a provider cannot be instantiated (binary missing, auth issues)."""


class ProviderExecutionError(ProviderError):
    """Raised when a provider command returns a non-retryable error."""


class ProviderTimeoutError(ProviderError):
    """Raised when a provider exceeds its allotted execution time.

    This error indicates the provider did not respond within the configured
    timeout period. It includes timing information to help with debugging
    and timeout configuration tuning.

    Attributes:
        provider: Provider that timed out
        elapsed: Actual elapsed time in seconds before timeout
        timeout: Configured timeout value in seconds
    """

    def __init__(
        self,
        message: str,
        *,
        provider: Optional[str] = None,
        elapsed: Optional[float] = None,
        timeout: Optional[float] = None,
    ):
        """Initialize provider timeout error.

        Args:
            message: Error message describing the timeout
            provider: Provider that timed out
            elapsed: Actual elapsed time in seconds
            timeout: Configured timeout in seconds
        """
        super().__init__(message, provider=provider)
        self.elapsed = elapsed
        self.timeout = timeout


class ContextWindowError(ProviderExecutionError):
    """Raised when prompt exceeds the model's context window limit.

    This error indicates the prompt/context size exceeded what the model
    can process. It includes token counts to help with debugging and
    provides actionable guidance for resolution.

    Attributes:
        prompt_tokens: Estimated tokens in the prompt (if known)
        max_tokens: Maximum context window size (if known)
        provider: Provider that raised the error
        truncation_needed: How many tokens need to be removed
    """

    def __init__(
        self,
        message: str,
        *,
        provider: Optional[str] = None,
        prompt_tokens: Optional[int] = None,
        max_tokens: Optional[int] = None,
    ):
        """Initialize context window error.

        Args:
            message: Error message describing the issue
            provider: Provider ID that raised the error
            prompt_tokens: Number of tokens in the prompt (if known)
            max_tokens: Maximum tokens allowed (if known)
        """
        super().__init__(message, provider=provider)
        self.prompt_tokens = prompt_tokens
        self.max_tokens = max_tokens
        self.truncation_needed = (
            (prompt_tokens - max_tokens) if prompt_tokens and max_tokens else None
        )


# =============================================================================
# Lifecycle Hooks
# =============================================================================


# Type aliases for hook callables
StreamChunkCallback = Callable[["StreamChunk", ProviderMetadata], None]
BeforeExecuteHook = Callable[[ProviderRequest, ProviderMetadata], None]
AfterResultHook = Callable[[ProviderResult, ProviderMetadata], None]


@dataclass(frozen=True)
class StreamChunk:
    """Represents a streamed fragment emitted by the provider."""

    content: str
    index: int
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ProviderHooks:
    """
    Optional lifecycle hooks wired by the registry.

    Hooks default to None (no-ops) so providers can invoke them unconditionally.
    Registries can wire hooks for observability, logging, or streaming.
    """

    before_execute: Optional[BeforeExecuteHook] = None
    on_stream_chunk: Optional[StreamChunkCallback] = None
    after_result: Optional[AfterResultHook] = None

    def emit_before(self, request: ProviderRequest, metadata: ProviderMetadata) -> None:
        """Emit before-execution hook if registered."""
        if self.before_execute:
            self.before_execute(request, metadata)

    def emit_stream(self, chunk: StreamChunk, metadata: ProviderMetadata) -> None:
        """Emit stream chunk hook if registered."""
        if self.on_stream_chunk:
            self.on_stream_chunk(chunk, metadata)

    def emit_after(self, result: ProviderResult, metadata: ProviderMetadata) -> None:
        """Emit after-result hook if registered."""
        if self.after_result:
            self.after_result(result, metadata)


# =============================================================================
# Abstract Base Class
# =============================================================================


class ProviderContext(ABC):
    """
    Base class for provider implementations.

    Subclasses should:
        * Resolve CLI/environment dependencies during initialization
        * Implement `_execute()` to run the underlying provider
        * Return a populated `ProviderResult` from `_execute()`
        * Emit streaming chunks via `self._emit_stream_chunk()` when
          `request.stream` is True and the provider supports streaming

    The `generate()` method is a template method that:
        1. Calls `_prepare_request()` for any request modifications
        2. Emits the `before_execute` hook
        3. Calls the abstract `_execute()` method
        4. Normalizes exceptions into typed ProviderErrors
        5. Emits the `after_result` hook
        6. Returns the result
    """

    def __init__(
        self,
        metadata: ProviderMetadata,
        hooks: Optional[ProviderHooks] = None,
    ):
        self._metadata = metadata
        self._hooks = hooks or ProviderHooks()

    @property
    def metadata(self) -> ProviderMetadata:
        """Return provider metadata."""
        return self._metadata

    def supports(self, capability: ProviderCapability) -> bool:
        """Return True if any registered model advertises the capability."""
        # Check provider-level capabilities first
        if capability in self._metadata.capabilities:
            return True
        # Then check model-level capabilities
        return any(capability in model.capabilities for model in self._metadata.models)

    def generate(self, request: ProviderRequest) -> ProviderResult:
        """
        Execute the provider with the supplied request (template method).

        Applies lifecycle hooks, normalizes errors, and ensures ProviderStatus
        is consistent across implementations.

        Args:
            request: The generation request

        Returns:
            ProviderResult with the generation output

        Raises:
            ProviderUnavailableError: If provider binary/auth unavailable
            ProviderTimeoutError: If request exceeds timeout
            ProviderExecutionError: For other execution errors
        """
        normalized_request = self._prepare_request(request)
        self._hooks.emit_before(normalized_request, self._metadata)

        try:
            result = self._execute(normalized_request)
        except ProviderTimeoutError:
            raise
        except ProviderUnavailableError:
            raise
        except ProviderError:
            raise
        except FileNotFoundError as exc:
            raise ProviderUnavailableError(
                str(exc), provider=self._metadata.provider_id
            ) from exc
        except TimeoutError as exc:
            raise ProviderTimeoutError(
                str(exc), provider=self._metadata.provider_id
            ) from exc
        except Exception as exc:  # noqa: BLE001 - intentionally wrap all provider exceptions
            raise ProviderExecutionError(
                str(exc), provider=self._metadata.provider_id
            ) from exc

        self._hooks.emit_after(result, self._metadata)
        return result

    def _prepare_request(self, request: ProviderRequest) -> ProviderRequest:
        """
        Allow subclasses to adjust request metadata before execution.

        The default implementation simply returns the request unchanged.
        Subclasses can override to inject defaults, normalize parameters, etc.
        """
        return request

    def _emit_stream_chunk(self, chunk: StreamChunk) -> None:
        """Helper for subclasses to publish streaming output through hooks."""
        self._hooks.emit_stream(chunk, self._metadata)

    async def _run_blocking(
        self,
        func: Callable[..., Any],
        *args: Any,
        timeout: Optional[float] = None,
        **kwargs: Any,
    ) -> Any:
        """Run a blocking operation in the provider executor.

        CLI providers should use this to wrap subprocess.run() and other
        blocking calls to prevent event loop starvation.

        Args:
            func: Blocking function to execute
            *args: Positional arguments for func
            timeout: Optional timeout in seconds
            **kwargs: Keyword arguments for func

        Returns:
            Result of the function call

        Example:
            result = await self._run_blocking(
                subprocess.run,
                ["claude", "--prompt", prompt],
                capture_output=True,
                timeout=30.0,
            )
        """
        try:
            from foundry_mcp.core.executor import get_provider_executor

            executor = get_provider_executor()
            return await executor.run_blocking(func, *args, timeout=timeout, **kwargs)
        except ImportError:
            # Executor not available, run inline
            return func(*args, **kwargs)

    @abstractmethod
    def _execute(self, request: ProviderRequest) -> ProviderResult:
        """
        Subclasses must implement the actual provider invocation.

        Args:
            request: The (possibly modified) generation request

        Returns:
            ProviderResult with generated content and metadata
        """
        raise NotImplementedError


# =============================================================================
# Module Exports
# =============================================================================


__all__ = [
    # Enums
    "ProviderCapability",
    "ProviderStatus",
    # Request/Response dataclasses
    "ProviderRequest",
    "ProviderResult",
    "TokenUsage",
    "StreamChunk",
    # Metadata dataclasses
    "ModelDescriptor",
    "ProviderMetadata",
    # Hooks
    "ProviderHooks",
    "StreamChunkCallback",
    "BeforeExecuteHook",
    "AfterResultHook",
    # Errors
    "ProviderError",
    "ProviderUnavailableError",
    "ProviderExecutionError",
    "ProviderTimeoutError",
    "ContextWindowError",
    # ABC
    "ProviderContext",
]
