"""Starter code for the monitoring homework.

Sets up the text-search RAG from homework 1 and a shared OpenAI client.
"""

import os
from openai import OpenAI

from gitsource import GithubRepositoryDataReader
from minsearch import Index

from rag_helper import RAGBase
from telemetry import tracer

COMMIT = "8c1834d"

# --- Load the course lessons (same as HW1, HW2, HW4) ---
reader = GithubRepositoryDataReader(
    repo_owner="DataTalksClub",
    repo_name="llm-zoomcamp",
    commit_id=COMMIT,
    allowed_extensions={"md"},
    filename_filter=lambda path: "/lessons/" in path,
)
documents = [file.parse() for file in reader.read()]

index = Index(text_fields=["content"], keyword_fields=["filename"])
index.fit(documents)

client = OpenAI(
    api_key=os.environ.get("MISTRAL_API_KEY"),
    base_url="https://api.mistral.ai/v1"
)
rag = RAGBase(index=index, llm_client=client, model="ministral-3b-2512",)


class RAGTraced(RAGBase):

    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):

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

    def llm(self, prompt):

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

            response = super().llm(prompt)

            usage = response.usage

            span.set_attribute(
                "input_tokens",
                usage.prompt_tokens,
            )

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

            return response


rag_traced = RAGTraced(index=index, llm_client=client, model="ministral-3b-2512",)

if __name__ == "__main__":
    query = "How does the agentic loop keep calling the model until it stops?"
    answer = rag.rag(query)
    print(answer)