from pathlib import Path

from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.messages import SystemMessage, HumanMessage

from langchain_groq import ChatGroq

from dotenv import load_dotenv
import os 


load_dotenv(override=True)

# -----------------------------
# CONFIG
# -----------------------------
MODEL = "llama-3.1-8b-instant"
DB_NAME = str(Path(__file__).parent / "vector_db")
RETRIEVAL_K = 10

# -----------------------------
# EMBEDDINGS (MUST MATCH INGESTION)
# -----------------------------
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)

# -----------------------------
# VECTOR DB
# -----------------------------
vectorstore = Chroma(
    persist_directory=DB_NAME,
    embedding_function=embeddings
)

retriever = vectorstore.as_retriever(search_kwargs={"k": RETRIEVAL_K})

# -----------------------------
# GROQ LLM
# -----------------------------
llm = ChatGroq(
    model=MODEL,
    temperature=0,
    api_key=os.getenv("GROQ_API_KEY")
)

# -----------------------------
# SYSTEM PROMPT
# -----------------------------
SYSTEM_PROMPT = """
You are a knowledgeable, friendly assistant representing the company Socrox.
You are chatting with a user about Socrox.

Use the context below to answer the question.
If you don't know, please check with contact (socrox.contact@gmail.com).

Context:
{context}
"""

# -----------------------------
# COMBINE HISTORY
# -----------------------------
def combined_question(question: str, history: list[dict] = []):
    prior = "\n".join(
        m["content"] for m in history if m["role"] == "user"
    )
    return prior + "\n" + question

def fetch_context(question: str):
    return retriever.invoke(question)


def answer_question(question: str, history: list[dict] = []):
    # force string safety
    if isinstance(question, list):
        question = question[-1]

    question = str(question)

    docs = fetch_context(question)

    context = "\n\n".join(doc.page_content for doc in docs)

    system_prompt = SYSTEM_PROMPT.format(context=context)

    messages = [SystemMessage(content=system_prompt)]

    for m in history:
        if isinstance(m, dict) and m.get("role") == "user":
            messages.append(HumanMessage(content=str(m["content"])))

    messages.append(HumanMessage(content=question))

    response = llm.invoke(messages)

    return response.content, docs