import streamlit as st
from datasets import load_dataset
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel, pipeline
from sentence_transformers import SentenceTransformer
import time
from datetime import datetime
import json
import os
import pickle
import random
import re
import warnings
warnings.filterwarnings('ignore')

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

print(f"CUDA доступна: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

MODEL_NAME = "sberbank-ai/rugpt3small_based_on_gpt2"
EMBEDDING_MODEL = "all-MiniLM-L6-v2"
SCIENCE_DATASET = "RafaelUI/ru_science"
ARTICLE_LIMIT = 50
MAX_LENGTH = 512
TEMPERATURE = 0.8
TOP_P = 0.95

LOG_FILE = "query_logs.json"
EMBEDDINGS_FILE = "science_embeddings.npy"
ARTICLES_FILE = "science_articles.pkl"

AI_NAME = "OpenAirAI"
COMPANY_NAME = "OpenRussianAI"
CREATORS = ["Грибков Евгений", "RootLinux21"]
WEBSITE = "https://sites.google.com/view/opruai/home"
HUGGINGFACE = "https://huggingface.co/OpenRussianAI"
CREATION_DATE = "2026"

st.set_page_config(
    page_title=f"{AI_NAME} - Научный AI-ассистент",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ===================================================================
# 2. БАЗА ЗНАНИЙ ДЛЯ КОРРЕКТНЫХ ОТВЕТОВ
# ===================================================================

KNOWLEDGE = {
    "education": """
### Образование и специальность

Я, OpenAirAI, не учусь в традиционном смысле, но постоянно обучаюсь на новых данных.

**Моя "специальность":**
- Обработка естественного языка
- Поиск информации в научных статьях
- Анализ данных

**Что я умею:**
- Отвечать на научные вопросы
- Находить информацию в статьях ru_science
- Помогать с системным администрированием
- Генерировать уникальные ответы

**Что я не умею:**
- Не имею личного опыта
- Не могу принимать решения за людей
- Не имею физического воплощения
""",
    "mission": """
### Моя миссия

Я создан для того, чтобы:
1. Делать знания доступными
2. Помогать в научных исследованиях
3. Популяризировать науку

**Как я помогаю людям:**
- Отвечаю на вопросы
- Нахожу релевантные статьи
- Объясняю сложные концепции простым языком

**Почему это важно:**
Наука должна быть доступна каждому, а я помогаю убрать барьеры между человеком и знаниями.
""",
    "internet": """
### О существовании интернета

Интернет существует как глобальная сеть, объединяющая миллиарды устройств.

**Что такое интернет:**
- Физическая инфраструктура (кабели, серверы)
- Протоколы передачи данных (TCP/IP)
- Миллионы веб-сайтов и сервисов

**Интернет - это:**
- Инструмент для обмена информацией
- Платформа для общения
- Источник знаний

Я существую в интернете, поэтому могу с уверенностью сказать, что он реален.
""",
    "projects": """
### О проектах и сотрудничестве

OpenRussianAI - открытый проект.

**Как присоединиться:**
1. Посетите наш сайт: https://sites.google.com/view/opruai/home
2. Подпишитесь на Hugging Face: https://huggingface.co/OpenRussianAI
3. Свяжитесь с нами: info@openrussianai.ru

**Что мы делаем:**
- Разрабатываем AI-помощников
- Исследуем NLP
- Популяризируем науку

**Присоединяйтесь к сообществу!**
"""
}

# ===================================================================
# 3. НЕЙРОСЕТЬ
# ===================================================================

class NeuralChatbot:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.tokenizer = None
        self.model = None
        self.generator = None
        self.is_loaded = False
        
        self.system_prompt = f"""Ты - {AI_NAME}, научный AI-ассистент от компании {COMPANY_NAME}.
Ты создан в {CREATION_DATE} командой {', '.join(CREATORS)}.
Ты отвечаешь на русском языке, четко и по делу.
Если вопрос про образование, миссию или интернет - используй знания из базы.
Если вопрос про науку - ищи в статьях.
Всегда давай структурированный ответ.

Вопрос: """
    
    def load_model(self):
        with st.spinner("🧠 Загружаю нейросеть..."):
            try:
                from transformers import GPT2LMHeadModel, GPT2Tokenizer
                
                self.tokenizer = GPT2Tokenizer.from_pretrained(MODEL_NAME)
                self.model = GPT2LMHeadModel.from_pretrained(MODEL_NAME)
                self.model.to(self.device)
                self.model.eval()
                
                self.generator = pipeline(
                    'text-generation',
                    model=self.model,
                    tokenizer=self.tokenizer,
                    device=0 if torch.cuda.is_available() else -1,
                    max_length=300,
                    temperature=TEMPERATURE,
                    top_p=TOP_P,
                    do_sample=True,
                    repetition_penalty=1.2,
                    pad_token_id=self.tokenizer.eos_token_id
                )
                
                self.is_loaded = True
                return True
            except Exception as e:
                st.warning(f"Не удалось загрузить нейросеть: {e}")
                return False
    
    def generate(self, query):
        if not self.is_loaded:
            return self.fallback_response(query)
        
        try:
            prompt = self.system_prompt + query
            
            response = self.generator(
                prompt,
                max_new_tokens=250,
                temperature=TEMPERATURE,
                top_p=TOP_P,
                do_sample=True,
                repetition_penalty=1.2
            )[0]['generated_text']
            
            response = response.replace(prompt, "").strip()
            
            if len(response) < 15:
                return self.fallback_response(query)
            
            return response
            
        except:
            return self.fallback_response(query)
    
    def fallback_response(self, query):
        return f"""Я {AI_NAME} от {COMPANY_NAME}.

Извините, нейросеть временно недоступна. Вот что я знаю по вашему вопросу:

{self.get_knowledge(query)}

Если нужна более подробная информация, попробуйте позже. 🧠"""
    
    def get_knowledge(self, query):
        q = query.lower()
        if any(w in q for w in ["образован", "специальн", "учусь", "умею"]):
            return KNOWLEDGE["education"]
        elif any(w in q for w in ["мисси", "цель", "помога", "задача"]):
            return KNOWLEDGE["mission"]
        elif any(w in q for w in ["интернет", "существу"]):
            return KNOWLEDGE["internet"]
        elif any(w in q for w in ["проект", "сотруднич", "присоедин"]):
            return KNOWLEDGE["projects"]
        else:
            return "Я помогаю с научными вопросами. Что именно вас интересует?"

# ===================================================================
# 4. ОЧИСТКА ЗАПРОСОВ ОТ СПАМА
# ===================================================================

def clean_query(query):
    """Очищает запрос от спама и мусора"""
    # Удаляем ссылки
    query = re.sub(r'http[s]?://\S+', '', query)
    query = re.sub(r'www\.\S+', '', query)
    
    # Удаляем email
    query = re.sub(r'\S+@\S+', '', query)
    
    # Удаляем телефоны
    query = re.sub(r'\+7\s*\(?\d{3}\)?\s*\d{3}\s*\d{2}\s*\d{2}', '', query)
    query = re.sub(r'8\s*\(?\d{3}\)?\s*\d{3}\s*\d{2}\s*\d{2}', '', query)
    
    # Удаляем странные символы
    query = re.sub(r'[^\w\s\.\?\!,;:]', ' ', query)
    
    # Удаляем множественные пробелы
    query = ' '.join(query.split())
    
    # Если запрос слишком длинный - берем первое осмысленное предложение
    if len(query) > 300:
        sentences = re.split(r'[.!?]', query)
        query = '. '.join([s for s in sentences[:3] if len(s) > 10]) + '.'
    
    return query.strip()

def extract_main_question(query):
    """Извлекает основной вопрос из спама"""
    q = query.lower()
    
    # Ключевые слова вопросов
    question_keywords = ['что', 'как', 'почему', 'где', 'когда', 'кто', 'зачем', 'какой']
    
    # Ищем предложения с вопросительными словами
    sentences = re.split(r'[.!?]', q)
    questions = []
    
    for s in sentences:
        if any(keyword in s for keyword in question_keywords):
            questions.append(s.strip())
    
    if questions:
        return '. '.join(questions[:2]) + '?'
    
    # Если вопросов нет, берем первые 100 символов
    return query[:150] + ('...' if len(query) > 150 else '')

# ===================================================================
# 5. ЗАГРУЗКА СТАТЕЙ
# ===================================================================

@st.cache_resource
def load_science_articles():
    if os.path.exists(ARTICLES_FILE):
        with open(ARTICLES_FILE, 'rb') as f:
            return pickle.load(f)
    
    with st.spinner("📚 Загружаю научные статьи..."):
        try:
            dataset = load_dataset(SCIENCE_DATASET, split="train", streaming=True)
            articles = []
            for i, row in enumerate(dataset):
                if i >= ARTICLE_LIMIT:
                    break
                text = row.get('content', '') or row.get('text', '') or str(row)
                title = row.get('title', f"Статья {i}")
                articles.append({
                    "id": i,
                    "title": title[:200],
                    "text": text[:2000],
                    "source": "ru_science"
                })
            with open(ARTICLES_FILE, 'wb') as f:
                pickle.dump(articles, f)
            return articles
        except:
            return []

@st.cache_resource
def load_embedder():
    try:
        return SentenceTransformer(EMBEDDING_MODEL)
    except:
        return None

@st.cache_resource
def create_embeddings(_articles, _embedder):
    if os.path.exists(EMBEDDINGS_FILE):
        return np.load(EMBEDDINGS_FILE)
    if not _articles or _embedder is None:
        return np.array([])
    texts = [f"{a['title']}\n\n{a['text']}" for a in _articles]
    embeddings = _embedder.encode(texts, normalize_embeddings=True, show_progress_bar=True, batch_size=64)
    np.save(EMBEDDINGS_FILE, embeddings)
    return embeddings

def search_articles(query, _articles, _embeddings, _embedder):
    if not _articles or len(_embeddings) == 0 or _embedder is None:
        return []
    try:
        query_vector = _embedder.encode([query], normalize_embeddings=True)[0]
        scores = _embeddings @ query_vector
        top_indices = np.argsort(-scores)[:2]
        results = []
        for idx in top_indices:
            score = float(scores[int(idx)])
            if score > 0.15:
                article = _articles[int(idx)]
                results.append({"title": article['title'], "score": score, "text": article['text'][:500]})
        return results
    except:
        return []

# ===================================================================
# 6. ОСНОВНОЙ КЛАСС
# ===================================================================

class OpenAirAI:
    def __init__(self):
        self.name = AI_NAME
        self.company = COMPANY_NAME
        self.creators = CREATORS
        self.chatbot = NeuralChatbot()
        self.is_ready = False
        
    def initialize(self):
        self.is_ready = self.chatbot.load_model()
        return self.is_ready
    
    def generate_answer(self, query):
        # Очищаем запрос
        clean_q = clean_query(query)
        
        # Если запрос стал пустым
        if not clean_q or len(clean_q) < 3:
            return f"Я {self.name} от {self.company}. Чем могу помочь? Задайте конкретный вопрос. 🧠"
        
        # Извлекаем основной вопрос
        main_q = extract_main_question(clean_q)
        
        # Ищем статьи
        articles_context = search_articles(main_q, articles, embeddings, embedder)
        
        # Если есть статьи и вопрос научный
        if articles_context and any(w in main_q.lower() for w in ["науч", "исслед", "стать", "эксперимент"]):
            response = self.chatbot.generate(main_q)
            if articles_context:
                response += "\n\n📄 **Релевантные статьи:**\n"
                for i, art in enumerate(articles_context, 1):
                    response += f"{i}. {art['title']}\n"
            return response
        
        # Генерируем ответ через нейросеть
        return self.chatbot.generate(main_q)

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

articles = load_science_articles()
embedder = load_embedder()
embeddings = create_embeddings(articles, embedder)

if 'ai' not in st.session_state:
    st.session_state.ai = OpenAirAI()
    st.session_state.ai.initialize()

ai = st.session_state.ai

if "messages" not in st.session_state:
    st.session_state.messages = []
    greeting = ai.generate_answer("Привет! Представься")
    st.session_state.messages.append({"role": "assistant", "content": greeting})

with st.sidebar:
    st.image("https://cdn-icons-png.flaticon.com/512/4248/4248455.png", width=80)
    st.title(f"🧠 {AI_NAME}")
    
    st.markdown(f"""
    **{ai.name}** | {CREATION_DATE}
    
    **Компания:** {ai.company}
    **Разработчики:** {', '.join(ai.creators)}
    
    ---
    **🔗 Ссылки:**
    [🌐 Сайт]({WEBSITE})
    [🤗 Hugging Face]({HUGGINGFACE})
    
    ---
    **📊 Статистика:**
    - Статей: {len(articles)}
    - Сообщений: {len(st.session_state.messages)}
    """)
    
    if st.button("🗑️ Очистить чат"):
        st.session_state.messages = []
        greeting = ai.generate_answer("Привет! Представься")
        st.session_state.messages.append({"role": "assistant", "content": greeting})
        st.rerun()

st.title(f"🧠 {AI_NAME} - Нейросетевой AI-ассистент")
st.markdown(f"**{AI_NAME}** от **{COMPANY_NAME}** | Очищает запросы от спама")

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("Задайте вопрос..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)
    
    with st.chat_message("assistant"):
        with st.spinner("🧠 Думаю..."):
            response = ai.generate_answer(prompt)
            st.markdown(response)
            st.session_state.messages.append({"role": "assistant", "content": response})
    
    st.rerun()

st.divider()
st.caption(f"🧠 {AI_NAME} от {COMPANY_NAME} | Создан в {CREATION_DATE}")