core/cloud_storage.py

import io
import json
import os
import re
import atexit
import threading
import queue
import time
from datetime import datetime, timezone
from urllib.parse import unquote

import boto3
from botocore.exceptions import ClientError
from PIL import Image

FLUSH_INTERVAL_SEC = 15
FLUSH_MAX_ITEMS = 4
NUMBER_WIDTH = 4
INDEX_KEY = "_meta/index.json"
FAVS_KEY = "_meta/favorites.json"

_q = queue.Queue()
_lock = threading.RLock()
_worker_started = False
_counters = {}
_logs = {}


def _safe(name):
    name = str(name or "unknown").strip()
    name = re.sub(r"[^\w.\-+ ]", "_", name)
    name = re.sub(r"\s+", "_", name)
    return name.strip("._ ")[:100] or "unknown"


def _config_ready():
    return all(os.environ.get(k) for k in ("R2_ENDPOINT", "R2_KEY", "R2_SECRET"))


def is_enabled():
    return _config_ready()


def _client():
    return boto3.client(
        "s3",
        endpoint_url=os.environ["R2_ENDPOINT"].rstrip("/"),
        aws_access_key_id=os.environ["R2_KEY"],
        aws_secret_access_key=os.environ["R2_SECRET"],
        region_name="auto",
    )


def _bucket():
    return os.environ.get("R2_BUCKET", "imagegen")


def _read_text(c, key):
    try:
        obj = c.get_object(Bucket=_bucket(), Key=key)
        return obj["Body"].read().decode("utf-8", errors="replace")
    except ClientError as e:
        if e.response.get("Error", {}).get("Code") in ("NoSuchKey", "404", "NotFound"):
            return ""
        raise


def _read_index(c):
    raw = _read_text(c, INDEX_KEY)
    if not raw:
        return []
    try:
        data = json.loads(raw)
        return data if isinstance(data, list) else []
    except Exception:
        print("⚠️ [R2Storage] _meta/index.json is not valid JSON; keeping a fresh index.")
        return []


def _write_json(c, key, value):
    c.put_object(
        Bucket=_bucket(), Key=key,
        Body=json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8"),
        ContentType="application/json; charset=utf-8",
    )


def _init_folder(c, folder):
    if folder in _counters:
        return
    last = 0
    log_text = _read_text(c, f"{folder}/log.txt")
    nums = re.findall(r"^#(\d+)\b", log_text, flags=re.M)
    if nums:
        last = max(int(n) for n in nums)
    try:
        paginator = c.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=_bucket(), Prefix=f"{folder}/"):
            for obj in page.get("Contents", []):
                match = re.match(rf"^{re.escape(folder)}/(\d+)\.png$", obj["Key"])
                if match:
                    last = max(last, int(match.group(1)))
    except Exception as e:
        print(f"⚠️ [R2Storage] Could not inspect folder {folder}: {e}")
    _counters[folder] = last
    _logs[folder] = log_text


def _split_params(params):
    params = (params or "").strip()
    settings_text = ""
    idx = params.rfind("\nSteps: ")
    if idx != -1:
        settings_text = params[idx + 1:].strip()
        params = params[:idx]
    negative = ""
    nidx = params.find("\nNegative prompt: ")
    if nidx != -1:
        negative = params[nidx + len("\nNegative prompt: "):].strip()
        params = params[:nidx]
    return params.strip(), negative, settings_text


def _parse_settings(s):
    out = {}
    # Common A1111-style values. This deliberately keeps parsing conservative.
    patterns = [
        ("Steps", r"(?:^|,\s*)Steps:\s*([^,]+)"),
        ("Sampler", r"(?:^|,\s*)Sampler:\s*([^,]+)"),
        ("Scheduler", r"(?:^|,\s*)Scheduler:\s*([^,]+)"),
        ("CFG scale", r"(?:^|,\s*)CFG scale:\s*([^,]+)"),
        ("Seed", r"(?:^|,\s*)Seed:\s*([^,]+)"),
        ("Size", r"(?:^|,\s*)Size:\s*([^,]+)"),
        ("Base Model", r"(?:^|,\s*)Base Model:\s*(.*?)(?=,\s*[A-Z][^,:]*:|$)"),
    ]
    for key, pattern in patterns:
        m = re.search(pattern, s or "")
        if m:
            out[key] = m.group(1).strip()
    return out


def _webp(data, max_side=None, quality=90):
    with Image.open(io.BytesIO(data)) as im:
        im = im.convert("RGB")
        if max_side:
            im.thumbnail((max_side, max_side))
        buf = io.BytesIO()
        im.save(buf, "WEBP", quality=quality, method=4)
        return buf.getvalue(), im.size


def _flush(items):
    if not items:
        return
    c = _client()
    touched = set()
    new_records = []
    stamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
    with _lock:
        for it in items:
            folder = it["folder"]
            _init_folder(c, folder)
            _counters[folder] += 1
            num = _counters[folder]
            filename = f"{num:0{NUMBER_WIDTH}d}.png"
            rel_path = f"{folder}/{filename}"
            base = rel_path[:-4]
            thumb_key = f"thumbs/{base}.webp"
            full_key = f"full/{base}.webp"
            prompt, negative, settings_text = _split_params(it.get("params", ""))
            settings = _parse_settings(settings_text)
            try:
                thumb, dimensions = _webp(it["data"], 512, 80)
                full, _ = _webp(it["data"], None, 92)
                c.put_object(Bucket=_bucket(), Key=rel_path, Body=it["data"], ContentType="image/png")
                c.put_object(Bucket=_bucket(), Key=thumb_key, Body=thumb, ContentType="image/webp")
                c.put_object(Bucket=_bucket(), Key=full_key, Body=full, ContentType="image/webp")

                entry = (
                    f"#{num:0{NUMBER_WIDTH}d} | {stamp} | {filename}\n"
                    f"Prompt: {prompt}\n"
                    f"Negative prompt: {negative}\n"
                    f"Settings: {settings_text}\n"
                    f"{'-' * 60}\n\n"
                )
                _logs[folder] += entry
                touched.add(folder)

                model = it["model"]
                family = it["family"]
                size = settings.get("Size") or f"{dimensions[0]}x{dimensions[1]}"
                seed = settings.get("Seed", "")
                steps = settings.get("Steps", "")
                cfg = settings.get("CFG scale", "")
                try: steps = int(float(steps))
                except (TypeError, ValueError): steps = 0
                try: cfg = float(cfg)
                except (TypeError, ValueError): cfg = 0
                orient = "مربع" if dimensions[0] == dimensions[1] else ("أفقي" if dimensions[0] > dimensions[1] else "عمودي")
                record = {
                    "id": rel_path,
                    "path": rel_path,
                    "t": thumb_key,
                    "f": full_key,
                    "prompt": prompt,
                    "negative": negative,
                    "family": family,
                    "model": model,
                    "file": filename,
                    "time": datetime.now(timezone.utc).isoformat(),
                    "seed": seed,
                    "size": size,
                    "steps": steps,
                    "sampler": settings.get("Sampler", ""),
                    "scheduler": settings.get("Scheduler", ""),
                    "cfg": cfg,
                    "base_model": settings.get("Base Model", model),
                    "orient": orient,
                    "settings": settings,
                }
                new_records.append(record)
            except Exception as e:
                print(f"⚠️ [R2Storage] Failed uploading {rel_path}: {type(e).__name__}: {e}")

        for folder in touched:
            c.put_object(
                Bucket=_bucket(), Key=f"{folder}/log.txt",
                Body=_logs[folder].encode("utf-8"),
                ContentType="text/plain; charset=utf-8",
            )

        if new_records:
            index = _read_index(c)
            existing_ids = {r.get("id") for r in index if isinstance(r, dict)}
            index.extend(r for r in new_records if r["id"] not in existing_ids)
            _write_json(c, INDEX_KEY, index)

    if new_records:
        print(f"☁️ [R2Storage] Saved {len(new_records)} image(s) directly to R2 bucket '{_bucket()}'.")


def _worker():
    pending = []
    last_flush = time.time()
    while True:
        try:
            item = _q.get(timeout=2)
            if item is None:
                _flush(pending)
                return
            pending.append(item)
        except queue.Empty:
            pass
        due = time.time() - last_flush >= FLUSH_INTERVAL_SEC
        if pending and (due or len(pending) >= FLUSH_MAX_ITEMS):
            batch, pending = pending, []
            try:
                _flush(batch)
            except Exception as e:
                print(f"⚠️ [R2Storage] Flush error: {type(e).__name__}: {e}")
            last_flush = time.time()


def _shutdown():
    if _worker_started:
        _q.put(None)
        time.sleep(1)


def _ensure_worker():
    global _worker_started
    if _worker_started:
        return
    with _lock:
        if _worker_started:
            return
        threading.Thread(target=_worker, daemon=True, name="r2-storage").start()
        atexit.register(_shutdown)
        _worker_started = True


def queue_image(architecture, model, image_path, params):
    """Queue a generated image for direct upload to R2; never raises into generation."""
    try:
        if not is_enabled():
            print("ℹ️ [R2Storage] Skipped: R2_ENDPOINT/R2_KEY/R2_SECRET are not all configured.")
            return
        with open(image_path, "rb") as f:
            data = f.read()
        _ensure_worker()
        _q.put({
            "folder": f"{_safe(architecture)}/{_safe(model)}",
            "family": _safe(architecture),
            "model": _safe(model),
            "data": data,
            "params": params or "",
        })
    except Exception as e:
        print(f"⚠️ [R2Storage] Could not queue image: {type(e).__name__}: {e}")