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


HOST_COLORS = {
    "tozt": "\x1b[1;33m",
    "mail": "\x1b[32m",
    "partofme": "\x1b[35m",
}


def color_host(host: str):
    return HOST_COLORS[host] + host + "\x1b[m"


async def unlock_rbw():
    proc = await asyncio.create_subprocess_exec("rbw", "unlock")
    await proc.wait()


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 == f"[sudo] password for {os.environ['USER']}: ".encode():
            await sudo_cb
            sudo_cb = None
            buf = b""


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

    proc = await asyncio.create_subprocess_exec(
        "ssh",
        host,
        "sudo",
        "--stdin",
        *sys.argv[1:],
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    assert proc.stdout is not None
    assert proc.stderr is not None

    colored_host = color_host(host)

    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"[{colored_host}:out] {line}"),
            None,
        ),
        read_stream(
            proc.stderr,
            lambda line: print(f"[{colored_host}:\x1b[31merr\x1b[m] {line}"),
            sudo_cb(),
        ),
    )

    ret = await proc.wait()
    if ret == 0:
        print(f"[{colored_host}] Exited successfully")
    else:
        print(f"[{colored_host}] \x1b[31mExited with code {ret}\x1b[m")


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


asyncio.run(main())