import json
import chromadb
from agentic_workflow.config import LLM_MODEL
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage

llm = ChatOllama(model=LLM_MODEL, temperature=0)

def generate_llm_response(query: str, context: list[str], history: list[dict], suggestion: str = "") -> str:
    system_parts = ["You are a helpful assistant."]
    if suggestion:
        system_parts.append(f"Improvement hint from previous attempt: {suggestion}")

    messages = [SystemMessage(content=" ".join(system_parts))]

    # Convert history dicts to LangChain message objects
    for turn in history:
        if turn["role"] == "user":
            messages.append(HumanMessage(content=turn["content"]))
        elif turn["role"] == "assistant":
            messages.append(AIMessage(content=turn["content"]))

    # Append context to the user query if available
    user_content = query
    if context:
        user_content = "Context:\n" + "\n".join(context) + "\n\nQuestion: " + query

    messages.append(HumanMessage(content=user_content))

    response = llm.invoke(messages)
    return response.content

# result = generate_llm_response(
#     query="What is the refund policy?",
#     context=[],
#     history=[]
# )
# print("Test 2:", result)