#!/usr/bin/python3

import libtmux
import re
import subprocess

from testutils import DONE_MARKER, DONE_MARKER_EXPRESSION, search_for_pattern, dump_screen

def get_start_timestamp(service: str) -> int:
    output = subprocess.run(
            ['systemctl', 'show', '--property', 'ExecMainStartTimestampMonotonic', service],
            capture_output=True
        ).stdout
    return int(output.split(b'=')[-1])

def get_session(server: libtmux.Server, name) -> libtmux.Session:
    return server.new_session(session_name=name, attach=False, window_command='bash -l', x=132, y=80)

def test_no_interaction(cli: str, session: libtmux.Session):
    dbus_start = get_start_timestamp('dbus.service')
    cron_start = get_start_timestamp('cron.service')
    pane = session.active_pane
    assert pane
    pane.send_keys(f"{cli} && {DONE_MARKER_EXPRESSION}")
    assert search_for_pattern(pane, DONE_MARKER)
    try:
        lines = pane.capture_pane()
        assert not isinstance(lines, str)
        anchor, restarted, deferred = None, None, None
        for i, line in enumerate(lines):
            if line.startswith('Restarting services...'):
                anchor = i
                continue
            if line.startswith('Service restarts being deferred:'):
                assert anchor is not None
                restarted = ' '.join(l.strip() for l in lines[anchor+1:i])
                continue
            if restarted is not None and line.strip():
                deferred = '\n'.join(l.strip() for l in lines[anchor+1:i])

        # Just quick check that needrestart retstarts things (using cron as our canary)
        assert restarted is not None and re.match(r'systemctl restart .*cron\.service', restarted)
        assert cron_start != get_start_timestamp('cron.service')

        # Also check that some are explictly *not* restarted, using dbus as our example
        assert deferred is not None and 'dbus.service' in deferred
        assert dbus_start == get_start_timestamp('dbus.service')
    except:
        dump_screen(pane)
        raise

def test_default_case(server: libtmux.Server):
    print()
    print("default case")
    print()
    session = get_session(server, 'needrestart-default')
    try:
        test_no_interaction('LANG=C DEBIAN_FRONTEND=dialog apt reinstall libc6', session)
    finally:
        session.kill()

def test_noninteractive(server: libtmux.Server):
    print()
    print("non interactive")
    print()
    session = get_session(server, 'needrestart-default')
    try:
        test_no_interaction('true | LANG=C DEBIAN_FRONTEND=noninteractive apt reinstall libc6 2>&1 | cat', session)
    finally:
        session.kill()

def test_debconf(server: libtmux.Server):
    print()
    print("debconf")
    print()
    session = get_session(server, 'needrestart-medium')
    try:
        dbus_start = get_start_timestamp('dbus.service')
        cron_start = get_start_timestamp('cron.service')
        pane = session.active_pane
        assert pane
        # autopkgtest environment sets DEBIAN_FRONTEND, so we need to explicitly set it back
        pane.send_keys(f"LANG=C DEBIAN_FRONTEND=dialog NEEDRESTART_UI=NeedRestart::UI::Debconf apt reinstall libc6 && {DONE_MARKER_EXPRESSION}")

        # At this point we should be in the interactive needrestart prompt
        assert search_for_pattern(pane, 'Which services should be restarted?', timeout=300)

        pane.send_keys('\t') # Just validates the default prompt

        assert search_for_pattern(pane, DONE_MARKER)
        lines = pane.capture_pane()
        try:
            assert not isinstance(lines, str)
            anchor, restarted, deferred = None, None, None
            for i, line in enumerate(lines):
                if line.startswith('Restarting services...'):
                    anchor = i
                    continue
                if line.startswith('Service restarts being deferred:'):
                    assert anchor is not None
                    restarted = ' '.join(l.strip() for l in lines[anchor+1:i])
                    continue
                if restarted is not None and line.strip():
                    deferred = '\n'.join(l.strip() for l in lines[anchor+1:i])

            # Just quick check that needrestart retstarts things (using cron as our canary)
            assert restarted is not None and re.match(r'systemctl restart .*cron\.service', restarted)
            assert cron_start != get_start_timestamp('cron.service')

            # Also check that some are explictly *not* restarted, using dbus as our example
            assert deferred is not None and 'dbus.service' in deferred
            assert dbus_start == get_start_timestamp('dbus.service')
        except:
            dump_screen(pane)
            raise
    finally:
        session.kill()

if __name__ == '__main__':
    SOCKET_NAME='needrestart-tmux'
    server = libtmux.Server(socket_name=SOCKET_NAME)
    try:
        test_default_case(server)
        test_noninteractive(server)
        test_debconf(server)
    finally:
        server.kill()

