#!/usr/bin/env python3

import os
import sys
import telnetlib

import common

class Main(common.LkmcCliFunction):
    def __init__(self):
        super().__init__(
            description='''\
Run a command on the QEMU monitor of a running QEMU instance

If the stdin is a terminal, open an interact shell. Otherwise,
run commands from stdin and quit.
''',
        )
        self.add_argument(
            'command',
            help='If given, run this command and quit',
            nargs='*',
        )

    def timed_main(self):
        def write_and_read(tn, cmd, prompt):
            tn.write(cmd.encode('utf-8'))
            return '\n'.join(tn.read_until(prompt).decode('utf-8').splitlines()[1:])[:-len(prompt)]

        with telnetlib.Telnet('localhost', self.env['qemu_monitor_port']) as tn:
            prompt = b'\n(qemu) '
            # Couldn't disable server echo, so just removing the write for now.
            # https://stackoverflow.com/questions/12421799/how-to-disable-telnet-echo-in-python-telnetlib
            # sock = tn.get_socket()
            # sock.send(telnetlib.IAC + telnetlib.WILL + telnetlib.ECHO)
            if os.isatty(sys.stdin.fileno()):
                if self.env['command'] == []:
                    print(tn.read_until(prompt).decode('utf-8'), end='')
                    tn.interact()
                else:
                    tn.read_until(prompt)
                    print(write_and_read(tn, ' '.join(self.env['command']) + '\n', prompt))
            else:
                tn.read_until(prompt)
                print(write_and_read(tn, sys.stdin.read() + '\n', prompt))

if __name__ == '__main__':
    Main().cli()