import os
import requests
from retriever import retrieve_docs
from langchain_community.embeddings import HuggingFaceEmbeddings
from numpy import dot
from numpy.linalg import norm

API_KEY = "AIzaSyClqQssVMjt02qKrGKnghYAK9RkGf0lkS4"

def filter_relevant_docs(docs, query, top_k=3):
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    query_embedding = embeddings.embed_query(query)
    scores = []
    for doc in docs:
        doc_embedding = embeddings.embed_query(doc.page_content)
        cosine_sim = dot(query_embedding, doc_embedding) / (norm(query_embedding) * norm(doc_embedding))
        scores.append((doc, cosine_sim))
    scores.sort(key=lambda x: x[1], reverse=True)
    return [doc for doc, _ in scores[:top_k]]

def format_sources(docs):
    sources = set()
    for doc in docs:
        section = doc.metadata.get("section")
        if section:
            sources.add(section.strip())
        else:
            filename = os.path.basename(doc.metadata.get("source", "Nguồn không xác định"))
            sources.add(filename)
    return "\n".join(f"- {src}" for src in sorted(sources))

def answer_query(query, model="Gemini Pro", temperature=0.2):
    all_docs = retrieve_docs(query)
    if not all_docs:
        return "Không tìm thấy tài liệu liên quan để trả lời.", []

    docs = filter_relevant_docs(all_docs, query)
    context = "\n\n".join([doc.page_content for doc in docs])

    prompt = f"""Dựa trên tài liệu sau, hãy trả lời câu hỏi theo phong cách trang trọng, lịch sự và chuyên nghiệp:

{context}

Câu hỏi: {query}

Yêu cầu:
- Sử dụng từ ngữ lịch sự ("Bạn cần...", "Vui lòng...", "Sau khi...")
- Không sử dụng từ nói miệng như "nhé", "nha", "ok".
- Câu trúc câu đầy đủ, rõ ràng.
- Chỉ trả lời dựa trên thông tin trong tài liệu. Nếu không tìm thấy thông tin liên quan, trả lời: "Thông tin không có trong tài liệu được cung cấp."
- Không tự thêm "Nguồn tham khảo" trong phần trả lời.

Trả lời:"""

    url = f"https://generativelanguage.googleapis.com/v1/models/gemini-1.5-pro:generateContent?key={API_KEY}"
    headers = {"Content-Type": "application/json"}
    payload = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {"temperature": temperature}
    }

    response = requests.post(url, headers=headers, json=payload)
    data = response.json()

    try:
        answer = data['candidates'][0]['content']['parts'][0]['text']
    except Exception as e:
        print("🔴 Response từ Gemini:", data)
        answer = "Lỗi khi gọi Gemini API: " + str(e)

    return answer, docs