"""Production RAG classes: RAGTraced and RAGProduction."""

from dotenv import load_dotenv

load_dotenv()

from config import make_llm_client
from embedder import Embedder
from ingest import build_index, build_vector_index, load_faq_data
from rag.metrics import RAGWithMetrics
from rag.rag_helper import rewrite_query
from telemetry.tracer import tracer




class RAGTraced(RAGWithMetrics):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def rag(self, query):

        with tracer.start_as_current_span("rag"):
            return super().rag(query)

    def search(self, query, num_results=10):

        with tracer.start_as_current_span("search"):
            return super().search(query, num_results)

    def llm(self, prompt):

        with tracer.start_as_current_span("llm") as span:

            response = super().llm(prompt)

            usage = self.usages[-1] if self.usages else None
            if usage is not None:
                span.set_attribute(
                    "input_tokens",
                    usage.prompt_tokens,
                )

                span.set_attribute(
                    "output_tokens",
                    usage.completion_tokens,
                )

            return response


class RAGProduction(RAGTraced):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.last_rewritten_question = None

    def rag(self, query):
        self.last_rewritten_question = rewrite_query(query, self.llm_client)
        return super().rag(self.last_rewritten_question)

if __name__ == "__main__":
    litellm_client = make_llm_client()

    documents = load_faq_data()

    text_index = build_index(documents)

    embedder = Embedder()
    vector_index = build_vector_index(embedder, documents)

    rag_traced = RAGProduction(
        index=vector_index,
        llm_client=litellm_client,
        text_index=text_index,
        vector_index=vector_index,
        embedder=embedder,
    )

    query = "I just got hit by a car and my friend is lying there bleeding badly—do I need to drag them out of the way or just wait for the ambulance?"
    answer = rag_traced.rag(query)
    print(answer)