#!/usr/bin/env python3
"""Seed the local dev homeserver (scripts/dev-homeserver.sh) with two users
and a busy unencrypted room: alice + bob, "Busy Room", N messages, one image
every 10th message. Idempotent for the users; every run creates a new room.

    python3 scripts/dev-seed.py [N=400]
"""
import json, struct, sys, time, urllib.error, urllib.parse, urllib.request, zlib

HS = "http://127.0.0.1:8008"
PASSWORD = "password123"


def req(method, path, data=None, token=None, raw=None, ctype="application/json"):
    headers = {"Content-Type": ctype}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    body = raw if raw is not None else (json.dumps(data).encode() if data is not None else None)
    r = urllib.request.Request(HS + path, data=body, headers=headers, method=method)
    return json.load(urllib.request.urlopen(r))


def register_or_login(user):
    try:
        return req("POST", "/_matrix/client/v3/register",
                   {"username": user, "password": PASSWORD, "auth": {"type": "m.login.dummy"}})
    except urllib.error.HTTPError:
        return req("POST", "/_matrix/client/v3/login",
                   {"type": "m.login.password", "identifier": {"type": "m.id.user", "user": user}, "password": PASSWORD})


def png(w, h, rgb):
    raw = b"".join(b"\x00" + bytes(rgb) * w for _ in range(h))
    def chunk(t, d):
        return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d) & 0xFFFFFFFF)
    return (b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
            + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b""))


def main():
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 400
    alice, bob = register_or_login("alice"), register_or_login("bob")
    ta, tb = alice["access_token"], bob["access_token"]
    room = req("POST", "/_matrix/client/v3/createRoom",
               {"name": "Busy Room", "preset": "public_chat", "visibility": "public"}, ta)["room_id"]
    req("POST", f"/_matrix/client/v3/join/{urllib.parse.quote(room)}", {}, tb)
    txn = int(time.time() * 1000)
    for i in range(n):
        tok = ta if i % 2 == 0 else tb
        if i % 10 == 0:
            data = png(64, 48, ((i * 37) % 256, (i * 91) % 256, (i * 17) % 256))
            up = req("POST", f"/_matrix/media/v3/upload?filename=img{i}.png", token=tok, raw=data, ctype="image/png")
            content = {"msgtype": "m.image", "body": f"img{i}.png", "url": up["content_uri"],
                       "info": {"mimetype": "image/png", "w": 64, "h": 48, "size": len(data)}}
        else:
            content = {"msgtype": "m.text", "body": f"message #{i}"}
        txn += 1
        req("PUT", f"/_matrix/client/v3/rooms/{urllib.parse.quote(room)}/send/m.room.message/{txn}", content, tok)
    print(json.dumps({"room": room, "alice": alice["user_id"], "bob": bob["user_id"], "count": n}))


if __name__ == "__main__":
    main()