import os
import sys
import time
import google.generativeai as genai
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import SupabaseVectorStore
from supabase.client import create_client

# Класс для работы с эмбеддингами и контроля лимитов (ваш рабочий код)
class RateLimitedEmbeddings:
    def __init__(self, api_key, model_name):
        genai.configure(api_key=api_key)
        self.model_name = model_name

    def embed_documents(self, texts):
        embeddings = []
        batch_size = 30 
        
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            print(f"--- Обработка батча: {i} из {len(texts)} ---")
            
            try:
                res = genai.embed_content(
                    model=self.model_name,
                    content=batch,
                    task_type="retrieval_document"
                )
                embeddings.extend(res['embedding'])
                # Пауза 5 секунд между батчами для стабильности на бесплатном тарифе
                time.sleep(5) 
            except Exception as e:
                if "429" in str(e):
                    print("Лимит Google превышен. Ждем 70 секунд...")
                    time.sleep(70)
                    res = genai.embed_content(
                        model=self.model_name,
                        content=batch,
                        task_type="retrieval_document"
                    )
                    embeddings.extend(res['embedding'])
                else:
                    raise e
                    
        return embeddings

    def embed_query(self, text):
        res = genai.embed_content(
            model=self.model_name,
            content=text,
            task_type="retrieval_query"
        )
        return res['embedding']

def load_book_to_supabase():
    # Получаем секреты из окружения GitHub Actions
    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!")
        sys.exit(1)

    print("--- Шаг 1: Подключение к Supabase ---")
    supabase = create_client(supabase_url, supabase_key)

    print("--- Шаг 1.1: Очистка таблицы перед полной загрузкой ---")
    try:
        # Очищаем таблицу, так как id теперь имеет тип text
        supabase.table("documents").delete().neq("id", "0").execute()
        print("Таблица очищена.")
    except Exception as e:
        print(f"Предупреждение при очистке: {e}")

    print("--- Шаг 2: Загрузка файла книги ---")
    if not os.path.exists("book.txt"):
        print("ОШИБКА: Файл book.txt не найден в корне репозитория!")
        sys.exit(1)
        
    try:
        loader = TextLoader("book.txt", encoding="utf-8")
        documents = loader.load()
    except Exception as e:
        print(f"ОШИБКА при чтении файла: {e}")
        sys.exit(1)

    print("--- Шаг 3: Нарезка текста ---")
    # Параметры нарезки (3000 символов с перекрытием 300)
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=3000, chunk_overlap=300)
    docs = text_splitter.split_documents(documents)
    
    # ОГРАНИЧЕНИЕ УБРАНО: Загружаем все фрагменты
    print(f"К ЗАГРУЗКЕ: {len(docs)} фрагментов.")

    print("--- Шаг 4: Инициализация эмбеддингов (Модель 001 / 3072 dim) ---")
    custom_embeddings = RateLimitedEmbeddings(gemini_api_key, "models/gemini-embedding-001")

    print("--- Шаг 5: Запуск процесса загрузки в базу ---")
    try:
        # LangChain сам разобьет docs на вызовы embed_documents по 30 штук
        SupabaseVectorStore.from_documents(
            docs,
            custom_embeddings,
            client=supabase,
            table_name="documents",
            query_name="match_documents",
        )
        print("--- УСПЕХ: Вся книга полностью загружена в Supabase! ---")
    except Exception as e:
        print(f"КРИТИЧЕСКАЯ ОШИБКА при загрузке: {e}")
        sys.exit(1)

if __name__ == "__main__":
    load_book_to_supabase()