#!/usr/bin/env python3
"""
fills in the empty descriptions in files/config/reactions.json by showing every sticker / emoji / media to the
local gemma-4 on ollama. run it on the box with the files, ollama has to be up (main.py doesn't).
stickers and emojis are downloaded from discord with the bot token (login only, the bot doesn't need to be running).

    describe.py                  only entries with no description yet
    describe.py --all            redo everything (overwrites the ones you wrote by hand too)
    describe.py --model <name>   ollama model, defaults to chatmodel from main.py (it has to be able to see images)
    describe.py --gemini         use the gemini api instead (chatmodel from gemini.py, key from GEMINI_API_KEY or config/gemini_key.txt)
"""
import argparse
import asyncio
import io
import json
import os
import re
import shutil
import subprocess
import tempfile
import time
from pathlib import Path

from PIL import Image

SCRIPT_DIR = Path(__file__).resolve().parent
CONFIG_DIR = SCRIPT_DIR / "files" / "config"
MEDIA_DIR = SCRIPT_DIR / "files" / "media"
REACTIONS_FILE = CONFIG_DIR / "reactions.json"
VIDEO_EXTS = (".mp4", ".webm", ".mov", ".mkv")

PROMPT = ("this is a {kind} from a discord server, its name is \"{hint}\" (the name is a hint, it may or may not help). "
          "describe what is literally in the image in one short line, under 15 words: the character or object (name the character if you recognize them, "
          "the name above can help), what they're doing, and any text written in the image. "
          "don't guess when to use it, no preamble, no quotes.")

# only for media, the bot sends [media: name] so the file name is what she sees, a hash tells her nothing
NAME_PROMPT = (" answer in exactly two lines. line 1: a short file name for it, 1 to 3 lowercase words joined by underscores, no extension "
               "(the character's name if you recognize them, otherwise the main subject). line 2: the description.")


def model_from(script: str) -> str:
    # same trick as colab_setup.py, read the model name straight out of main.py / gemini.py so they can't drift apart
    text = (SCRIPT_DIR / script).read_text(encoding="utf-8")
    return re.search(r'^chatmodel\s*=\s*["\'](.+?)["\']', text, re.M).group(1)


def read_lines(path: Path) -> list[str]:
    return [l.strip() for l in path.read_text(encoding="utf-8").splitlines() if l.strip() and not l.strip().startswith("#")]


async def fetch_discord(wanted: dict) -> dict:
    # {(kind, name): bytes}. login only, no gateway, so this works while the bot is running too
    import discord
    token, guild_id = read_lines(CONFIG_DIR / "tokens.txt")[:2]
    found = {}
    client = discord.Client(intents=discord.Intents.none())
    await client.login(token)
    try:
        guild = await client.fetch_guild(int(guild_id))
        for kind, items in (("stickers", await guild.fetch_stickers()), ("emojis", await guild.fetch_emojis())):
            for item in items:
                if item.name not in wanted[kind] or not item.available:  # locked ones (boost level dropped) can't be sent anyway
                    continue
                if kind == "stickers" and item.format == discord.StickerFormatType.lottie:  # json animation, nothing to look at
                    continue
                found[(kind, item.name)] = await item.read()
    finally:
        await client.close()
    return found


def to_png(data: bytes, dest: Path) -> bool:
    # first frame on a white background (transparent stickers come out black otherwise), small so gemma doesn't chew on it
    try:
        img = Image.open(io.BytesIO(data))
        img.seek(getattr(img, "n_frames", 1) // 2)  # middle frame, the first one of an animation is often blank or half drawn
        img = img.convert("RGBA")
    except Exception:
        return False
    bg = Image.new("RGBA", img.size, "white")
    bg.alpha_composite(img)
    bg = bg.convert("RGB")
    bg.thumbnail((512, 512))
    bg.save(dest, "PNG")
    return True


def video_frame(path: Path, dest: Path) -> bool:
    if not shutil.which("ffmpeg"):
        return False
    r = subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(path), "-frames:v", "1", "-vf", "thumbnail,scale=512:-1", str(dest)])
    return r.returncode == 0 and dest.exists()


def describe(model: str, png: Path, prompt: str) -> str:
    import ollama
    r = ollama.chat(model=model, messages=[{"role": "user", "content": prompt, "images": [str(png)]}], options={"temperature": 0.2})
    return strip_think(r["message"]["content"])


def describe_gemini(model: str, png: Path, prompt: str) -> str:
    from google import genai
    from google.genai import types
    key = os.environ.get("GEMINI_API_KEY", "").strip() or (SCRIPT_DIR / "config" / "gemini_key.txt").read_text(encoding="utf-8").strip()
    # a sticker with a slur in it would get blocked by the default filters, same reason gemini.py turns them off
    safety = [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")]
    client = genai.Client(api_key=key)  # has to stay in a variable, as a temporary it gets garbage collected mid call and closes the connection
    r = client.models.generate_content(
        model=model,
        contents=[types.Part.from_bytes(data=png.read_bytes(), mime_type="image/png"), prompt],
        config=types.GenerateContentConfig(safety_settings=safety, temperature=0.2),
    )
    time.sleep(4)  # free tier rate limit, rerun the script if some fail, it only does the empty ones
    return strip_think(r.text or "")


def strip_think(text: str) -> str:
    return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)


def clean(text: str) -> str:
    return " ".join(strip_think(text).split()).strip(' "')


def split_name(text: str) -> tuple[str, str]:
    # media answers come back as "name\ndescription", no name (model ignored the format) = keep the file as it is
    lines = [l for l in strip_think(text).splitlines() if l.strip()]
    if len(lines) < 2:
        return "", clean(text)
    name = re.sub(r"[^a-z0-9]+", "_", lines[0].lower()).strip("_")[:40]
    return name, clean(" ".join(lines[1:]))


def rename_media(old: str, name: str) -> str:
    # keeps the extension, adds _2 _3 if the name is taken
    suffix = Path(old).suffix
    new = f"{name}{suffix}"
    n = 2
    while (MEDIA_DIR / new).exists() and new != old:
        new = f"{name}_{n}{suffix}"
        n += 1
    if new == old:
        return old
    (MEDIA_DIR / old).rename(MEDIA_DIR / new)
    return new


def save_description(kind: str, name: str, text: str, force: bool, new_name: str = None):
    # read it fresh every time, the bot adds new entries to this file while we run
    data = json.loads(REACTIONS_FILE.read_text(encoding="utf-8"))
    if name in data.get(kind, {}) and (force or not data[kind][name]):
        if new_name and new_name != name:
            data[kind][new_name] = data[kind].pop(name)
            name = new_name
        data[kind][name] = text
        REACTIONS_FILE.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--all", action="store_true")
    ap.add_argument("--model")
    ap.add_argument("--gemini", action="store_true")
    args = ap.parse_args()
    model = args.model or model_from("gemini.py" if args.gemini else "main.py")

    if not REACTIONS_FILE.exists():
        raise SystemExit("no files/config/reactions.json yet, start bot.py once so it lists your stickers/emojis/media")
    data = json.loads(REACTIONS_FILE.read_text(encoding="utf-8"))
    todo = {k: [n for n, d in data.get(k, {}).items() if args.all or not d] for k in ("stickers", "emojis", "media")}
    if not any(todo.values()):
        raise SystemExit("nothing without a description")

    remote = asyncio.run(fetch_discord(todo)) if todo["stickers"] or todo["emojis"] else {}

    with tempfile.TemporaryDirectory() as tmp:
        png = Path(tmp) / "thing.png"
        for kind in ("stickers", "emojis", "media"):
            for name in todo[kind]:
                png.unlink(missing_ok=True)
                if kind == "media":
                    path = MEDIA_DIR / name
                    if not path.is_file():
                        print(f"skip media/{name}: file is gone")
                        continue
                    ok = video_frame(path, png) if path.suffix.lower() in VIDEO_EXTS else to_png(path.read_bytes(), png)
                else:
                    ok = (kind, name) in remote and to_png(remote[(kind, name)], png)
                if not ok:
                    print(f"skip {kind}/{name}: couldn't get an image out of it")
                    continue
                try:
                    hint = re.sub(r"^\d+_", "", Path(name).stem) if kind == "media" else name  # media is saved as id_filename
                    prompt = PROMPT.format(kind=kind.rstrip("s") if kind != "media" else "media file", hint=hint)
                    if kind == "media":
                        prompt += NAME_PROMPT
                    text = (describe_gemini if args.gemini else describe)(model, png, prompt)
                except Exception as e:
                    print(f"skip {kind}/{name}: describe failed: {e}")
                    continue
                new_name = None
                if kind == "media":
                    new_name, text = split_name(text)
                else:
                    text = clean(text)
                if not text:
                    continue
                if new_name:
                    new_name = rename_media(name, new_name)
                save_description(kind, name, text, args.all, new_name)
                print(f"{kind}/{name}{f' -> {new_name}' if new_name and new_name != name else ''}: {text}")


if __name__ == "__main__":
    main()