"""PromptForge image library: tags, themes and NSFW flags for the images in the input folder.

A local vision model (a GGUF with its mmproj) looks at each image once and writes keywords, a one line caption,
how many people are in it and whether it is NSFW. The results live in the user directory
(`promptforge/image_tags.json`) and are used by the Load Image picker for search and blurring. A second job moves
images into theme folders (female, male, group, nsfw ...) and can be undone.

Everything here runs locally. Jobs are started from the picker or the sidebar tab, one at a time, in a thread.
"""
import asyncio
import base64
import gc
import importlib
import io
import json
import os
import re
import threading
import time
from pathlib import Path

from aiohttp import web
from PIL import Image, ImageOps
from server import PromptServer

import folder_paths
import node_helpers

EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tif", ".tiff", ".avif"}
THEMES = ("female", "male", "group", "people", "nsfw", "animals", "scenery", "objects", "other")
SUBJECT_THEMES = {"animal": "animals", "landscape": "scenery", "building": "scenery", "object": "objects", "vehicle": "objects", "food": "objects", "text": "other"}
MAX_SIDE = 768
FORCE_CPU = bool(os.environ.get("PROMPTFORGE_TAGGER_CPU"))

PROMPT = """You catalogue images for a personal library. Look at the image and answer with ONE JSON object and nothing else:
{"people": <number of people visible, 3 means three or more>, "gender": "female" | "male" | "mixed" | "none", "nsfw": true | false, "subject": "person" | "animal" | "landscape" | "building" | "object" | "vehicle" | "food" | "text" | "other", "style": "photo" | "illustration" | "anime" | "3d" | "screenshot" | "other", "tags": [10 to 15 short lowercase keywords: hair colour, clothing, setting, pose, mood, notable objects], "caption": "one plain sentence describing the image"}
"nsfw" is true when nudity, sexual content or explicit imagery is visible, in any style including drawings. "gender" is about the main person. Output the JSON only."""


# ---------------------------------------------------------------- stores

def store_dir():
    return Path(folder_paths.get_user_directory()) / "promptforge"


def _read(name):
    try:
        return json.loads((store_dir() / name).read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}


def _write(name, data):
    store_dir().mkdir(parents=True, exist_ok=True)
    (store_dir() / name).write_text(json.dumps(data, indent=1, ensure_ascii=False), encoding="utf-8")


descriptions = lambda: _read("image_descriptions.json")
save_descriptions = lambda data: _write("image_descriptions.json", data)
tags = lambda: _read("image_tags.json")
save_tags = lambda data: _write("image_tags.json", data)
_store_lock = threading.Lock()


def input_root():
    return Path(os.path.abspath(folder_paths.get_input_directory()))


def resolve(name):
    """Absolute path of `name` inside the input folder; raises when it points elsewhere or is not an image."""
    root = input_root()
    path = Path(os.path.abspath(root / name))
    if root not in path.parents or path.suffix.lower() not in EXTENSIONS:
        raise ValueError(f"[PromptForge] {name} is not an image inside the input folder.")
    return path


def scan():
    root = str(input_root())
    found = []
    for folder, dirs, files in os.walk(root, followlinks=True):
        dirs[:] = sorted(d for d in dirs if not d.startswith("."))
        for file in files:
            if os.path.splitext(file)[1].lower() in EXTENSIONS:
                full = os.path.join(folder, file)
                found.append({"name": os.path.relpath(full, root).replace(os.sep, "/"), "mtime": os.path.getmtime(full)})
    return found


def theme_of(result):
    if result.get("nsfw"):
        return "nsfw"
    people = int(result.get("people") or 0)
    if people >= 2:
        return "group"
    if people == 1:
        return {"female": "female", "male": "male"}.get(result.get("gender"), "people")
    return SUBJECT_THEMES.get(result.get("subject"), "other")


# ---------------------------------------------------------------- vision model

def _package():
    return importlib.import_module(__package__)


VISION_FAMILIES = ("vl", "qwen3.5", "qwen3.6", "qwen35", "gemma-3", "gemma3", "gemma-4", "gemma4", "llava", "minicpm")


def _is_vision(path, mmproj):
    """An mmproj named after the model always counts; a generic one only for model families that have vision."""
    stem = re.split(r"[-_.](?:q\d|f16|bf16|f32|iq\d)", path.stem.lower())[0]
    if len(stem) >= 8 and mmproj.name.lower().startswith(stem):
        return True
    return any(family in path.name.lower() for family in VISION_FAMILIES)


def vision_models():
    """GGUF models on this machine that have a matching mmproj next to them: [{"id", "label"}]."""
    package = _package()
    models = []
    for entry in package.LOCAL_TEXT_MODELS:
        path = Path(entry["path"])
        mmproj = package._local_mmproj(path, None) if entry["kind"] == "gguf" and path.is_file() else None
        if mmproj is not None and _is_vision(path, mmproj):
            models.append({"id": str(path), "label": path.name})
    return models


class Tagger:
    def __init__(self, model_path):
        package = _package()
        if model_path not in {m["id"] for m in vision_models()}:
            raise ValueError("[PromptForge] pick one of the listed vision models.")
        from llama_cpp import Llama, llama_chat_format
        path = Path(model_path)
        mmproj = package._local_mmproj(path, None)
        handler = getattr(llama_chat_format, package._gguf_vision_handler(path.name), None)
        if handler is None:
            raise ValueError("[PromptForge] this llama-cpp-python build cannot run this vision model.")
        package._free_vram()
        layers = 0 if FORCE_CPU else -1
        if layers and package._gguf_memory_estimate(path, 4096, layers, mmproj) > _free_bytes():
            layers = package._gguf_layers_that_fit(path, 4096, mmproj)
        handler_instance = handler(clip_model_path=str(mmproj), verbose=False)
        for gpu_layers, kv_on_gpu in ((layers, layers != 0), (0, False)):
            try:
                self.llm = Llama(model_path=str(path), chat_handler=handler_instance, n_ctx=4096, n_gpu_layers=gpu_layers, offload_kqv=kv_on_gpu, verbose=False)
                break
            except (ValueError, RuntimeError):
                if gpu_layers == 0:
                    raise
                print("[PromptForge] the vision model does not fit the free VRAM; running it on the CPU.")
        self.split_thinking = package._split_thinking

    def ask(self, path):
        with node_helpers.pillow(Image.open, str(path)) as opened:
            image = node_helpers.pillow(ImageOps.exif_transpose, opened).convert("RGB")
        image.thumbnail((MAX_SIDE, MAX_SIDE))
        buffer = io.BytesIO()
        image.save(buffer, "JPEG", quality=88)
        url = "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode()
        reply = self.llm.create_chat_completion(
            messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}}, {"type": "text", "text": PROMPT}]}],
            max_tokens=500, temperature=0.0)
        return self.split_thinking(reply["choices"][0]["message"]["content"] or "")[0]

    def close(self):
        self.llm.close()
        del self.llm
        gc.collect()
        _package()._free_vram()


def _free_bytes():
    import torch
    try:
        return torch.cuda.mem_get_info()[0]
    except (RuntimeError, AssertionError):
        return 0


def parse(text):
    """The catalogue entry from the model's answer, or None when it is not usable."""
    start, end = text.find("{"), text.rfind("}")
    if start < 0 or end <= start:
        return None
    try:
        data = json.loads(text[start:end + 1])
    except ValueError:
        return None
    if not isinstance(data, dict) or "nsfw" not in data:
        return None
    words = data.get("tags") if isinstance(data.get("tags"), list) else str(data.get("tags", "")).split(",")
    keywords = list(dict.fromkeys(w.strip().lower()[:40] for w in map(str, words) if w.strip()))[:20]
    nsfw = data["nsfw"] if isinstance(data["nsfw"], bool) else str(data["nsfw"]).lower() == "true"
    try:
        people = max(0, int(data.get("people") or 0))
    except (TypeError, ValueError):
        people = 0
    entry = {"people": people, "gender": str(data.get("gender", "none")).lower(), "nsfw": nsfw, "subject": str(data.get("subject", "other")).lower(),
             "style": str(data.get("style", "")).lower(), "tags": keywords, "caption": str(data.get("caption", "")).strip()[:400]}
    entry["theme"] = theme_of(entry)
    return entry


# ---------------------------------------------------------------- jobs

_job = {"running": False, "kind": "", "total": 0, "done": 0, "current": "", "failed": [], "message": "", "cancel": False}
_job_lock = threading.Lock()


def _set(**values):
    with _job_lock:
        _job.update(values)


def _tag_worker(model, names):
    tagger = None
    try:
        _set(message="Loading the vision model...")
        tagger = Tagger(model)
        store = tags()
        for index, name in enumerate(names):
            if _job["cancel"]:
                break
            _set(done=index, current=name, message="Tagging")
            try:
                path = resolve(name)
                text = tagger.ask(path)
                entry = parse(text) or parse(tagger.ask(path))
                if entry is None:
                    raise ValueError("no usable answer")
            except Exception as error:
                with _job_lock:
                    _job["failed"].append(f"{name}: {str(error)[:80]}")
                continue
            entry.update({"mtime": path.stat().st_mtime, "model": Path(model).name, "tagged_at": int(time.time())})
            with _store_lock:
                store = tags()
                store[name] = entry
                save_tags(store)
        _set(done=len(names), message="Cancelled" if _job["cancel"] else "Done")
    except Exception as error:
        _set(message=f"Failed: {error}")
    finally:
        if tagger is not None:
            tagger.close()
        _set(running=False, current="")


def _start(kind, target, *args):
    with _job_lock:
        if _job["running"]:
            return False
        _job.update(running=True, kind=kind, total=len(args[1]), done=0, current="", failed=[], message="Starting", cancel=False)
    threading.Thread(target=target, args=args, daemon=True).start()
    return True


def _targets(names, retag):
    known = tags()
    root = input_root()
    chosen = names if isinstance(names, list) else [image["name"] for image in scan()]
    result = []
    for name in chosen:
        try:
            path = resolve(name)
            current = known.get(name)
            if retag or current is None or abs(current.get("mtime", 0) - path.stat().st_mtime) > 1:
                result.append(path.relative_to(root).as_posix())
        except (ValueError, OSError):
            continue
    return result


def _plan(names, include_subfolders):
    """[(name, theme)] of tagged images that are not already in their theme folder."""
    known = tags()
    chosen = names if isinstance(names, list) else [image["name"] for image in scan()]
    plan = []
    for name in chosen:
        entry = known.get(name)
        if not entry or not (include_subfolders or "/" not in name):
            continue
        if name.split("/")[0] != entry["theme"]:
            plan.append((name, entry["theme"]))
    return plan


def _unique(target):
    if not target.exists():
        return target
    for number in range(2, 10000):
        candidate = target.with_name(f"{target.stem}_{number}{target.suffix}")
        if not candidate.exists():
            return candidate
    raise ValueError("[PromptForge] too many files with the same name.")


def _organize(plan):
    root = input_root()
    moves = []
    with _store_lock:
        store, notes = tags(), descriptions()
        for name, theme in plan:
            source = resolve(name)
            if not source.is_file():
                continue
            target = _unique(Path(os.path.abspath(root / theme / source.name)))
            target.parent.mkdir(parents=True, exist_ok=True)
            source.rename(target)
            new = target.relative_to(root).as_posix()
            for table in (store, notes):
                if name in table:
                    table[new] = table.pop(name)
            moves.append([name, new])
        save_tags(store)
        save_descriptions(notes)
        _write("organize_log.json", {"moves": moves, "time": int(time.time())})
    return moves


def _undo():
    log = _read("organize_log.json").get("moves", [])
    root = input_root()
    restored = 0
    with _store_lock:
        store, notes = tags(), descriptions()
        for old, new in reversed(log):
            source, target = Path(os.path.abspath(root / new)), Path(os.path.abspath(root / old))
            if not source.is_file() or target.exists():
                continue
            target.parent.mkdir(parents=True, exist_ok=True)
            source.rename(target)
            for table in (store, notes):
                if new in table:
                    table[old] = table.pop(new)
            restored += 1
            try:
                source.parent.rmdir()
            except OSError:
                pass
        save_tags(store)
        save_descriptions(notes)
        _write("organize_log.json", {"moves": [], "time": int(time.time())})
    return restored


# ---------------------------------------------------------------- routes

routes = PromptServer.instance.routes


@routes.get("/promptforge/library/models")
async def _models_route(request):
    return web.json_response({"models": vision_models()})


@routes.get("/promptforge/library/status")
async def _status_route(request):
    with _job_lock:
        return web.json_response({key: value for key, value in _job.items() if key != "cancel"})


@routes.get("/promptforge/library/info")
async def _info_route(request):
    return web.json_response(tags().get(request.query.get("name", ""), {}))


@routes.post("/promptforge/library/tag")
async def _tag_route(request):
    body = await request.json()
    names = _targets(body.get("names") if isinstance(body.get("names"), list) else "all", bool(body.get("retag")))
    if not names:
        return web.json_response({"started": False, "message": "Nothing to tag."})
    try:
        started = _start("tag", _tag_worker, str(body.get("model", "")), names)
    except ValueError as error:
        return web.json_response({"started": False, "message": str(error)}, status=400)
    return web.json_response({"started": started, "count": len(names), "message": "" if started else "A job is already running."})


@routes.post("/promptforge/library/cancel")
async def _cancel_route(request):
    _set(cancel=True)
    return web.json_response({})


@routes.post("/promptforge/library/organize")
async def _organize_route(request):
    body = await request.json()
    with _job_lock:
        if _job["running"]:
            return web.json_response({"error": "A job is running."}, status=409)
    plan = _plan(body.get("names") if isinstance(body.get("names"), list) else "all", bool(body.get("include_subfolders")))
    if body.get("dry_run", True):
        counts = {}
        for _, theme in plan:
            counts[theme] = counts.get(theme, 0) + 1
        return web.json_response({"counts": counts, "total": len(plan), "sample": [[a, b] for a, b in plan[:6]]})
    moves = await asyncio.get_running_loop().run_in_executor(None, _organize, plan)
    return web.json_response({"moved": len(moves)})


@routes.post("/promptforge/library/undo")
async def _undo_route(request):
    return web.json_response({"restored": await asyncio.get_running_loop().run_in_executor(None, _undo)})


@routes.post("/promptforge/library/edit")
async def _edit_route(request):
    body = await request.json()
    name = str(body.get("name", ""))
    try:
        resolve(name)
    except ValueError as error:
        return web.json_response({"error": str(error)}, status=400)
    with _store_lock:
        store = tags()
        entry = store.setdefault(name, {"people": 0, "gender": "none", "subject": "other", "style": "", "tags": [], "caption": "", "model": "manual"})
        if "nsfw" in body:
            entry["nsfw"] = bool(body["nsfw"])
        entry.setdefault("nsfw", False)
        if "tags" in body:
            entry["tags"] = list(dict.fromkeys(w.strip().lower() for w in str(body["tags"]).split(",") if w.strip()))[:30]
        if "caption" in body:
            entry["caption"] = str(body["caption"])[:400]
        if body.get("theme") in THEMES:
            entry["theme"] = body["theme"]
        else:
            entry["theme"] = theme_of(entry)
        entry["mtime"] = resolve(name).stat().st_mtime
        save_tags(store)
    return web.json_response(entry)