from llama_index.core.query_pipeline import (
    QueryPipeline,
    InputComponent,
    ArgPackComponent,
)
from llama_index.core.prompts import PromptTemplate
from llama_index.llms.ollama import Ollama
from typing import Any, Dict, List, Optional
from llama_index.core.bridge.pydantic import Field
from llama_index.core.llms import ChatMessage
from llama_index.core.query_pipeline import CustomQueryComponent
from llama_index.core.schema import NodeWithScore


DEFAULT_CONTEXT_PROMPT = (
    "Here is some context that may be relevant:\n"
    "-----\n"
    "{node_context}\n"
    "-----chat_history-----\n"
    "{chat_history}\n"
    "-----\n"
    "Please write a response to the following question, using the above context:\n"
    "{query_str}\n"
)


class ResponseComponent(CustomQueryComponent):
    llm: Ollama = Field(..., description="The language model to use for generating responses.")
    system_prompt: Optional[str] = Field(default=None, description="System prompt to use for the LLM")
    context_prompt: str = Field(
        default=DEFAULT_CONTEXT_PROMPT,
        description="Context prompt use for llm",
    )
    
    @property
    def _input_keys(self) -> set:
        return {"chat_history", "nodes", "query_str"}
    
    @property
    def _output_keys(self) -> set:
        return {"response"}

    def _prepare_context(
        self,
        chat_history: List[ChatMessage],
        nodes: List[NodeWithScore],
        query_str: str,
    ) -> List[ChatMessage]:
        node_context = ""
        for idx, node in enumerate(nodes):
            node_text = node.get_content(metadata_mode="llm")
            node_context += f"Context Chunk {idx}:\n{node_text}\n\n"
        
        formatted_context = self.context_prompt.format(
            node_context=node_context,
            query_str=query_str,
            chat_history=chat_history
        )
        user_message = ChatMessage(role="user", content=formatted_context)
        chat_history.append(user_message)
        
        if self.system_prompt is not None:
            chat_history = [
                ChatMessage(role="system", content=self.system_prompt),
            ] + chat_history
            
        return chat_history
    
    def _run_component(self, **kwargs) -> Dict[str, Any]:
        chat_history = kwargs["chat_history"]
        nodes = kwargs["nodes"]
        query_str = kwargs["query_str"]

        prepared_context = self._prepare_context(
            chat_history=chat_history,
            nodes=nodes,
            query_str=query_str,
        )

        response = self.llm.chat(prepared_context)
        
        return {"response": response}
        
    async def _arun_component(self, **kwargs: Any) -> Dict[str, Any]:
        chat_history = kwargs["chat_history"]
        nodes = kwargs["nodes"]
        query_str = kwargs["query_str"]

        prepared_context = self._prepare_context(
            chat_history=chat_history,
            nodes=nodes,
            query_str=query_str,
        )

        response = await self.llm.chat(prepared_context)
        
        return {"response": response}