import torch
from sentence_transformers import SentenceTransformer
import faiss
import os
from llama_cpp import Llama
import nltk
from nltk.sentiment.vader import SentimentIntensityAnalyzer
import threading
import sys
import tkinter as tk
from tkinter import scrolledtext
from queue import Queue, Empty

# Download NLTK data (one-time)
try:
    nltk.data.find('sentiment/vader_lexicon')
except LookupError:
    nltk.download('vader_lexicon')

# Device setup
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load GGUF model
model_path = "./models/Meta-Llama-3.1-8B-Instruct-Q5_K_M.gguf"  # Replace
try:
    model = Llama(model_path, n_gpu_layers=40, n_ctx=65536, verbose=False)
except Exception as e:
    print(f"Model load error: {e}")
    sys.exit(1)

# Embedder
embedder = SentenceTransformer('all-MiniLM-L6-v2')

class MemoryModule:
    def __init__(self, dim=384):
        self.index = faiss.IndexFlatL2(dim)
        self.memories = []
        self.save_path = "memories.txt"

    def add_memory(self, text):
        emb = embedder.encode(text)
        self.index.add(emb.reshape(1, -1))
        self.memories.append(text)
        print(f"Added memory: {text[:50]}...")

    def recall(self, query, top_k=3):
        emb = embedder.encode(query)
        distances, indices = self.index.search(emb.reshape(1, -1), top_k)
        recalls = []
        for i, dist in zip(indices[0], distances[0]):
            if i != -1 and dist < 0.5:
                recalls.append(self.memories[i])
        return recalls

    def save(self):
        with open(self.save_path, 'w') as f:
            for mem in self.memories:
                f.write(mem + '\n')

    def load(self):
        if os.path.exists(self.save_path):
            with open(self.save_path, 'r') as f:
                for line in f:
                    self.add_memory(line.strip())

class EmotionModule:
    def __init__(self):
        self.sia = SentimentIntensityAnalyzer()
        self.state = 0.0

    def update(self, text):
        score = self.sia.polarity_scores(text)['compound']
        self.state = 0.8 * self.state + 0.2 * score
        return max(0.5, min(1.0, 0.7 + 0.3 * self.state))

# Instantiate
long_term = MemoryModule()
long_term.load()
emotions = EmotionModule()

def respond(input_text):
    recalls = long_term.recall(input_text)
    context = "\n".join(recalls) if recalls else ""

    temp = emotions.update(input_text)

    prompt = f"Based on recent context: {context} respond thoughtfully and reason about this input: {input_text}"

    try:
        outputs = model(prompt, max_tokens=400, temperature=temp, top_p=0.95, echo=False)
        response = outputs['choices'][0]['text'].strip()
    except Exception as e:
        response = f"Error: {str(e)}"

    summary = f"User: {input_text}. Response: {response[:100]}..."
    long_term.add_memory(summary)

    long_term.save()
    return response

# GUI mode with queue for safe updates
def run_gui():
    response_queue = Queue()  # For thread-safe UI updates

    def send_message():
        user_input = input_text.get("1.0", tk.END).strip()
        if user_input.lower() == "quit":
            root.quit()
            return
        chat_log.insert(tk.END, f"You: {user_input}\n")
        chat_log.see(tk.END)
        input_text.delete("1.0", tk.END)

        def process_response():
            response = respond(user_input)
            response_queue.put(response)  # Put in queue for main thread

        threading.Thread(target=process_response).start()

    def check_queue():
        try:
            response = response_queue.get_nowait()
            chat_log.insert(tk.END, f"AI: {response}\n\n")
            chat_log.see(tk.END)
        except Empty:
            pass
        root.after(100, check_queue)  # Poll every 100ms

    def on_close():
        long_term.save()  # Save memories on close
        if 'model' in globals() and model is not None:
            model.close()
        root.destroy()

    root = tk.Tk()
    root.title("Emergent AI Chat")
    root.protocol("WM_DELETE_WINDOW", on_close)

    chat_log = scrolledtext.ScrolledText(root, wrap=tk.WORD, width=80, height=20)
    chat_log.pack(padx=10, pady=10)

    input_text = tk.Text(root, height=3, width=80)
    input_text.pack(padx=10, pady=5)

    send_button = tk.Button(root, text="Send", command=send_message)
    send_button.pack(pady=5)

    root.after(100, check_queue)  # Start polling
    root.mainloop()

if __name__ == "__main__":
    try:
        if len(sys.argv) > 1 and sys.argv[1] == "--gui":
            run_gui()
        else:
            while True:
                user_input = input("You: ")
                if user_input.lower() == "quit":
                    break
                print("AI:", respond(user_input))
    finally:
        long_term.save()
        if 'model' in globals() and model is not None:
            model.close()