"""
RAG Generation Engine

Handles LLM generation for RAG responses, including
prompt assembly and provider management.
"""

import logging
from dataclasses import dataclass

from sage.llm import (
    BaseLLMProvider,
    ClaudeProvider,
    LLMConfig,
    LLMError,
    LLMProvider,
    LLMResponse,
    OpenAIProvider,
)
from sage.llm.config import AnthropicConfig, OpenAIConfig
from sage.rag.context import RAGContext
from sage.rag.prompts import (
    NCS_SYSTEM_PROMPT,
    SystemPrompt,
    get_ncs_system_prompt,
    get_qa_template,
)

logger = logging.getLogger(__name__)


@dataclass
class GenerationResult:
    """Result of RAG generation."""

    content: str
    model: str
    provider: LLMProvider
    latency_ms: float = 0.0

    # Token usage
    prompt_tokens: int = 0
    completion_tokens: int = 0

    # Generation metadata
    temperature: float = 0.0
    system_prompt_used: str = ""
    template_used: str = ""

    # Raw response for debugging
    raw_response: LLMResponse | None = None

    @property
    def success(self) -> bool:
        """Check if generation was successful."""
        return bool(self.content)


class GenerationEngine:
    """
    Engine for LLM generation in RAG pipeline.

    Manages LLM providers, prompt assembly, and generation
    with automatic fallback support.

    Example:
        engine = GenerationEngine(config)
        await engine.initialize()

        result = await engine.generate(
            context=rag_context,
            system_prompt=NCS_SYSTEM_PROMPT
        )
        print(result.content)
    """

    def __init__(self, config: LLMConfig):
        """
        Initialize generation engine.

        Args:
            config: LLM configuration with API keys and settings
        """
        self.config = config
        self._primary_provider: BaseLLMProvider | None = None
        self._fallback_provider: BaseLLMProvider | None = None
        self._initialized = False

    async def initialize(self) -> None:
        """Initialize LLM providers."""
        if self._initialized:
            return

        logger.info("Initializing generation engine...")

        # Initialize primary provider
        self._primary_provider = self._create_provider(
            self.config.primary_provider,
            self.config.get_primary_model(),
        )
        await self._primary_provider.initialize()
        logger.info(f"Primary provider initialized: {self.config.primary_provider}")

        # Initialize fallback if configured
        if self.config.fallback_provider:
            fallback_model = self.config.get_fallback_model()
            if fallback_model:
                try:
                    self._fallback_provider = self._create_provider(
                        self.config.fallback_provider,
                        fallback_model,
                    )
                    await self._fallback_provider.initialize()
                    logger.info(f"Fallback provider initialized: {self.config.fallback_provider}")
                except Exception as e:
                    logger.warning(f"Failed to initialize fallback provider: {e}")

        self._initialized = True

    def _create_provider(
        self, provider_type: LLMProvider, model: str
    ) -> BaseLLMProvider:
        """Create an LLM provider instance."""
        if provider_type == LLMProvider.ANTHROPIC:
            if not self.config.anthropic_api_key:
                raise ValueError("Anthropic API key required")
            config = AnthropicConfig(
                api_key=self.config.anthropic_api_key,
                model=model,
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
                timeout=self.config.timeout_seconds,
            )
            return ClaudeProvider(config)

        elif provider_type == LLMProvider.OPENAI:
            if not self.config.openai_api_key:
                raise ValueError("OpenAI API key required")
            config = OpenAIConfig(
                api_key=self.config.openai_api_key,
                model=model,
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
                timeout=self.config.timeout_seconds,
            )
            return OpenAIProvider(config)

        raise ValueError(f"Unsupported provider: {provider_type}")

    async def close(self) -> None:
        """Close providers and clean up."""
        if self._primary_provider:
            await self._primary_provider.close()
        if self._fallback_provider:
            await self._fallback_provider.close()
        self._initialized = False

    async def __aenter__(self) -> "GenerationEngine":
        """Async context manager entry."""
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit."""
        await self.close()

    async def generate(
        self,
        context: RAGContext,
        system_prompt: SystemPrompt | None = None,
        task_type: str = "general",
        temperature: float | None = None,
        max_tokens: int | None = None,
        conversation_history: list[dict] | None = None,
    ) -> GenerationResult:
        """
        Generate a response using RAG context.

        Args:
            context: RAGContext with retrieved chunks
            system_prompt: System prompt to use (default: NCS expert)
            task_type: Task type for system prompt selection
            temperature: Override temperature
            max_tokens: Override max tokens
            conversation_history: Optional previous conversation turns

        Returns:
            GenerationResult with generated content
        """
        if not self._initialized:
            await self.initialize()

        # Select system prompt
        if system_prompt is None:
            system_prompt = get_ncs_system_prompt(task_type)

        # Select QA template
        has_history = conversation_history is not None and len(conversation_history) > 0
        low_confidence = context.confidence_level == "low"
        template = get_qa_template(has_history=has_history, low_confidence=low_confidence)

        # Build the prompt
        template_vars = {
            "context": context.formatted_context,
            "question": context.query,
        }

        if has_history:
            history_str = self._format_history(conversation_history)
            template_vars["history"] = history_str

        user_prompt = template.render(**template_vars)

        # Set generation parameters
        temp = temperature if temperature is not None else self.config.temperature
        tokens = max_tokens or self.config.max_tokens

        # Try primary provider first
        try:
            response = await self._primary_provider.generate(
                prompt=user_prompt,
                system=system_prompt.to_string(),
                temperature=temp,
                max_tokens=tokens,
            )

            return GenerationResult(
                content=response.content,
                model=response.model,
                provider=response.provider,
                latency_ms=response.latency_ms,
                prompt_tokens=response.usage.prompt_tokens,
                completion_tokens=response.usage.completion_tokens,
                temperature=temp,
                system_prompt_used=system_prompt.name,
                template_used=template.name,
                raw_response=response,
            )

        except LLMError as e:
            logger.warning(f"Primary provider failed: {e}")

            # Try fallback if available
            if self._fallback_provider:
                logger.info("Attempting fallback provider...")
                try:
                    response = await self._fallback_provider.generate(
                        prompt=user_prompt,
                        system=system_prompt.to_string(),
                        temperature=temp,
                        max_tokens=tokens,
                    )

                    return GenerationResult(
                        content=response.content,
                        model=response.model,
                        provider=response.provider,
                        latency_ms=response.latency_ms,
                        prompt_tokens=response.usage.prompt_tokens,
                        completion_tokens=response.usage.completion_tokens,
                        temperature=temp,
                        system_prompt_used=system_prompt.name,
                        template_used=template.name,
                        raw_response=response,
                    )

                except LLMError as fallback_error:
                    logger.error(f"Fallback provider also failed: {fallback_error}")
                    raise

            raise

    async def generate_with_messages(
        self,
        messages: list[dict[str, str]],
        system_prompt: SystemPrompt | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ) -> GenerationResult:
        """
        Generate from a message list (for multi-turn conversations).

        Args:
            messages: List of {"role": "user"|"assistant", "content": "..."}
            system_prompt: System prompt to use
            temperature: Override temperature
            max_tokens: Override max tokens

        Returns:
            GenerationResult
        """
        if not self._initialized:
            await self.initialize()

        if system_prompt is None:
            system_prompt = NCS_SYSTEM_PROMPT

        temp = temperature if temperature is not None else self.config.temperature
        tokens = max_tokens or self.config.max_tokens

        try:
            response = await self._primary_provider.generate_with_messages(
                messages=messages,
                system=system_prompt.to_string(),
                temperature=temp,
                max_tokens=tokens,
            )

            return GenerationResult(
                content=response.content,
                model=response.model,
                provider=response.provider,
                latency_ms=response.latency_ms,
                prompt_tokens=response.usage.prompt_tokens,
                completion_tokens=response.usage.completion_tokens,
                temperature=temp,
                system_prompt_used=system_prompt.name,
                template_used="messages",
                raw_response=response,
            )

        except LLMError:
            if self._fallback_provider:
                response = await self._fallback_provider.generate_with_messages(
                    messages=messages,
                    system=system_prompt.to_string(),
                    temperature=temp,
                    max_tokens=tokens,
                )

                return GenerationResult(
                    content=response.content,
                    model=response.model,
                    provider=response.provider,
                    latency_ms=response.latency_ms,
                    prompt_tokens=response.usage.prompt_tokens,
                    completion_tokens=response.usage.completion_tokens,
                    temperature=temp,
                    system_prompt_used=system_prompt.name,
                    template_used="messages",
                    raw_response=response,
                )
            raise

    def _format_history(self, history: list[dict]) -> str:
        """Format conversation history for prompt."""
        lines = []
        for turn in history:
            role = turn.get("role", "user").upper()
            content = turn.get("content", "")
            lines.append(f"{role}: {content}")
        return "\n\n".join(lines)

    @property
    def primary_model(self) -> str:
        """Get the primary model name."""
        if self._primary_provider:
            return self._primary_provider.model_name
        return self.config.get_primary_model()

    @property
    def fallback_model(self) -> str | None:
        """Get the fallback model name."""
        if self._fallback_provider:
            return self._fallback_provider.model_name
        return self.config.get_fallback_model()
