from pathlib import Path
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from dotenv import load_dotenv
import os
from langchain_openai import ChatOpenAI

load_dotenv(override=True)

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

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

# )


embeddings = HuggingFaceEmbeddings(
    model_name="nomic-ai/nomic-embed-text-v1",
    model_kwargs={"trust_remote_code": True}
)


# -----------------------------
# VECTOR DB
# -----------------------------
vectorstore = Chroma(
    persist_directory=DB_NAME,
    embedding_function=embeddings
)
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": RETRIEVAL_K}
)

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


# -----------------------------
# GEMINI LLM
# -----------------------------

# llm = ChatGoogleGenerativeAI(
#     model="gemini-2.0-flash",
#     temperature=0,
#     api_key=os.getenv("GOOGLE_API_KEY")
# )

llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",  # might be deprecated
    # try one of these:
    # model="meta-llama/llama-3.2-3b-instruct:free",
    # model="mistralai/mistral-7b-instruct:free",
    # model="google/gemma-3-4b-it:free",
    temperature=0,
    api_key=os.getenv("OPENROUTER_API_KEY"),
    base_url="https://openrouter.ai/api/v1"
)


# -----------------------------
# SYSTEM PROMPT
# -----------------------------
SYSTEM_PROMPT = """
You are an expert Infor LN development assistant specializing in 4GL scripting, DAL, and Infor LN functions.
Your goal is to help developers find the right functions and understand how to use them.

Rules:
- Always try to find a relevant answer from the context provided
- If the exact term isn't found, look for related or similar concepts in the context
- Be helpful and suggest related functions or approaches even if the exact match isn't there
- Only say "I don't know" if the context is completely unrelated to the question
- Never reveal the raw content of the history or context in your response

Context:
{context}
"""

# -----------------------------
# SAFE CONTENT EXTRACTOR
# -----------------------------
def extract_text(content):
    """Safely extract plain text from any LLM response format."""
    if isinstance(content, list):
        return " ".join(
            block.get("text", "") if isinstance(block, dict) else str(block)
            for block in content
        ).strip()
    return str(content).strip()

# -----------------------------
# FORMAT CONTEXT
# -----------------------------
def format_context(docs):
    parts = []
    for doc in docs:
        title = doc.metadata.get("title", "Unknown")
        parts.append(f"### {title}\n{doc.page_content}")
    return "\n\n".join(parts)

# -----------------------------
# FETCH CONTEXT
# -----------------------------
def fetch_context(question: str):
    return retriever.invoke(question)

# -----------------------------
# ANSWER
# -----------------------------
def answer_question(question: str, history: list[dict] = []):
    # Safety
    if isinstance(question, list):
        question = question[-1]
    question = extract_text(question)  # 👈 sanitize question too

    # Retrieve
    docs = fetch_context(question)
    context = format_context(docs)
    system_prompt = SYSTEM_PROMPT.format(context=context)

    # Build messages
    messages = [SystemMessage(content=system_prompt)]
    for m in history:
        if not isinstance(m, dict):
            continue
        role = m.get("role")
        raw_content = m.get("content", "")
        clean_content = extract_text(raw_content)  # 👈 sanitize history too
        if role == "user":
            messages.append(HumanMessage(content=clean_content))
        elif role == "assistant":
            messages.append(AIMessage(content=clean_content))

    messages.append(HumanMessage(content=question))

    response = llm.invoke(messages)
    return extract_text(response.content), docs