#!/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"

            # 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
            messages.append({"role": "user", "content": f"[responding to]: {nots(user_input)}\n(that is the newest message and the only one you answer, the lines above are just context. only the name that starts it is speaking to you, everyone else is someone else)"})

            # 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())