import os
import sys
import google.generativeai as genai
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.vectorstores import SupabaseVectorStore
from langchain.chains import RetrievalQA
from supabase.client import create_client

# Обертка для эмбеддингов (аналогично загрузчику)
class SimpleEmbeddings:
    def __init__(self, api_key):
        genai.configure(api_key=api_key)
    def embed_documents(self, texts):
        res = genai.embed_content(model="models/gemini-embedding-001", content=texts, task_type="retrieval_document")
        return res['embedding']
    def embed_query(self, text):
        res = genai.embed_content(model="models/gemini-embedding-001", content=text, task_type="retrieval_query")
        return res['embedding']

def ask_question(question):
    supabase_url = os.environ.get("SUPABASE_URL")
    supabase_key = os.environ.get("SUPABASE_SERVICE_ROLE_KEY")
    gemini_api_key = os.environ.get("GEMINI_API_KEY")

    if not all([supabase_url, supabase_key, gemini_api_key]):
        print("ОШИБКА: Проверьте секреты GitHub!")
        return

    # 1. Инициализация поиска
    supabase = create_client(supabase_url, supabase_key)
    embeddings = SimpleEmbeddings(gemini_api_key)
    
    vector_store = SupabaseVectorStore(
        client=supabase,
        embedding=embeddings,
        table_name="documents",
        query_name="match_documents",
    )

    # 2. Настройка модели ответов (Gemini 1.5 Flash - быстрая и умная)
    llm = ChatGoogleGenerativeAI(
        model="gemini-1.5-flash",
        google_api_key=gemini_api_key,
        temperature=0.2 # Низкая температура для точных ответов по фактам
    )

    # 3. Создание цепочки вопрос-ответ
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever(search_kwargs={'k': 5}) # Берем 5 самых похожих отрывков
    )

    print(f"\nВопрос: {question}")
    print("Ответ:", qa_chain.run(question))

if __name__ == "__main__":
    if len(sys.argv) > 1:
        user_query = " ".join(sys.argv[1:])
        ask_question(user_query)
    else:
        print("Пожалуйста, введите вопрос. Пример: python chat.py 'О чем эта книга?'")