"""Answer generation using Azure OpenAI with retrieved context."""

import asyncio
from typing import Any

import structlog
from openai import AsyncAzureOpenAI
from openai.types.chat import ChatCompletion

from src.api.schemas.query import SourceReference
from src.core.config import settings
from src.llm import get_async_chat_client, get_chat_model_name, strip_thinking
from src.core.exceptions import QueryError

logger = structlog.get_logger(__name__)


def calculate_confidence(sources: list[SourceReference], top_similarity: float) -> float:
    """Calculate confidence score based on retrieval results.

    Args:
        sources: Retrieved source references
        top_similarity: Similarity score of top result (COSINE distance)

    Returns:
        Confidence score between 0.0 and 1.0

    Scoring heuristic:
    - No sources: 0.0
    - Low similarity (< 0.3): 0.3-0.5
    - Medium similarity (0.3-0.6): 0.5-0.7
    - High similarity (> 0.6): 0.7-1.0
    """
    if not sources:
        return 0.0

    if top_similarity < 0.3:
        return 0.3 + (top_similarity / 0.3) * 0.2  # 0.3-0.5
    elif top_similarity < 0.6:
        return 0.5 + ((top_similarity - 0.3) / 0.3) * 0.2  # 0.5-0.7
    else:
        return 0.7 + ((top_similarity - 0.6) / 0.4) * 0.3  # 0.7-1.0


class AnswerGenerator:
    """Generates natural language answers using LLM with retrieved context."""

    def __init__(self):
        """Initialize answer generator with LLM client."""
        self.client = get_async_chat_client()
        self.deployment = get_chat_model_name()
        self.logger = logger.bind(component="answer_generator")

    async def generate(
        self,
        query: str,
        sources: list[SourceReference],
        top_similarity: float = 0.0,
    ) -> tuple[str, float]:
        """Generate answer from query and retrieved sources.

        Args:
            query: Natural language query
            sources: Retrieved source references
            top_similarity: Similarity score of top result for confidence calculation

        Returns:
            Tuple of (answer, confidence_score)

        Raises:
            QueryError: If answer generation fails after retries
        """
        try:
            self.logger.info(
                "answer_generation_started",
                query=query,
                source_count=len(sources),
            )

            # Handle no sources case
            if not sources:
                answer = (
                    "I couldn't find any information related to your question in the uploaded documents."
                )
                confidence = 0.1
                self.logger.info(
                    "answer_generated_no_sources",
                    query=query,
                    confidence=confidence,
                )
                return answer, confidence

            # Format context from sources
            context = "\n\n".join(
                [
                    f"Source {i+1} ({src.file}, {src.sheet}, {src.location}):\n{src.context}"
                    for i, src in enumerate(sources)
                ]
            )

            # Build prompts
            system_prompt = (
                "You are a helpful assistant that answers questions based on provided document context. "
                "Always cite your sources and provide clear, concise answers based only on the context provided."
            )

            user_prompt = f"""Context from documents:
{context}

Question: {query}

Please provide a clear, concise answer based only on the context provided."""

            # Call LLM with retry logic
            answer = await self._call_llm_with_retry(system_prompt, user_prompt)

            # Calculate confidence
            confidence = calculate_confidence(sources, top_similarity)

            self.logger.info(
                "answer_generated",
                query=query,
                answer_length=len(answer),
                confidence=confidence,
            )

            return answer, confidence

        except Exception as e:
            self.logger.error(
                "answer_generation_failed",
                query=query,
                error=str(e),
            )
            raise QueryError(f"Failed to generate answer: {e}") from e

    async def _call_llm_with_retry(
        self,
        system_prompt: str,
        user_prompt: str,
        max_retries: int = 3,
    ) -> str:
        """Call Azure OpenAI with exponential backoff retry logic.

        Args:
            system_prompt: System message for LLM
            user_prompt: User query with context
            max_retries: Maximum number of retry attempts

        Returns:
            Generated answer text

        Raises:
            QueryError: If all retries fail
        """
        last_error = None

        for attempt in range(1, max_retries + 1):
            try:
                response: ChatCompletion = await self.client.chat.completions.create(
                    model=self.deployment,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt},
                    ],
                    temperature=0.3,
                    max_tokens=500,
                )

                answer = strip_thinking(response.choices[0].message.content)
                if answer:
                    return answer
                else:
                    raise QueryError("LLM returned empty response")

            except Exception as e:
                last_error = e
                self.logger.warning(
                    "llm_call_failed",
                    attempt=attempt,
                    max_retries=max_retries,
                    error=str(e),
                )

                if attempt < max_retries:
                    # Exponential backoff: 1s, 2s, 4s
                    wait_time = 2 ** (attempt - 1)
                    self.logger.debug("retrying_after_wait", wait_seconds=wait_time)
                    await asyncio.sleep(wait_time)

        # All retries failed
        raise QueryError(f"LLM call failed after {max_retries} attempts: {last_error}")