#!/usr/bin/env python3
import asyncio
import json
import ollama
import websockets
from websockets.server import serve
from datetime import datetime

#chatmodel = "hf.co/Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2-GGUF:Llama-3.1-8B-Lexi-Uncensored_V2_Q4.gguf"
embedmodel = 'hf.co/CompendiumLabs/bge-base-en-v1.5-gguf'

# these are test models just ignore
#chatmodel = 'hf.co/bartowski/Llama-3.2-1B-Instruct-GGUF'
#chatmodel = 'hf.co/VibeStudio/Nidum-Llama-3.2-3B-Uncensored-GGUF'
#chatmodel = "faggotchan"
chatmodel = "qwen3.5:4b"

NUM_THREADS = 6  # hardcoded since threads.txt is gone, change if needed

def times() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def nots(line: str) -> str:
    if line.startswith("[") and "]" in line:
        return line.split("]", 1)[1].lstrip()
    return line


# THANK YOU GUY FROM HUGGING FACE THAT MADE ALL THIS SO THAT I DON'T HAVE TO THINK I LOVE YOU

# Each element in the VECTOR_DB will be a tuple (chunk, embedding)
# The embedding is a list of floats, for example: [0.1, 0.04, -0.34, 0.21, ...]
VECTOR_DB = []

def add_chunk_to_database(chunk):
    embedding = ollama.embed(model=embedmodel, input=chunk)['embeddings'][0]
    VECTOR_DB.append((chunk, embedding))

def cosine_similarity(a, b):
    dot_product = sum([x * y for x, y in zip(a, b)])
    norm_a = sum([x ** 2 for x in a]) ** 0.5
    norm_b = sum([x ** 2 for x in b]) ** 0.5
    return dot_product / (norm_a * norm_b)

def retrieve(query, top_n=5):
    query_embedding = ollama.embed(model=embedmodel, input=query)['embeddings'][0]
    similarities = []
    for chunk, embedding in VECTOR_DB:
        similarity = cosine_similarity(query_embedding, embedding)
        similarities.append((chunk, similarity))
    similarities.sort(key=lambda x: x[1], reverse=True)
    return similarities[:top_n]

def rebuild_vector_db(longterm_lines: list[str]):
    VECTOR_DB.clear()
    for i, chunk in enumerate(longterm_lines):
        add_chunk_to_database(chunk)
        print(f"[{times()}] Added chunk {i+1}/{len(longterm_lines)} to the database")

interaction_count = 0


async def ollamachat(websocket):
    global interaction_count
    try:
        # bot sends one json payload with everything main needs
        raw = await websocket.recv()
        payload = json.loads(raw)

        system_prompt = payload["system_prompt"]
        user_input    = payload["user_input"]
        shortterm     = payload.get("shortterm", [])   # list of strings
        longterm      = payload.get("longterm", [])    # list of strings

        interaction_count += 1

        # rebuild vector db if bot signals it (every 150 interactions bot sends full longterm)
        if payload.get("rebuild_vectors"):
            rebuild_vector_db(longterm)
            print(f"[{times()}] Vector DB rebuilt with {len(longterm)} chunks")

        user_query = nots(user_input)
        retrieved_knowledge = retrieve(user_query) if VECTOR_DB else []

        context_message = (
            "These are memory notes to help you continue naturally. "
            "Do NOT imitate tone, style, structure, or wording.\n\n"
            "Short-term memory, these are the most 15 recent chats sent by various people:\n"
            + "\n".join(f"- {line}" for line in shortterm)
            + "\n\nLong-term memory, these are retrieved older messages using cosine similarity:\n"
            + "\n".join(f"- {chunk}" for chunk, _ in retrieved_knowledge)
        )

        messages = [
            {"role": "system", "content": system_prompt + "\n\n" + context_message},
            {"role": "user",   "content": user_input},
        ]

        response = ollama.chat(
            model=chatmodel,
            messages=messages,
            options={
                "temperature":    0.8,
                "top_p":          0.9,
                "top_k":          40,
                "repeat_penalty": 1.1,
                "num_predict":    4096,
                "num_thread":     NUM_THREADS,
            }
        )

        bot_reply = response["message"]["content"].strip()

        print(messages)
        print(f"[{times()}] reply: {bot_reply}\n")

        await websocket.send(json.dumps({"ok": bot_reply}))

    except Exception as e:
        error_msg = f"[{times()}] error in ollamachat: {e}"
        print(error_msg)
        try:
            await websocket.send(json.dumps({"error": error_msg}))
        except Exception:
            pass  # if we can't even send the error, nothing we can do


async def main():
    print(f"[{times()}] hosting on ws://0.0.0.0:8765")
    async with serve(ollamachat, "0.0.0.0", 8765, ping_interval=None, ping_timeout=None):
        await asyncio.Future()


if __name__ == "__main__":
    asyncio.run(main())