import os
import sys
from pathlib import Path
from typing import List

import requests
from bs4 import BeautifulSoup
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.llms import OpenAI
from langchain.chains.question_answering import load_qa_chain
from rich.console import Console
from rich.prompt import Prompt

# --------------------------------------------------------------------------- #
# Конфигурация и глобальные объекты
# --------------------------------------------------------------------------- #

console = Console()

CHROMA_DIR = Path("./chroma_db")
EMBEDDINGS_MODEL = "text-embedding-ada-002"
LLM_MODEL = "gpt-3.5-turbo"

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    console.print("[red]Ошибка: переменная окружения OPENAI_API_KEY не задана[/red]")
    sys.exit(1)

embeddings = OpenAIEmbeddings(model=EMBEDDINGS_MODEL, openai_api_key=OPENAI_API_KEY)
llm = OpenAI(model_name=LLM_MODEL, temperature=0.2, openai_api_key=OPENAI_API_KEY)


def init_vector_store() -> Chroma:
    """
    Инициализирует или загружает существующую базу данных Chroma.
    """
    return Chroma(
        persist_directory=str(CHROMA_DIR),
        embedding_function=embeddings,
    )


# --------------------------------------------------------------------------- #
# Веб‑парсинг и загрузка контента
# --------------------------------------------------------------------------- #

def fetch_text_from_url(url: str) -> str:
    """
    Загружает страницу по URL, удаляет скрипты/стили и возвращает чистый текст.
    """
    try:
        resp = requests.get(url, timeout=10)
        resp.raise_for_status()
    except Exception as e:
        console.print(f"[yellow]Не удалось загрузить {url}: {e}[/yellow]")
        return ""

    soup = BeautifulSoup(resp.text, "html.parser")

    # Удаляем скрипты и стили
    for tag in soup(["script", "style"]):
        tag.decompose()

    text = soup.get_text(separator="\n")
    lines = [line.strip() for line in text.splitlines()]
    cleaned = "\n".join([l for l in lines if l])
    return cleaned


def ingest_urls(urls: List[str], vector_store: Chroma) -> None:
    """
    Загружает тексты по списку URL и сохраняет их в векторный хранилище.
    """
    docs = []
    metadatas = []

    for url in urls:
        console.print(f"[cyan]Обрабатываем {url}[/cyan]")
        content = fetch_text_from_url(url)
        if not content:
            continue
        # Разбиваем на части по 1000 символов для удобства векторизации
        chunk_size = 1000
        for i in range(0, len(content), chunk_size):
            chunk = content[i : i + chunk_size]
            docs.append(chunk)
            metadatas.append({"source": url})

    if docs:
        vector_store.add_texts(docs, metadatas=metadatas)
        console.print(f"[green]Добавлено {len(docs)} фрагментов в базу[/green]")
    else:
        console.print("[red]Нет данных для добавления[/red]")


# --------------------------------------------------------------------------- #
# RAG‑обработчик запросов
# --------------------------------------------------------------------------- #

def answer_query(query: str, vector_store: Chroma) -> str:
    """
    Получает ответ на вопрос, используя векторное хранение и LLM.
    """
    # Поиск похожих документов
    docs = vector_store.similarity_search(query, k=5)
    if not docs:
        return "Извините, я не нашёл подходящей информации."

    chain = load_qa_chain(llm, chain_type="stuff")
    result = chain.run(input_documents=docs, question=query)
    return result


# --------------------------------------------------------------------------- #
# Основная логика программы
# --------------------------------------------------------------------------- #

def main() -> None:
    vector_store = init_vector_store()

    # Если база пуста – запрашиваем URL для загрузки
    if not list(vector_store.get_all_ids()):
        console.print("[yellow]База данных пустая. Пожалуйста, укажите URL-адреса для индексации[/yellow]")
        urls_input = Prompt.ask("Введите URL через запятую")
        urls = [u.strip() for u in urls_input.split(",") if u.strip()]
        ingest_urls(urls, vector_store)
        vector_store.persist()

    console.print("[bold green]RAG‑агент готов к работе![/bold green]")
    while True:
        try:
            query = Prompt.ask("\nВведите ваш вопрос (или 'выход' для завершения)")
            if query.lower() in ("выход", "exit", "quit"):
                console.print("[blue]До свидания![/blue]")
                break
            answer = answer_query(query, vector_store)
            console.print(f"\n[bold]Ответ:[/bold]\n{answer}\n")
        except KeyboardInterrupt:
            console.print("\n[blue]Прервано пользователем. Завершение работы.[/blue]")
            break


if __name__ == "__main__":
    main()