"""
Download every still-retrievable WhatsApp media file for the support number and
store it permanently on disk.

WhatsApp serves media from its own CDN for a limited window and then drops it —
measured at roughly two weeks on this account. Evolution keeps only the encrypted
URL, so anything past the window is already unrecoverable. Everything inside it
is downloaded here once, decrypted by Evolution, and written to local storage so
the support inbox can keep showing it indefinitely.

Idempotent: a file already on disk is skipped, so this can be re-run safely and
scheduled to catch anything the webhook missed.
"""
import base64, json, os, re, subprocess, sys, time, urllib.request, urllib.error

INSTANCE = os.environ.get("MEDIA_INSTANCE", "birdie")
MEDIA_DIR = os.environ.get("WHATSAPP_MEDIA_DIR", "/var/www/whatsapp-media")
SINCE = os.environ.get("MEDIA_SINCE", "2026-07-01")
DELAY = float(os.environ.get("MEDIA_DELAY", "0.4"))

EVO_BASE = os.environ["EVOLUTION_API_BASE"].rstrip("/")
EVO_KEY = os.environ["EVOLUTION_API_KEY"]
SB_URL = os.environ["SUPABASE_URL"].rstrip("/")
SB_KEY = os.environ["SUPABASE_SERVICE_ROLE_KEY"]

MEDIA_TYPES = ("imageMessage", "audioMessage", "videoMessage", "documentMessage", "stickerMessage")

EXT = {
    "image/jpeg": "jpg", "image/png": "png", "image/webp": "webp", "image/gif": "gif",
    "video/mp4": "mp4", "video/3gpp": "3gp", "video/quicktime": "mov",
    "audio/ogg": "ogg", "audio/mpeg": "mp3", "audio/mp4": "m4a", "audio/amr": "amr",
    "application/pdf": "pdf",
}


def sh(sql):
    out = subprocess.run(
        ["docker", "exec", "evolution_postgres", "psql", "-U", "evolution", "-d", "evolution",
         "-t", "-A", "-F", "\t", "-c", sql],
        capture_output=True, text=True, timeout=300,
    )
    if out.returncode != 0:
        raise RuntimeError(out.stderr[:500])
    return [l.split("\t") for l in out.stdout.strip().splitlines() if l.strip()]


def evo_media(wa_id):
    body = json.dumps({"message": {"key": {"id": wa_id}}, "convertToMp4": False}).encode()
    r = urllib.request.Request(
        f"{EVO_BASE}/chat/getBase64FromMediaMessage/{INSTANCE}",
        data=body, method="POST",
        headers={"apikey": EVO_KEY, "Content-Type": "application/json"},
    )
    try:
        return json.loads(urllib.request.urlopen(r, timeout=180).read().decode())
    except urllib.error.HTTPError as e:
        return {"_error": f"HTTP {e.code}"}
    except Exception as e:
        return {"_error": str(e)[:120]}


def sb(method, path, body=None, prefer=None):
    h = {"apikey": SB_KEY, "Authorization": "Bearer " + SB_KEY, "Content-Type": "application/json"}
    if prefer:
        h["Prefer"] = prefer
    r = urllib.request.Request(SB_URL + "/rest/v1/" + path, method=method,
                               data=json.dumps(body).encode() if body else None, headers=h)
    raw = urllib.request.urlopen(r, timeout=120).read().decode()
    return json.loads(raw) if raw.strip() else None


rows = sh(f"""
select m.key->>'id', m."messageType"
from "Message" m join "Instance" i on i.id = m."instanceId"
where i.name = '{INSTANCE}'
  and m."messageType" in {MEDIA_TYPES}
  and to_timestamp(m."messageTimestamp") >= '{SINCE}'
order by m."messageTimestamp" desc
""")
print(f"candidates: {len(rows)}", flush=True)

os.makedirs(os.path.join(MEDIA_DIR, INSTANCE), exist_ok=True)

saved = skipped = expired = failed = 0
meta_updates = []

for n, (wa_id, mtype) in enumerate(rows, 1):
    wa_id = wa_id.strip()
    if not wa_id or not re.fullmatch(r"[A-Za-z0-9_-]{5,80}", wa_id):
        failed += 1
        continue

    existing = [f for f in os.listdir(os.path.join(MEDIA_DIR, INSTANCE)) if f.startswith(wa_id + ".")] \
        if os.path.isdir(os.path.join(MEDIA_DIR, INSTANCE)) else []
    if existing:
        skipped += 1
        continue

    res = evo_media(wa_id)
    if res.get("_error") or not res.get("base64"):
        # A 400 here means WhatsApp has already dropped the file — not a bug.
        expired += 1
    else:
        mime = res.get("mimetype") or "application/octet-stream"
        mime_base = mime.split(";")[0].strip()
        ext = EXT.get(mime_base) or (res.get("fileName", "").rsplit(".", 1) + ["bin"])[1][:5]
        path = os.path.join(MEDIA_DIR, INSTANCE, f"{wa_id}.{ext}")
        try:
            data = base64.b64decode(res["base64"])
            with open(path, "wb") as fh:
                fh.write(data)
            saved += 1
            meta_updates.append({
                "wa_id": wa_id,
                "media": {
                    "kind": mtype.replace("Message", ""),
                    "mimetype": mime_base,
                    "filename": res.get("fileName") or f"{wa_id}.{ext}",
                    "size": len(data),
                    "ext": ext,
                    "stored": True,
                },
            })
        except Exception as e:
            print(f"  write failed {wa_id}: {e}", flush=True)
            failed += 1

    if n % 25 == 0:
        print(f"  {n}/{len(rows)}  saved={saved} skipped={skipped} expired={expired} failed={failed}",
              flush=True)
    time.sleep(DELAY)

print(f"DOWNLOAD DONE saved={saved} skipped={skipped} expired={expired} failed={failed}", flush=True)

# ---- record what we stored on the message rows so the UI knows what to render ----
print(f"tagging {len(meta_updates)} messages ...", flush=True)
tagged = 0
for u in meta_updates:
    try:
        got = sb("GET", f"whatsapp_messages?wa_message_id=eq.{u['wa_id']}&select=id,raw_payload") or []
        for row in got:
            rp = row.get("raw_payload") or {}
            rp["media"] = u["media"]
            sb("PATCH", f"whatsapp_messages?id=eq.{row['id']}", {"raw_payload": rp}, prefer="return=minimal")
            tagged += 1
    except Exception as e:
        print(f"  tag failed {u['wa_id']}: {str(e)[:100]}", flush=True)
print(f"TAGGED {tagged}", flush=True)
print("ALL DONE", flush=True)