"""
Chatbot logic — context-stuffs the extracted study material into the
system prompt and streams the response back.
"""

from prompts import CHAT_SYSTEM_PROMPT
from llm import groq, CHAT_MODEL
from rag import retrieve_relevant_chunks, find_direct_reference, clean_math_output

def extract_text_content(content) -> str:
    """
    Gradio's Chatbot can hand back message content as either a plain string
    or a list of content-part dicts (e.g. [{"type": "text", "text": "..."}])
    depending on version/message type. Normalize to a plain string either way.
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return " ".join(
            part.get("text", "") for part in content
            if isinstance(part, dict) and part.get("type") == "text"
        )
    return str(content)

def chat(history, study_context, vector_store):
    """
    Streams a response from the LLM based on the chat history and study context.
    """
    latest_message = extract_text_content(history[-1]["content"])

    relevant_chunks = retrieve_relevant_chunks(vector_store, latest_message, k=4)
    direct_match = find_direct_reference(study_context, latest_message)

    pieces = relevant_chunks + ([direct_match] if direct_match else [])
    retrieved_context = "\n\n".join(pieces)

    system_message = CHAT_SYSTEM_PROMPT.format(study_context=retrieved_context)

    history_for_api = [{"role": h["role"], "content": h["content"]} for h in history]
    messages = [{"role": "system", "content": system_message}] + history_for_api

    stream = groq.chat.completions.create(model=CHAT_MODEL, messages=messages, stream=True)
    response = ""
    for chunk in stream:
        response += chunk.choices[0].delta.content or ""
        yield history + [{"role": "assistant", "content": clean_math_output(response)}]