#!/usr/bin/env python3
import asyncio
import os
from typing import Any, Callable, Coroutine


async def get_password(host: str):
    proc = await asyncio.create_subprocess_exec(
        "rbw",
        "get",
        host,
        os.environ["USER"],
        stdout=asyncio.subprocess.PIPE,
    )

    stdout, _ = await proc.communicate()
    return stdout.rstrip()


async def read_stream(
    stream: asyncio.StreamReader,
    print_cb: Callable[[str], None],
    sudo_cb: Coroutine[Any, Any, None] | None,
):
    buf = b""
    while True:
        read = await stream.read(1024)
        if len(read) == 0:
            if len(buf) > 0:
                print_cb(buf.decode())
            break

        buf += read
        lines = buf.split(b"\n")
        buf = lines.pop()
        for line in lines:
            print_cb(line.decode())

        if sudo_cb and buf == b"[sudo] password for doy: ":
            await sudo_cb
            sudo_cb = None


async def puppet_tozt(host: str):
    password = await get_password(host)

    proc = await asyncio.create_subprocess_exec(
        "ssh",
        host,
        "sudo",
        "--stdin",
        "puppet-tozt",
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    assert proc.stdout is not None
    assert proc.stderr is not None

    async def sudo_cb():
        assert proc.stdin is not None
        proc.stdin.write(password + b"\n")
        await proc.stdin.drain()

    await asyncio.gather(
        read_stream(
            proc.stdout,
            lambda line: print(f"[{host}:out] {line}"),
            None,
        ),
        read_stream(
            proc.stderr,
            lambda line: print(f"[{host}:err] {line}"),
            sudo_cb(),
        ),
    )

    ret = await proc.wait()
    if ret == 0:
        print(f"[{host}] Exited successfully")
    else:
        print(f"[{host}] Exited with code {ret}")


async def main():
    await asyncio.gather(
        puppet_tozt("tozt"),
        puppet_tozt("mail"),
        puppet_tozt("partofme"),
    )


asyncio.run(main())