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

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