#!/usr/bin/env python3
import asyncio
import base64
import hmac
import json
import os
import re  # thank you claude
import ollama
import websockets
from pathlib import Path
from websockets.server import serve
from utils import nots, now_str, make_chunks, pick_chunks

chatmodel = "hf.co/mradermacher/Huihui-gemma-4-E4B-it-abliterated-GGUF:IQ4_XS"
embedmodel = 'hf.co/CompendiumLabs/bge-base-en-v1.5-gguf'
RAG_THRESHOLD = 0.65  # chunks scoring under this aren't injected at all, watch the "best rag score" print and tweak it
PORT = int(os.environ.get("PORT", 8765))  # PORT=8766 python3 main.py when it runs as the fallback next to gemini.py

# resolve config relative to this script, not whatever dir it happens to get launched from
SCRIPT_DIR = Path(__file__).resolve().parent
CONFIG_DIR = SCRIPT_DIR / "config"

try:
    AUTH_TOKEN = (CONFIG_DIR / "auth_token.txt").read_text(encoding="utf-8").strip()
except FileNotFoundError:
    raise SystemExit("no config/auth_token.txt found — run setup_wizard.py first")

# only one generation at a time, we've got one cpu running one model here
chat_semaphore = asyncio.Semaphore(1)


# 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 = []
rag_loaded = False  # False until the bot sends longterm at least once, after a restart we ask for it

def add_chunk_to_database(chunk):
    embedding = ollama.embed(model=embedmodel, input=chunk)['embeddings'][0]
    return 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)
    if similarities:
        print(f"[{now_str()}] best rag score: {similarities[0][1]:.2f}")
    return pick_chunks(similarities, top_n, RAG_THRESHOLD)

def rebuild_vector_db(longterm_lines: list[str]):
    global VECTOR_DB
    existing = dict(VECTOR_DB)  # chunk -> embedding, so we don't re-embed stuff we already have
    new_db = []
    added = 0
    for i, chunk in enumerate(longterm_lines):
        embedding = existing.get(chunk)
        if embedding is None:
            embedding = add_chunk_to_database(chunk)
            added += 1
            print(f"[{now_str()}] embedded new chunk {added} ({i+1}/{len(longterm_lines)})")
        new_db.append((chunk, embedding))
    VECTOR_DB = new_db


async def ollamachat(websocket):
    global rag_loaded
    # loop over every message on this connection instead of handling one and returning —
    # returning from this function is what makes the websockets lib close the connection,
    # which was breaking bot.py's persistent connection after every single reply
    async for raw in websocket:
        try:
            payload = json.loads(raw)

            if not hmac.compare_digest(str(payload.get("auth_token", "")).encode(), AUTH_TOKEN.encode()):
                await websocket.send(json.dumps({"error": "bad auth token"}))
                continue

            system_prompt = payload["system_prompt"]
            user_input    = payload["user_input"]
            shortterm     = payload.get("shortterm", [])   # list of strings
            longterm      = payload.get("longterm", [])    # list of strings, only there on a rebuild
            media         = payload.get("media", [])       # list of {"mime_type", "data" (base64)}

            # rebuild vector db if bot signals it (every 150 interactions bot sends full longterm),
            # if we never got longterm (just restarted, or first time as a fallback) ask for it and the bot resends
            if payload.get("rebuild_vectors"):
                await asyncio.to_thread(rebuild_vector_db, make_chunks(longterm))
                rag_loaded = True
            elif not rag_loaded:
                await websocket.send(json.dumps({"need_longterm": True}))
                continue
                print(f"[{now_str()}] Vector DB reconciled, {len(VECTOR_DB)} chunks")

            user_query = nots(user_input)
            retrieved_knowledge = await asyncio.to_thread(retrieve, user_query) if VECTOR_DB else []

            # rag chunks go in system so they don't pollute turn history
            rag_block = ""
            if retrieved_knowledge:
                rag_block = "\n\nOlder relevant messages for context:\n" + "\n".join(f"- {chunk.replace(chr(10), ' | ')}" for chunk, _ in retrieved_knowledge)

            # lines to skip — not real dialogue
            skip_fragments = ("sent sticker:", "sent media:", "sent a random dm to ", "saved media:")
            def is_skip(line):
                c = nots(line)
                if any(f in c for f in skip_fragments):
                    return True
                # empty said: like "username said: "
                if "said:" in c and c.split("said:", 1)[1].strip() == "":
                    return True
                return False

            BOT_NAME = "Faggot-chan" #i'll make this a variable later you can edit 

            # parse shortterm lines into (speaker, text) pairs
            def parse_line(line):
                c = nots(line).strip()
                # "X said: Y", "X said in DM: Y", "X replied to ... with reply: Y"
                # look at what comes right after the username (no spaces in those) instead of searching the whole line,
                # otherwise a message that happens to contain " with reply: " gets read as a reply
                speaker, _, rest = c.partition(" ")
                if rest.startswith("said in DM: "):
                    return speaker.strip(), rest[len("said in DM: "):].strip()
                if rest.startswith("replied to ") and ' with reply: "' in rest:
                    # "X replied to Y saying: "Z" with reply: "W""
                    text = rest.rsplit(' with reply: "', 1)[-1]
                    if text.endswith('"'):
                        text = text[:-1]
                    target = rest[len("replied to "):].split(" saying: ", 1)[0]  # keep who they were talking to, otherwise replies lose it
                    return speaker.strip(), f"(replying to {target}) {text.strip()}"
                if rest.startswith("said: "):
                    return speaker.strip(), rest[len("said: "):].strip()
                return None, None

            # build alternating user/assistant turns from shortterm
            # group consecutive non-bot lines into one user blob, bot lines become assistant turns
            messages = []
            current_user_lines = []
            seen_last = None  # for dedup

            for line in shortterm:
                if is_skip(line):
                    continue
                speaker, text = parse_line(line)
                if not speaker or not text:
                    continue
                entry = f"{speaker}: {text}"
                if entry == seen_last:  # skip duplicate consecutive lines
                    continue
                seen_last = entry

                if speaker == BOT_NAME:
                    # flush accumulated user lines first
                    if current_user_lines:
                        messages.append({"role": "user", "content": "\n".join(current_user_lines)})
                        current_user_lines = []
                    messages.append({"role": "assistant", "content": text})
                else:
                    current_user_lines.append(entry)

            # flush any trailing user lines before the final turn
            if current_user_lines:
                messages.append({"role": "user", "content": "\n".join(current_user_lines)})

            # final isolated turn — the specific message she needs to reply to, with the speaker named
            # so she doesn't anchor on whoever talked last in the history above
            speaker, text = parse_line(user_input)
            if speaker and text:
                final_turn = f"[responding to {speaker}]: {text}\n(reply only to {speaker}, the lines before this are other people's messages for context)"
            else:
                final_turn = f"[responding to]: {nots(user_input)}"  # couldn't parse it, old format
            messages.append({"role": "user", "content": final_turn})

            # ollama only takes images (no audio/video/pdf), and only on the final turn so history stays text
            images = [base64.b64decode(m["data"]) for m in media if m["mime_type"].startswith("image/")]
            if images:
                messages[-1]["images"] = images

            # make sure we start with a user turn (model requirement)
            if messages and messages[0]["role"] != "user":
                messages.insert(0, {"role": "user", "content": ""})

            messages.insert(0, {"role": "system", "content": system_prompt + rag_block})

            try:
                params = json.loads((CONFIG_DIR / "params.json").read_text())
            except FileNotFoundError:
                params = {}  # no params.json, just use ollama's defaults

            async with chat_semaphore:
                response = await asyncio.to_thread(
                    ollama.chat,
                    model=chatmodel,
                    messages=messages,
                    options=params,
                )

            # strip <think> blocks in case the model ignores think=False, thank you claude
            bot_reply = re.sub(r'<think>.*?</think>', '', response["message"]["content"], flags=re.DOTALL).strip()

            print([{k: v for k, v in m.items() if k != "images"} for m in messages])  # skip the raw image bytes
            print(f"[{now_str()}] reply: {bot_reply}\n")

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

        except Exception as e:
            error_msg = f"[{now_str()}] 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"[{now_str()}] hosting on ws://0.0.0.0:{PORT}")
    async with serve(ollamachat, "0.0.0.0", PORT, ping_interval=None, ping_timeout=None, max_size=32 * 1024 * 1024):  # default 1mb would kill any media payload
        await asyncio.Future()


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