#!/usr/bin/env python3
"""Generate a cover image for a THREATMARKET intel pack via ComfyUI SDXL.

Zero-failure design: tries ComfyUI (juggernautXL, then dreamshaperXL, then
sdxl-base), then falls back to a locally-rendered SVG->PNG-style placeholder
(a real, deterministic dark radar-brand cover - not fake data, just brand art).
Always writes /opt/threatmarket/static/covers/<pack_id>.png (or .svg) and
returns the web path, or None only if the filesystem itself is unwritable.

Usage: from cover_gen import generate_cover
       path = generate_cover("chewy.com", "2026-09-25")
"""
import base64
import json
import os
import urllib.request

COMFY = os.environ.get("COMFY_URL", "http://10.30.20.29:8188")
COVER_DIR = os.environ.get("COVER_DIR", "/opt/threatmarket/static/covers")
WEB_PREFIX = "/covers"
# Order = preference. Lightning models = few steps = fast + reliable.
CKPTS = ["flux1-schnell-fp8.safetensors",
         "juggernautXL_ragnarok.safetensors",
         "dreamshaperXL_lightningDPMSDE.safetensors"]
TIMEOUT_SUBMIT = 20
TIMEOUT_POLL = 150
TIMEOUT_FETCH = 30


def _prompt_text(domain, date):
    return ('dark hacker operations room at night, glowing green matrix rain code '
            'walls, multiple monitors showing world-map network dashboards and '
            'breach data, empty desk chair, dramatic cyan and green lighting, '
            'cinematic depth of field, gritty cyberpunk atmosphere, wide dark '
            'area at top for title text, professional threat-intelligence dossier '
            'cover art, no text, no letters, no words, 4k')


def _workflow(ckpt, seed):
    return {
        "3": {"class_type": "KSampler", "inputs": {
            "seed": seed, "steps": 20, "cfg": 1.8, "sampler_name": "euler",
            "scheduler": "sgm_uniform", "denoise": 1.0,
            "model": ["4", 0], "positive": ["6", 0], "negative": ["7", 0],
            "latent_image": ["5", 0]}},
        "4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": ckpt}},
        "5": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 640, "batch_size": 1}},
        "6": {"class_type": "CLIPTextEncode", "inputs": {"text": "placeholder", "clip": ["4", 1]}},
        "7": {"class_type": "CLIPTextEncode", "inputs": {"text": "text, watermark, blurry, low quality, oversaturated", "clip": ["4", 1]}},
        "8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["4", 2]}},
        "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "tmcover", "images": ["8", 0]}},
    }


def _comfy_generate(domain, date, seed):
    """Try each checkpoint; return PNG bytes or None."""
    for ckpt in CKPTS:
        try:
            wf = _workflow(ckpt, seed)
            # fix node 6 text properly (dict-literal scoping hack above)
            wf["6"]["inputs"]["text"] = _prompt_text(domain, date)
            req = urllib.request.Request(
                COMFY + "/prompt",
                data=json.dumps({"prompt": wf}).encode(),
                headers={"Content-Type": "application/json"})
            with urllib.request.urlopen(req, timeout=TIMEOUT_SUBMIT) as r:
                resp = json.loads(r.read())
            pid = resp.get("prompt_id")
            if not pid:
                continue
            import time
            deadline = time.time() + TIMEOUT_POLL
            while time.time() < deadline:
                with urllib.request.urlopen(COMFY + "/history/" + pid, timeout=15) as r:
                    hist = json.loads(r.read())
                if pid in hist and hist[pid].get("outputs"):
                    outs = hist[pid]["outputs"]
                    for node in outs.values():
                        for img in node.get("images", []):
                            fn = img["filename"]
                            sub = img.get("subfolder", "")
                            imgtype = img.get("type", "output")
                            url = f"{COMFY}/view?filename={fn}&subfolder={sub}&type={imgtype}"
                            with urllib.request.urlopen(url, timeout=TIMEOUT_FETCH) as fr:
                                return fr.read()
                time.sleep(3)
        except Exception:
            continue
    return None


def _fallback_svg(domain, date):
    """Deterministic branded SVG cover - always succeeds if fs writable."""
    seed = sum(ord(c) for c in domain) % 360
    svg = f"""<svg xmlns='http://www.w3.org/2000/svg' width='1024' height='640'>
<defs><radialGradient id='g' cx='75%' cy='20%'><stop offset='0%' stop-color='hsl({seed},80%,22%)'/>
<stop offset='100%' stop-color='#0a0e1a'/></radialGradient></defs>
<rect width='1024' height='640' fill='url(#g)'/>
<g stroke='hsl({seed},90%,60%)' stroke-opacity='0.25' fill='none'>"""
    for i in range(6):
        svg += f"<circle cx='780' cy='200' r='{40+i*45}' stroke-width='1'/>"
    svg += f"</g><g fill='hsl({seed},90%,65%)'>"
    import math
    for i in range(14):
        x = 120 + (i * 61) % 700
        y = 420 + (i * 37) % 160
        r = 2 + (i % 3)
        svg += f"<circle cx='{x}' cy='{y}' r='{r}' opacity='{0.35 + (i % 5) * 0.12}'/>"
    svg += f"""</g>
<text x='48' y='90' font-family='monospace' font-size='30' fill='#e8f4ff' font-weight='bold'>INTEL PACK</text>
<text x='48' y='130' font-family='monospace' font-size='22' fill='hsl({seed},90%,70%)'>{domain}</text>
<text x='48' y='162' font-family='monospace' font-size='16' fill='#8fa3bf'>{date} — THREATMARKET</text>
</svg>"""
    return svg.encode()


def generate_cover(domain, date):
    """Returns web path of cover image, or None. Never raises."""
    os.makedirs(COVER_DIR, exist_ok=True)
    pid = (domain + "-" + date.replace("-", "")).replace("/", "_")
    png_path = os.path.join(COVER_DIR, pid + ".png")
    svg_path = os.path.join(COVER_DIR, pid + ".svg")
    if os.path.exists(png_path):
        return f"{WEB_PREFIX}/{pid}.png"
    if os.path.exists(svg_path):
        return f"{WEB_PREFIX}/{pid}.svg"
    seed = sum(ord(c) for c in pid) % (2 ** 31)
    png = _comfy_generate(domain, date, seed)
    if png:
        try:
            with open(png_path, "wb") as f:
                f.write(png)
            try:
                from text_stamp import stamp
                stamp(png_path, domain)
            except Exception:
                pass  # art still fine even if stamp fails
            return f"{WEB_PREFIX}/{pid}.png"
        except OSError:
            pass
    try:
        with open(svg_path, "wb") as f:
            f.write(_fallback_svg(domain, date))
        return f"{WEB_PREFIX}/{pid}.svg"
    except OSError:
        return None


if __name__ == "__main__":
    import sys
    d = sys.argv[1] if len(sys.argv) > 1 else "mozilla.org"
    dt = sys.argv[2] if len(sys.argv) > 2 else "2026-09-24"
    print(generate_cover(d, dt) or "FAILED")