"""LlmInstructionExecutor for LLM-based text processing."""

import time
from typing import Any

from pydantic import ValidationError

from agent_flows.core.runtime.variable_interpolator import VariableInterpolator
from agent_flows.exceptions import ConfigurationError, ExecutorError
from agent_flows.executors.base import BaseExecutor
from agent_flows.executors.llm_instruction.models import LlmInstructionExecutorConfig
from agent_flows.integrations import LLMProviderManager
from agent_flows.models.execution import ExecutionContext, ExecutorResult
from agent_flows.models.primitives import ResponseFormat
from agent_flows.services.llm import LLMRequest, LLMService
from agent_flows.utils.logging import get_logger

log = get_logger(__name__)


class LlmInstructionExecutor(BaseExecutor):
    """Executor for LLM instruction steps."""

    def __init__(self) -> None:
        """Initialize the LLM instruction executor."""
        self.interpolator = VariableInterpolator()
        self.llm_service = LLMService()

    def get_required_fields(self) -> list[str]:
        """Get list of required configuration fields."""
        return ["instruction"]

    def get_optional_fields(self) -> list[str]:
        """Get list of optional configuration fields."""
        return [
            "provider",
            "model",
            "temperature",
            "maxTokens",
            "systemPrompt",
            "responseFormat",
            "jsonSchema",
            "widget",
            "apiBase",
            "resultVariable",
            "directOutput",
            "timeout",
            "maxRetries",
        ]

    def validate_config(self, config: dict[str, Any]) -> bool:
        """Validate LLM instruction configuration."""
        try:
            LlmInstructionExecutorConfig(**config)
            return True
        except ValidationError:
            raise
        except Exception as e:
            raise ConfigurationError(
                f"LLM instruction executor configuration validation failed: {e}"
            ) from e

    async def execute(
        self,
        config: dict[str, Any],
        context: ExecutionContext,
    ) -> ExecutorResult:
        """Execute LLM instruction."""
        start_time = time.time()

        try:
            validated_config = LlmInstructionExecutorConfig(**config)

            log.info(
                "Starting LLM instruction execution",
                provider=validated_config.provider,
                model=validated_config.model,
                temperature=validated_config.temperature,
                response_format=validated_config.responseFormat.value,
            )

            # Interpolate instruction and system prompt
            interpolated_instruction = self.interpolator.interpolate(
                validated_config.instruction, context.variables
            )
            interpolated_system_prompt = None
            if validated_config.systemPrompt:
                interpolated_system_prompt = self.interpolator.interpolate(
                    validated_config.systemPrompt, context.variables
                )

            # Emit streaming update
            if context.streaming_handler:
                truncated = (
                    interpolated_instruction[:100] + "..."
                    if len(interpolated_instruction) > 100
                    else interpolated_instruction
                )
                content = f"Processing with {validated_config.provider}/{validated_config.model}: {truncated}"
                context.streaming_handler.stream_step("LLM Instruction", "Brain", content, content)

            # Configure provider credentials
            if validated_config.provider:
                providers_config = {
                    name: cfg.credentials
                    for name, cfg in context.config.llm_providers.providers.items()
                }
                LLMProviderManager.configure(validated_config.provider, providers=providers_config)

            # Build LLM request
            request = self._build_request(
                validated_config, interpolated_instruction, interpolated_system_prompt
            )

            # Execute via LLMService
            response = await self.llm_service.complete(request)

            # Determine output based on response format
            if validated_config.responseFormat in {
                ResponseFormat.JSON,
                ResponseFormat.JSON_SCHEMA,
                ResponseFormat.WIDGET,
            }:
                output_content = response.parsed_content
                raw_content = response.content
            else:
                output_content = response.content
                raw_content = None

            # Update variables if configured
            variables_updated: dict[str, Any] = {}
            if validated_config.resultVariable:
                variables_updated[validated_config.resultVariable] = output_content

            execution_time = time.time() - start_time

            log.info(
                "LLM instruction completed",
                execution_time=execution_time,
                tokens_used=response.usage.get("total_tokens"),
            )

            return ExecutorResult(
                success=True,
                data=output_content,
                variables_updated=variables_updated,
                direct_output=validated_config.directOutput,
                execution_time=execution_time,
                metadata={
                    "provider": validated_config.provider,
                    "model": response.model or validated_config.model,
                    "temperature": validated_config.temperature,
                    "response_format": validated_config.responseFormat.value,
                    "usage": response.usage,
                    "tokens_used": response.usage.get("total_tokens"),
                    "raw_content": raw_content,
                    "step_type": "llmInstruction",
                },
            )

        except ConfigurationError:
            raise
        except Exception as e:
            raise ExecutorError(f"LLM instruction execution failed: {e}") from e

    def _build_request(
        self,
        config: LlmInstructionExecutorConfig,
        instruction: str,
        system_prompt: str | None,
    ) -> LLMRequest:
        """Build LLMRequest from executor config."""
        messages = self.llm_service.build_messages(instruction, system_prompt)

        # Build response format for structured outputs
        response_format: dict[str, Any] | None = None
        if config.responseFormat == ResponseFormat.JSON:
            response_format = {"type": "json_object"}
        elif config.responseFormat in {ResponseFormat.JSON_SCHEMA, ResponseFormat.WIDGET}:
            if config.jsonSchema is None:
                raise ExecutorError(
                    "jsonSchema must be provided when responseFormat is 'json_schema' or 'widget'"
                )
            response_format = {"type": "json_schema", "json_schema": config.jsonSchema}

        return LLMRequest(
            messages=messages,
            model=config.model,
            temperature=config.temperature,
            max_tokens=config.maxTokens,
            timeout=config.timeout or 30,
            max_retries=config.maxRetries or 0,
            response_format=response_format,
            api_base=config.apiBase if config.apiBase and not config.provider else None,
        )

    def get_step_type(self) -> str:
        """Return the step type identifier for this executor."""
        return "llmInstruction"
