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

chatmodel = "gemini-3.5-flash-lite"
embedmodel = "gemini-embedding-001"
embeddim = 768  # default is 3072, the cosine loop below is pure python so smaller = way faster
RAG_THRESHOLD = 0.6  # 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))
RETRIES = 5  # how many times we retry a 503 (5s apart) before giving up
QUOTA_MSG = "usage limit reached, see you in a month"

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

# api key comes from env var or config/gemini_key.txt, whichever exists
GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "").strip()
if not GEMINI_KEY:
    try:
        GEMINI_KEY = (CONFIG_DIR / "gemini_key.txt").read_text(encoding="utf-8").strip()
    except FileNotFoundError:
        raise SystemExit("no gemini key — set GEMINI_API_KEY or put it in config/gemini_key.txt")

client = genai.Client(api_key=GEMINI_KEY)

# gemini free tier is rate limited, keeping this at one generation at a time
chat_semaphore = asyncio.Semaphore(1)

# bot is edgy, don't let gemini's default filters eat the replies
SAFETY_SETTINGS = [
    types.SafetySetting(category=c, threshold="BLOCK_NONE")
    for c in (
        "HARM_CATEGORY_HARASSMENT",
        "HARM_CATEGORY_HATE_SPEECH",
        "HARM_CATEGORY_SEXUALLY_EXPLICIT",
        "HARM_CATEGORY_DANGEROUS_CONTENT",
    )
]

# Each element in the VECTOR_DB will be a tuple (chunk, embedding)
VECTOR_DB = []
RAG_CHUNKS = []  # chunks from the last rebuild, so a partial (quota limited) rebuild can be finished without the bot resending everything
rag_loaded = False  # False until the bot sends longterm at least once, after a restart we ask for it

def embed(texts, task_type):
    # batches of 100 since that's the api limit per call
    out = []
    for i in range(0, len(texts), 100):
        res = client.models.embed_content(
            model=embedmodel,
            contents=texts[i:i + 100],
            config=types.EmbedContentConfig(task_type=task_type, output_dimensionality=embeddim),
        )
        out.extend(e.values for e in res.embeddings)
    return out

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 = embed([query], "RETRIEVAL_QUERY")[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
    missing = [c for c in dict.fromkeys(longterm_lines) if c not in existing]
    # free tier counts every chunk in a batch as a request (100/min), so only do 90 per rebuild
    # and catch up on the next messages instead of failing the whole reply
    batch = missing[:90]
    if batch:
        print(f"[{now_str()}] embedding {len(batch)} new chunks ({len(missing)} missing total)")
        try:
            for chunk, embedding in zip(batch, embed(batch, "RETRIEVAL_DOCUMENT")):
                existing[chunk] = embedding
        except Exception as e:
            print(f"[{now_str()}] embed failed, will retry next message: {e}")
    VECTOR_DB = [(chunk, existing[chunk]) for chunk in longterm_lines if chunk in existing]


async def geminichat(websocket):
    global RAG_CHUNKS, 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
    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)}

            # bot only sends longterm on a rebuild, if we never got it (just restarted) ask for it and the bot resends
            if payload.get("rebuild_vectors"):
                RAG_CHUNKS = make_chunks(longterm)
                rag_loaded = True
            elif not rag_loaded:
                await websocket.send(json.dumps({"need_longterm": True}))
                continue

            # rebuild vector db if bot signals it (every 150 interactions bot sends full longterm)
            # also keeps going while the db is still catching up after a partial (quota limited) rebuild
            if payload.get("rebuild_vectors") or len(VECTOR_DB) < len(RAG_CHUNKS):
                await asyncio.to_thread(rebuild_vector_db, RAG_CHUNKS)
                print(f"[{now_str()}] Vector DB reconciled, {len(VECTOR_DB)} chunks")

            user_query = nots(user_input)
            try:
                retrieved_knowledge = await asyncio.to_thread(retrieve, user_query) if VECTOR_DB else []
            except Exception as e:
                # quota hit on the query embed, just reply without rag this time
                print(f"[{now_str()}] retrieve failed, skipping rag: {e}")
                retrieved_knowledge = []

            # 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/model turns from shortterm
            # group consecutive non-bot lines into one user blob, bot lines become model 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": "model", "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})

            # make sure we start with a user turn (gemini rejects empty parts so no empty string here)
            if messages and messages[0]["role"] != "user":
                messages.insert(0, {"role": "user", "content": "..."})

            contents = [
                types.Content(role=m["role"], parts=[types.Part(text=m["content"])])
                for m in messages
            ]

            # media only goes on the final turn (history stays text), put it before the text like gemini recommends
            if media:
                media_parts = [
                    types.Part.from_bytes(data=base64.b64decode(m["data"]), mime_type=m["mime_type"])
                    for m in media
                ]
                contents[-1].parts = media_parts + contents[-1].parts

            config = types.GenerateContentConfig(
                system_instruction=system_prompt + rag_block,
                safety_settings=SAFETY_SETTINGS,
            )

            quota_hit = False
            async with chat_semaphore:
                for attempt in range(RETRIES + 1):
                    try:
                        response = await asyncio.to_thread(
                            client.models.generate_content,
                            model=chatmodel,
                            contents=contents,
                            config=config,
                        )
                        break
                    except errors.APIError as e:
                        if e.code == 429 and "PerMinute" not in str(e):  # out of quota for real, not just a rate limit blip
                            quota_hit = True
                            break
                        if (e.code == 503 or e.code == 429) and attempt < RETRIES:  # overloaded / rate limited, usually clears up in a few secs
                            print(f"[{now_str()}] gemini {e.code}, retrying the same request in 5s ({attempt + 1}/{RETRIES})")
                            await asyncio.sleep(5)  # sleeping with the semaphore held on purpose, the retry goes first
                            continue
                        raise

            if quota_hit:
                print(f"[{now_str()}] gemini is out of quota")
                await websocket.send(json.dumps({"error": QUOTA_MSG, "quota": True}))  # bot.py switches to its fallback on this
                continue

            # response.text is None when gemini blocks the reply anyway
            if not response.text:
                reason = response.prompt_feedback.block_reason if response.prompt_feedback else None
                if not reason and response.candidates:
                    reason = response.candidates[0].finish_reason
                if "PROHIBITED_CONTENT" in str(reason): 
                    await websocket.send(json.dumps({"ok": "google doesn't want me to reply to this :stocking:"}))
                    continue
                raise RuntimeError(f"gemini returned no text ({reason})")

            # strip <think> blocks just in case, thank you claude
            bot_reply = re.sub(r'<think>.*?</think>', '', response.text, flags=re.DOTALL).strip()

            print(messages)
            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 geminichat: {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(geminichat, "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())