import faiss
import numpy as np
import os
import threading
import pickle
from datetime import datetime
from pymongo import MongoClient
from bson import ObjectId
from sentence_transformers import SentenceTransformer
from llama_cpp import Llama
from rich.console import Console

console = Console()

# --- AYARLAR ---
MONGO_URL = "mongodb://fikoyt:fikoyt1299@fikoyt.com:9001/?replicaSet=rs0&tls=false&authSource=fikoyt"
DB_NAME = "fikoyt"
GGUF_PATH = "./model/Turkish-Llama-8b-Instruct-v0.1.Q4_0.gguf"
INDEX_PATH = "./faiss.index"
DOCS_PATH = "./docs.pkl"
EMBED_DIM = 768*2
CHUNK_SIZE = 300*2
CHUNK_OVERLAP = 50*2

# --- VERİTABANI BAĞLANTISI ---
client = MongoClient(MONGO_URL)
db = client[DB_NAME]
col_docs = db["docs"]
col_history = db["chat_history"]
col_indexed = db["indexed_texts"]

# --- MODELLER ---
embedder = SentenceTransformer("dbmdz/bert-base-turkish-cased", device="cpu")
llm = Llama(model_path=GGUF_PATH, n_ctx=8192, n_threads=16)

# --- YARDIMCI FONKSİYONLAR ---
def chunk_text(text, size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        chunk = words[i:i+size]
        chunks.append(" ".join(chunk))
        i += size - overlap
    return chunks

def init_index():
    if os.path.exists(INDEX_PATH) and os.path.exists(DOCS_PATH):
        index = faiss.read_index(INDEX_PATH)
        with open(DOCS_PATH, "rb") as f:
            all_items = pickle.load(f)
        console.print(f"✅ [green]{len(all_items)} içerik yüklendi.")
    else:
        all_items = []
        index = faiss.IndexFlatL2(EMBED_DIM)
    return index, all_items

index, all_items = init_index()

def save_index():
    faiss.write_index(index, INDEX_PATH)
    with open(DOCS_PATH, "wb") as f:
        pickle.dump(all_items, f)

def embed_and_store(text, source, meta):
    emb = embedder.encode([text]).astype("float32")[0]
    obj = {
        "text": text,
        "embedding": emb.tolist(),
        "source": source,
        "meta": meta,
        "timestamp": datetime.utcnow()
    }
    res = col_indexed.insert_one(obj)
    index.add(np.array([emb], dtype="float32"))
    obj["_id"] = res.inserted_id
    all_items.append(obj)
    save_index()

def embed_pending_texts():
    threading.Timer(30.0, embed_pending_texts).start()

    new_docs = list(col_docs.find({"indexed": {"$ne": True}}))
    for doc in new_docs:
        chunks = chunk_text(doc["text"])
        for c in chunks:
            embed_and_store(c, source="document", meta={"doc_id": str(doc["_id"])})
        col_docs.update_one({"_id": doc["_id"]}, {"$set": {"indexed": True}})

    new_hist = list(col_history.find({"indexed": {"$ne": True}}))
    for item in new_hist:
        metin = f"Soru: {item['question']}\nCevap: {item['answer']}"
        embed_and_store(metin, source="history", meta={"user": item["user"]})
        col_history.update_one({"_id": item["_id"]}, {"$set": {"indexed": True}})

def search(query, k=7):
    if index.ntotal == 0: return []
    q_emb = embedder.encode([query]).astype("float32")
    D, I = index.search(q_emb, k)
    results = []
    for dist, idx in zip(D[0], I[0]):
        item = all_items[idx]
        weight = 1/(dist+1e-6)
        results.append((weight, item))
    results.sort(key=lambda x: x[0], reverse=True)
    return [x[1] for x in results]

def generate(contexts, question):
    metinler = []
    for c in contexts:
        source = c.get("source", "?")
        text = c["text"]
        meta = c.get("meta", {})
        prefix = f"[{source}]"
        metinler.append(prefix + " " + text)
    baglam = "\n---\n".join(metinler)

    prompt = f"""
Sen Türkçe dilinde çalışan güçlü bir yapay zeka asistanısın.  
Aşağıdaki bağlamlara dayanarak, verilen soruya açık, anlaşılır, kısa ve öz cevap ver.

Bağlamlar:
{baglam}

Soru:
{question}

Cevabın:
"""
    result = llm(prompt=prompt, max_tokens=300, stop=["\n\n", "Soru:", "Cevap:"])
    return result["choices"][0]["text"].strip()

def add_doc(text):
    col_docs.insert_one({"text": text, "indexed": False})
    console.print("📥 Belge eklendi, embedding bekleniyor.")

def update_doc(doc_id, new_text):
    try:
        col_docs.update_one({"_id": ObjectId(doc_id)}, {"$set": {"text": new_text, "indexed": False}})
        console.print("✏️ Belge güncellendi, embedding yapılacak.")
    except Exception as e:
        console.print(f"[red]Hata: {e}")

def delete_doc(doc_id):
    try:
        col_docs.delete_one({"_id": ObjectId(doc_id)})
        rebuild_index()
        console.print("🗑️ Belge silindi ve indeks güncellendi.")
    except Exception as e:
        console.print(f"[red]Hata: {e}")

def rebuild_index():
    global index, all_items
    index = faiss.IndexFlatL2(EMBED_DIM)
    all_items = list(col_indexed.find({}))
    if all_items:
        embs = np.array([d["embedding"] for d in all_items]).astype("float32")
        index.add(embs)
        save_index()
        console.print(f"♻️ FAISS yeniden oluşturuldu: {len(all_items)} içerik.")

def chat_session(user):
    console.rule(f"💬 {user} ile sohbet")
    while True:
        q = console.input("[bold blue]🗣️ Soru:[/] ").strip()
        if q.lower() in ["q", "exit", "çık"]: break

        baglamlar = search(q, k=7)
        cevap = generate(baglamlar, q)
        console.print(f"[green]🤖 Cevap:[/] {cevap}")

        col_history.insert_one({
            "user": user,
            "question": q,
            "answer": cevap,
            "timestamp": datetime.utcnow(),
            "indexed": False
        })

def main():
    console.rule("🧠 Ağlayan Bozkurt Gelişmiş RAG Sistemi")
    embed_pending_texts()

    user = console.input("👤 Kullanıcı adı: ").strip() or "anonim"

    while True:
        secim = console.input("\n[a] Ekle | [c] Konuş | [u] Güncelle | [d] Sil | [q] Çık:[/] ").lower()
        if secim == "q":
            console.print("👋 Hoşça kal!")
            break
        elif secim == "a":
            txt = console.input("📄 Metin gir: ").strip()
            if txt: add_doc(txt)
        elif secim == "c":
            chat_session(user)
        elif secim == "u":
            doc_id = console.input("🆔 Belge ID: ")
            new_text = console.input("✏️ Yeni metin: ")
            update_doc(doc_id, new_text)
        elif secim == "d":
            doc_id = console.input("🆔 Silinecek ID: ")
            delete_doc(doc_id)
        else:
            console.print("[red]Geçersiz seçim.")

if __name__ == "__main__":
    main()