import gradio as gr
import glob
import requests
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from docx import Document
from ddgs import DDGS
import trafilatura
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# ---------- 1. Загрузка документа ----------
def find_docx():
    for pattern in ["*.docx", "*.DOCX", "knowledge.docx"]:
        files = glob.glob(pattern)
        if files:
            return files[0]
        files = glob.glob(f"**/{pattern}", recursive=True)
        if files:
            return files[0]
    return None

def load_local_doc():
    path = find_docx()
    if not path:
        raise FileNotFoundError("Документ не найден.")
    doc = Document(path)
    full = "\n".join([p.text for p in doc.paragraphs if p.text.strip()])
    chunks, cur = [], ""
    for para in full.split("\n"):
        if len(cur) + len(para) < 1500:
            cur += para + "\n"
        else:
            if cur:
                chunks.append(cur.strip())
            cur = para + "\n"
    if cur:
        chunks.append(cur.strip())
    vec = TfidfVectorizer(stop_words=None, ngram_range=(1,2), max_features=5000)
    tfidf = vec.fit_transform(chunks)
    return chunks, vec, tfidf

print("Загрузка документа...")
local_chunks, vectorizer, tfidf_matrix = load_local_doc()
print(f"Документ загружен, {len(local_chunks)} фрагментов.")

def local_search(query, k=4):
    if not query.strip():
        return []
    qv = vectorizer.transform([query])
    sim = cosine_similarity(qv, tfidf_matrix).flatten()
    top = sim.argsort()[-k:][::-1]
    return [local_chunks[i] for i in top if sim[i] > 0.05]

# ---------- 2. Веб‑поиск (DuckDuckGo) ----------
def web_search(query, max_results=3):
    snippets = []
    try:
        with DDGS() as ddgs:
            results = ddgs.text(query, max_results=max_results)
            for r in results:
                url = r.get('href')
                title = r.get('title')
                if not url:
                    continue
                try:
                    resp = requests.get(url, timeout=10, headers={'User-Agent': 'Mozilla/5.0'})
                    if resp.status_code == 200:
                        text = trafilatura.extract(resp.text, include_comments=False, include_tables=False)
                        if text and len(text) > 200:
                            snippets.append(f"{title}\n{url}\n{text[:1500]}")
                except Exception:
                    continue
    except Exception as e:
        print(f"Веб-поиск: {e}")
    return snippets

# ---------- 3. Локальная генерация (flan-t5-small) ----------
print("Загрузка модели генерации (первый раз может занять минуту)...")
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
print("Модель готова.")

def generate_answer(question, context):
    context = context[:3000]
    prompt = f"""Ответь на вопрос, используя только информацию из Контекста. Если ответа нет, напиши "Не найдено".

Контекст:
{context}

Вопрос: {question}

Ответ:"""
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
    outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.3, do_sample=False)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# ---------- 4. Основная функция ----------
def answer(query):
    if not query.strip():
        return "Введите вопрос."
    local = local_search(query)
    web = web_search(query)
    all_context = local + web
    if not all_context:
        return "Ничего не найдено."
    return generate_answer(query, "\n\n---\n\n".join(all_context))

# ---------- 5. Интерфейс ----------
with gr.Blocks() as demo:
    gr.Markdown("# Поиск")
    with gr.Row():
        q = gr.Textbox(label="Вопрос", lines=2, placeholder="Введите запрос...")
        btn = gr.Button("Найти", variant="primary")
    out = gr.Markdown(label="Результат")
    btn.click(fn=answer, inputs=q, outputs=out)

demo.launch()