import logging

from config import MODEL_NAME

logger = logging.getLogger(__name__)

INSTRUCTIONS = """
Your task is to answer questions from the person at accident site
based on the provided context.

Use the context to find relevant information and provide accurate
answers. If the answer is not found in the context,
respond with "I don't know."
"""

PROMPT_TEMPLATE = """
QUESTION: {question}

CONTEXT:
{context}
""".strip()


def rewrite_query(user_query: str, client) -> str:
    system_prompt = (
        "You are an AI assistant specialized in optimizing search queries for RAG pipelines. "
        "Your task is to rewrite the user's raw input query to make it clearer, more descriptive, "
        "and packed with relevant keywords that match technical documentation. "
        "Output ONLY the final rewritten query text. Do not add explanations or quotes."
    )
    logger.debug("Original query: %s", user_query)
    response = client.completion(
        model=MODEL_NAME,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Original query: {user_query}"},
        ],
        temperature=0.0,
    )
    rewritten = response.choices[0].message.content.strip()
    logger.debug("Rewritten query: %s", rewritten)
    return rewritten


class RAGBase:

    def __init__(
        self,
        index=None,
        llm_client=None,
        instructions=INSTRUCTIONS,
        prompt_template=PROMPT_TEMPLATE,
        model=MODEL_NAME,
    ):
        self.index = index
        self.llm_client = llm_client
        self.instructions = instructions
        self.prompt_template = prompt_template
        self.model = model

    def search(self, query, num_results=10):
        boost_dict = {"question": 0.5, "answer": 1.0}

        return self.index.search(query, num_results=num_results, boost_dict=boost_dict)

    def build_context(self, search_results):
        lines = []

        for doc in search_results:
            lines.append("Q: " + doc["question"])
            lines.append("A: " + doc["answer"])
            lines.append("")

        return "\n".join(lines).strip()

    def build_prompt(self, query, search_results):
        context = self.build_context(search_results)

        return self.prompt_template.format(question=query, context=context)

    def _call_llm(self, prompt):
        return self.llm_client.completion(
            model=self.model,
            messages=[
                {"role": "developer", "content": self.instructions},
                {"role": "user", "content": prompt},
            ],
        )

    def llm(self, prompt):
        response = self._call_llm(prompt)
        return response.choices[0].message.content

    def rag(self, query):
        search_results = self.search(query, num_results=10)
        prompt = self.build_prompt(query, search_results)
        return self.llm(prompt)