import asyncio
import json
import os
from contextlib import asynccontextmanager
from typing import List, Optional

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

from rasa.builder.exceptions import DocumentRetrievalError
from rasa.builder.models import Document
from rasa.constants import PACKAGE_NAME
from rasa.shared.utils.io import read_json_file

INKEEP_API_KEY_ENV_VAR = "INKEEP_API_KEY"
INKEEP_RAG_RESPONSE_SCHEMA_PATH = str(
    importlib_resources.files(PACKAGE_NAME).joinpath(
        "builder/inkeep-rag-response-schema.json"
    )
)

INKEEP_DOCUMENT_RETRIEVAL_MODEL = "inkeep-rag"
INKEEP_BASE_URL = "https://api.inkeep.com/v1/"

structlogger = structlog.get_logger()


class InKeepDocumentRetrieval:
    """Handles the document retrieval from InKeep AI."""

    def __init__(
        self,
        api_key: Optional[str] = None,
    ):
        self._client: Optional[openai.AsyncOpenAI] = None
        self._rag_schema = read_json_file(INKEEP_RAG_RESPONSE_SCHEMA_PATH)
        self._api_key = api_key or os.getenv(INKEEP_API_KEY_ENV_VAR)

    async def retrieve_documents(
        self, query: str, temperature: float = 0.0, timeout: float = 30.0
    ) -> List[Document]:
        """Retrieve relevant documents using InKeep AI based on the given query.

        Args:
            query: The search query
            temperature: Controls randomness in generation (0.0 for deterministic)
            timeout: Timeout for the API call

        Returns:
            List of Document objects containing retrieved content

        Raises:
            LLMGenerationError: If the API call fails or returns invalid response
        """
        response = await self._call_inkeep_rag_api(
            query=query,
            temperature=temperature,
            timeout=timeout,
        )
        documents = self._parse_documents_from_response(response)
        return documents

    async def _call_inkeep_rag_api(
        self, query: str, temperature: float, timeout: float
    ) -> ChatCompletion:
        """Call InKeep AI RAG's API endpoint and return the response content.

        Args:
            query: The search query to send to InKeep
            temperature: Controls randomness in generation (0.0 for deterministic)
            timeout: Timeout for the API call

        Returns:
            The response content from InKeep AI. The response is made of the retrieved
            documents.

        Raises:
            LLMGenerationError: If the API call fails or returns invalid response
        """
        request_params = {
            "model": INKEEP_DOCUMENT_RETRIEVAL_MODEL,
            "messages": [{"role": "user", "content": query}],
            "temperature": temperature,
            "timeout": timeout,
            "response_format": {
                "type": "json_schema",
                "json_schema": self._rag_schema,
            },
        }
        try:
            async with self._get_client() as client:
                response = await client.chat.completions.create(**request_params)

                if not response.choices[0].message.content:
                    structlogger.warning(
                        "inkeep_document_retrieval.empty_response",
                        event_info="InKeep AI returned an empty response. ",
                        request_params=request_params,
                        response_content=response.choices[0].message.content,
                    )
                    raise DocumentRetrievalError(
                        "InKeep Document Retrieval: Empty response"
                    )

                return response

        except openai.OpenAIError as e:
            structlogger.error(
                "inkeep_document_retrieval.api_error",
                event_info="InKeep Document Retrieval: API error",
                request_params=request_params,
                error=e,
            )
            raise DocumentRetrievalError(f"InKeep Document Retrieval: API error: {e}")
        except asyncio.TimeoutError as e:
            structlogger.error(
                "inkeep_document_retrieval.timeout_error",
                event_info="InKeep Document Retrieval: Timeout error",
                request_params=request_params,
                error=e,
            )
            raise DocumentRetrievalError(f"InKeep AI request timed out: {e}")
        except Exception as e:
            structlogger.error(
                "inkeep_document_retrieval.error",
                event_info="InKeep Document Retrieval: Error",
                request_params=request_params,
                error=e,
            )
            raise DocumentRetrievalError(
                f"InKeep Document Retrieval: Unexpected error: {e}"
            )

    @asynccontextmanager
    async def _get_client(self):
        """Get or create client that handles the API calls to InKeep AI."""
        if self._client is None:
            self._client = openai.AsyncOpenAI(
                api_key=self._api_key,
                base_url=INKEEP_BASE_URL,
            )

        try:
            yield self._client
        except Exception as e:
            structlogger.error(
                "inkeep_document_retrieval.client_error",
                event_info="InKeep Document Retrieval: Client error",
                error=str(e),
            )
            raise

    def _parse_documents_from_response(
        self, response: ChatCompletion
    ) -> List[Document]:
        """Parse the InKeep AI response into Document objects.

        Args:
            response: ChatCompletion response from InKeep AI's RAG model.

        Returns:
            List of Document objects
        """
        try:
            content = response.choices[0].message.content
            if not content:
                return []

            response_data = json.loads(content)
            documents = []

            for item in response_data.get("content", []):
                try:
                    document = Document.from_inkeep_rag_response(item)
                    documents.append(document)
                except Exception as e:
                    structlogger.warning(
                        "inkeep_document_retrieval.invalid_document_skipped",
                        event_info=(
                            "InKeep Document Retrieval: Invalid document structure "
                            "skipped. Returning empty list for this item."
                        ),
                        error=str(e),
                        item=item,
                    )
                    # Continue processing other items, skip this invalid one
                    continue

            return documents

        except json.JSONDecodeError as e:
            structlogger.warning(
                "inkeep_document_retrieval.parse_response_failed",
                event_info=(
                    "InKeep Document Retrieval: Parse response failed. "
                    "Returning empty list.",
                ),
                error=str(e),
            )
            return []
        except Exception as e:
            structlogger.error(
                "inkeep_document_retrieval.parse_response_error",
                event_info=(
                    "InKeep Document Retrieval: Parse response error. "
                    "Returning empty list.",
                ),
                error=str(e),
            )
            return []
