"""Executor: LLM call execution with optional structured output validation.

Wraps LiteLLM (provider switching) and Instructor (schema validation + retry).
Requires the `execution` extra: pip install prompt-builder[execution]
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, TypeVar

from pydantic import BaseModel

from .types import CompilationResult

T = TypeVar("T", bound=BaseModel)


@dataclass
class ExecutionConfig:
    """Configuration for the executor."""

    model: str = "openai/gpt-4o"
    temperature: float = 0.0
    max_tokens: int = 4096
    max_retries: int = 3  # for Instructor validation retries
    extra_params: dict[str, Any] = field(default_factory=dict)


@dataclass
class ExecutionResult:
    """Result of an LLM execution."""

    content: str | None = None
    parsed: BaseModel | None = None
    raw_response: Any = None
    model: str = ""
    usage: dict[str, int] = field(default_factory=dict)


class Executor:
    """Executes compiled prompts against an LLM.

    Uses LiteLLM for provider-agnostic calls and Instructor for
    structured output validation.

    Usage:
        executor = Executor(ExecutionConfig(model="anthropic/claude-sonnet-4-5-20250514"))

        # Plain text response
        result = executor.run(compiled)

        # Structured output
        result = executor.run(compiled, response_model=MySchema)
    """

    def __init__(self, config: ExecutionConfig | None = None) -> None:
        self.config = config or ExecutionConfig()
        self._litellm = None
        self._instructor_client = None

    def _get_litellm(self):
        if self._litellm is None:
            try:
                import litellm
                self._litellm = litellm
            except ImportError:
                raise ImportError(
                    "litellm is required for execution. "
                    "Install with: pip install prompt-builder[execution]"
                )
        return self._litellm

    def _get_instructor(self):
        if self._instructor_client is None:
            try:
                import instructor
                litellm = self._get_litellm()
                self._instructor_client = instructor.from_litellm(
                    litellm.completion
                )
            except ImportError:
                raise ImportError(
                    "instructor is required for structured output. "
                    "Install with: pip install prompt-builder[execution]"
                )
        return self._instructor_client

    def run(
        self,
        compiled: CompilationResult,
        response_model: type[T] | None = None,
    ) -> ExecutionResult:
        """Execute a compiled prompt.

        Args:
            compiled: Output from the compiler.
            response_model: Optional Pydantic model for structured output.
                           When provided, Instructor handles validation + retry.

        Returns:
            ExecutionResult with content and/or parsed output.
        """
        if response_model is not None:
            return self._run_structured(compiled, response_model)
        return self._run_plain(compiled)

    def _run_plain(self, compiled: CompilationResult) -> ExecutionResult:
        """Plain text completion via LiteLLM."""
        litellm = self._get_litellm()

        response = litellm.completion(
            model=self.config.model,
            messages=compiled.messages,
            temperature=self.config.temperature,
            max_tokens=self.config.max_tokens,
            **self.config.extra_params,
        )

        choice = response.choices[0]
        usage = {}
        if hasattr(response, "usage") and response.usage:
            usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            }

        return ExecutionResult(
            content=choice.message.content,
            raw_response=response,
            model=self.config.model,
            usage=usage,
        )

    def _run_structured(
        self, compiled: CompilationResult, response_model: type[T]
    ) -> ExecutionResult:
        """Structured output via Instructor + LiteLLM."""
        client = self._get_instructor()

        result = client.chat.completions.create(
            model=self.config.model,
            messages=compiled.messages,
            response_model=response_model,
            temperature=self.config.temperature,
            max_tokens=self.config.max_tokens,
            max_retries=self.config.max_retries,
            **self.config.extra_params,
        )

        return ExecutionResult(
            parsed=result,
            content=result.model_dump_json(indent=2) if result else None,
            model=self.config.model,
        )
