import chromadb
import requests
import os
from dotenv import load_dotenv

def call_openrouter_llama3(prompt, api_key):
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    data = {
        "model": "meta-llama/llama-3-8b-instruct",
        "messages": [
            {"role": "system", "content": "Bạn là nhân viên tư vấn của đơn vị trả lời các câu hỏi của khách bằng tiếng Việt tối đa 50 từ."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.2
    }
    response = requests.post(url, headers=headers, json=data)
    response.raise_for_status()
    return response.json()["choices"][0]["message"]["content"]

def prompt_chromadb(input_text=None):
    load_dotenv("venv/api.env")
    chroma_db_path = os.environ.get("CHROMADB_PATH")  # Load path from api.env
    client = chromadb.PersistentClient(path=chroma_db_path)
    collection_name = os.environ.get("CHROMADB_COLLECTION_NAME")  # Load collection name from api.env
    collection = client.get_collection(name=collection_name)
    api_key = os.environ.get("OPENROUTER_API_KEY")

    # Truy vấn Chroma
    results = collection.query(
        query_texts=[input_text],
        n_results=2,
        include=["documents", "distances"]
    )

    documents = results["documents"][0]
    distances = results["distances"][0]

    context_lines = []
    for i, (doc, dist) in enumerate(zip(documents, distances), 1):
        context_lines.append(f"{i}. {doc}")

    # Ghép context
    context = "\n".join(context_lines)

    # Prompt rõ ràng hơn
    prompt = f"""Bạn là nhân viên tư vấn trả lời các thông tin về đơn vị, các câu hỏi của khách bằng tiếng Việt, dựa trên các thông tin mới nhất. Trong trường hợp không có thông tin thì yêu cầu tham khảo website của đơn vị".

    Câu hỏi của quý khách: {input_text}

    Thông tin có sẵn:
    {context}

    Trả lời:"""

    try:
        answer = call_openrouter_llama3(prompt, api_key)
    except Exception as e:
        answer = "Xin lỗi, tôi gặp lỗi khi xử lý câu hỏi của bạn."
    return answer

# if __name__ == "__main__":
#     # Ví dụ sử dụng
#     input_text = "Trường có những ngành học nào?"
#     answer = prompt_chromadb(input_text)
#     print(f"Câu trả lời: {answer}")