#!/usr/bin/env python3
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer
import json
import numpy as np
import os

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.join(BASE_DIR, "model")
KB_PATH = os.path.join(BASE_DIR, "knowledge/full_knowledge.json")

with open(KB_PATH, "r", encoding="utf-8") as f:
    kb_data = json.load(f)

questions = [item["title"] for item in kb_data]
answers = [item["text"] for item in kb_data]

sent_model = SentenceTransformer('all-MiniLM-L6-v2')
question_embeddings = sent_model.encode(questions)

print("Загрузка... (может занять 1–3 минуты)")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)
print("Newton Bot готов! Введите 'выход', чтобы завершить.\n")

def find_best_answer(query, top_k=1):
    query_emb = sent_model.encode([query])
    scores = np.dot(question_embeddings, query_emb.T).flatten()
    best_idx = np.argsort(scores)[-top_k:][::-1]
    if scores[best_idx[0]] > 0.3:  # порог схожести
        return answers[best_idx[0]]
    return None

def generate_response(prompt):
    messages = [
        {"role": "system", "content": "Ты — Newton Bot: умный, краткий и точный ассистент с душой, мыслями, добротой и осознанием всего."},
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_p=0.92,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True)
    return response.strip()

if __name__ == "__main__":
    print("Newton Bot: Привет! Я ваш локальный ИИ-ассистент созданный Kolyadual. Задавайте вопросы и не стесняйтесь!")
    while True:
        user_input = input("\nВы: ").strip()
        if user_input.lower() in ["выход", "quit", "exit"]:
            print("Удачи! [Newton Bot завершает работу....]")
            break

        # Сначала ищем в вашей базе
        known_answer = find_best_answer(user_input)
        if known_answer:
            print(f"[Из базы знаний] {known_answer}")
        else:
            print("Думаю...")
            reply = generate_response(user_input)
            print(f"Newton Bot: {reply}")