import streamlit as st
from datasets import load_dataset
import numpy as np
import torch
from sentence_transformers import SentenceTransformer

# ===================================================================
# 1. НАСТРОЙКИ
# ===================================================================

MODEL_NAME = "BAAI/bge-large-en-v1.5"
DATASET_NAME = "wikimedia/wikipedia"
LANGUAGE = "20231101.en"  # можно заменить на "20231101.ru" для русского
ARTICLE_LIMIT = 1000  # сколько статей загрузить (для экономии памяти)

st.set_page_config(
    page_title="Wikipedia Assistant",
    page_icon="🌖",
    layout="wide"
)

st.title("🌖 Wikipedia Assistant")
st.markdown("Задай вопрос — я найду ответ в Википедии!")

# ===================================================================
# 2. ЗАГРУЗКА ДАТАСЕТА
# ===================================================================

@st.cache_resource
def load_wikipedia():
    """Загружает датасет Википедии"""
    with st.spinner("📚 Загружаю Википедию..."):
        dataset = load_dataset(
            DATASET_NAME,
            LANGUAGE,
            split="train",
            streaming=True
        )
        
        articles = []
        for i, row in enumerate(dataset):
            if i >= ARTICLE_LIMIT:
                break
            articles.append({
                "id": row.get("id", i),
                "title": row.get("title", "Без названия"),
                "text": row.get("text", "")[:3000],
                "url": row.get("url", "")
            })
        
        return articles

@st.cache_resource
def load_embedder():
    """Загружает модель для эмбеддингов"""
    with st.spinner("🧠 Загружаю модель..."):
        return SentenceTransformer(MODEL_NAME)

# ===================================================================
# 3. ПОИСК
# ===================================================================

def search_wikipedia(query, articles, embedder):
    """Ищет ответ на вопрос в Википедии"""
    # Кодируем запрос
    query_vector = embedder.encode(
        [query],
        normalize_embeddings=True
    )[0]
    
    # Кодируем все статьи (если ещё не закодированы)
    if "embeddings" not in st.session_state:
        with st.spinner("🔢 Создаю эмбеддинги статей..."):
            texts = [f"{a['title']}\n\n{a['text']}" for a in articles]
            embeddings = embedder.encode(
                texts,
                normalize_embeddings=True
            )
            st.session_state.embeddings = embeddings
    
    # Ищем
    scores = st.session_state.embeddings @ query_vector
    top_indices = np.argsort(-scores)[:3]
    
    results = []
    for idx in top_indices:
        results.append({
            "article": articles[int(idx)],
            "score": float(scores[int(idx)])
        })
    
    return results

# ===================================================================
# 4. ИНТЕРФЕЙС
# ===================================================================

# Загружаем данные
articles = load_wikipedia()
embedder = load_embedder()

st.success(f"✅ Загружено {len(articles)} статей")

# Поле ввода
query = st.text_input("🔍 Что хочешь узнать?", placeholder="Например: Как измеряют расстояние до галактик?")

if query:
    with st.spinner("🔎 Ищу ответ..."):
        results = search_wikipedia(query, articles, embedder)
    
    # Показываем результаты
    for i, result in enumerate(results, 1):
        score = result["score"]
        article = result["article"]
        
        # Только если релевантно
        if score > 0.3:
            with st.expander(f"#{i} {article['title']} (сходство: {score:.2f})"):
                st.write(article["text"][:2000] + "...")
                if article["url"]:
                    st.link_button("🔗 Читать на Википедии", article["url"])
        else:
            st.warning("😕 Не нашёл подходящей статьи. Попробуй уточнить вопрос.")
else:
    st.info("💡 Напиши вопрос, и я найду ответ в Википедии")