"""Response builder class.

This class provides general functions for taking in a set of text
and generating a response.

Will support different modes, from 1) stuffing chunks into prompt,
2) create and refine separately over each chunk, 3) tree summarization.

"""
import logging
from abc import ABC, abstractmethod
from typing import Any, Dict, Generator, List, Optional, Sequence, Union

from llama_index.callbacks.schema import CBEventType, EventPayload
from llama_index.indices.query.schema import QueryBundle
from llama_index.indices.service_context import ServiceContext
from llama_index.response.schema import RESPONSE_TYPE, Response, StreamingResponse
from llama_index.schema import BaseNode, NodeWithScore, MetadataMode
from llama_index.types import RESPONSE_TEXT_TYPE

logger = logging.getLogger(__name__)

QueryTextType = Union[str, QueryBundle]


class BaseSynthesizer(ABC):
    """Response builder class."""

    def __init__(
        self,
        service_context: Optional[ServiceContext] = None,
        streaming: bool = False,
    ) -> None:
        """Init params."""
        self._service_context = service_context or ServiceContext.from_defaults()
        self._callback_manager = self._service_context.callback_manager
        self._streaming = streaming

    @property
    def service_context(self) -> ServiceContext:
        return self._service_context

    @abstractmethod
    def get_response(
        self,
        query_str: str,
        text_chunks: Sequence[str],
        **response_kwargs: Any,
    ) -> RESPONSE_TEXT_TYPE:
        """Get response."""
        ...

    @abstractmethod
    async def aget_response(
        self,
        query_str: str,
        text_chunks: Sequence[str],
        **response_kwargs: Any,
    ) -> RESPONSE_TEXT_TYPE:
        """Get response."""
        ...

    def _log_prompt_and_response(
        self,
        formatted_prompt: str,
        response: RESPONSE_TEXT_TYPE,
        log_prefix: str = "",
    ) -> None:
        """Log prompt and response from LLM."""
        logger.debug(f"> {log_prefix} prompt template: {formatted_prompt}")
        self._service_context.llama_logger.add_log(
            {"formatted_prompt_template": formatted_prompt}
        )
        logger.debug(f"> {log_prefix} response: {response}")
        self._service_context.llama_logger.add_log(
            {f"{log_prefix.lower()}_response": response or "Empty Response"}
        )

    def _get_metadata_for_response(
        self,
        nodes: List[BaseNode],
    ) -> Optional[Dict[str, Any]]:
        """Get metadata for response."""
        return {node.node_id: node.metadata for node in nodes}

    def _prepare_response_output(
        self,
        response_str: Optional[RESPONSE_TEXT_TYPE],
        source_nodes: List[NodeWithScore],
    ) -> RESPONSE_TYPE:
        """Prepare response object from response string."""
        response_metadata = self._get_metadata_for_response(
            [node_with_score.node for node_with_score in source_nodes]
        )

        if response_str is None or isinstance(response_str, str):
            return Response(
                response_str,
                source_nodes=source_nodes,
                metadata=response_metadata,
            )
        elif response_str is None or isinstance(response_str, Generator):
            return StreamingResponse(
                response_str,
                source_nodes=source_nodes,
                metadata=response_metadata,
            )
        else:
            raise ValueError(
                f"Response must be a string or a generator. Found {type(response_str)}"
            )

    def synthesize(
        self,
        query: QueryTextType,
        nodes: List[NodeWithScore],
        additional_source_nodes: Optional[Sequence[NodeWithScore]] = None,
    ) -> RESPONSE_TYPE:
        event_id = self._callback_manager.on_event_start(CBEventType.SYNTHESIZE)

        if isinstance(query, str):
            query = QueryBundle(query_str=query)

        response_str = self.get_response(
            query_str=query.query_str,
            text_chunks=[
                n.node.get_content(metadata_mode=MetadataMode.LLM) for n in nodes
            ],
        )

        additional_source_nodes = additional_source_nodes or []
        source_nodes = list(nodes) + list(additional_source_nodes)

        response = self._prepare_response_output(response_str, source_nodes)

        self._callback_manager.on_event_end(
            CBEventType.SYNTHESIZE,
            payload={EventPayload.RESPONSE: response},
            event_id=event_id,
        )

        return response

    async def asynthesize(
        self,
        query: QueryTextType,
        nodes: List[NodeWithScore],
        additional_source_nodes: Optional[Sequence[NodeWithScore]] = None,
    ) -> RESPONSE_TYPE:
        event_id = self._callback_manager.on_event_start(CBEventType.SYNTHESIZE)

        if isinstance(query, str):
            query = QueryBundle(query_str=query)

        response_str = await self.aget_response(
            query_str=query.query_str,
            text_chunks=[
                n.node.get_content(metadata_mode=MetadataMode.LLM) for n in nodes
            ],
        )

        additional_source_nodes = additional_source_nodes or []
        source_nodes = list(nodes) + list(additional_source_nodes)

        response = self._prepare_response_output(response_str, source_nodes)

        self._callback_manager.on_event_end(
            CBEventType.SYNTHESIZE,
            payload={EventPayload.RESPONSE: response},
            event_id=event_id,
        )

        return response
