"""Session startup regressions using fake services; never touches the host bus."""
import os
from pathlib import Path
import shutil
import signal
import subprocess
import tempfile
import time
import unittest

SCRIPT = Path(__file__).resolve().parents[2] / 'packaging/ferese-session-shell'


class SessionStartupTest(unittest.TestCase):
    def run_session(self, direct=True, fail_target=False, shell_status=0, terminate=False,
                     display=None, xauthority=None, extra_env=None, programs_override=None,
                     user_customization=None, restart_after_kill=False):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            helper = root / 'ferese-session-shell'
            shutil.copyfile(SCRIPT, helper)
            (root / 'rime').mkdir()
            shutil.copyfile(SCRIPT.parent / 'rime/default.custom.yaml', root / 'rime/default.custom.yaml')
            shared = root / 'rime-shared'
            shared.mkdir()
            (shared / 'luna_pinyin_simp.schema.yaml').write_text('fixture')
            units = root / 'units'
            units.mkdir()
            foreign_unit = units / 'ferese-rime-session@another-session.service'
            foreign_unit.write_text('foreign service')
            if user_customization is not None:
                user_rime = root / 'data/ferese/rime'
                user_rime.mkdir(parents=True)
                (user_rime / 'default.custom.yaml').write_text(user_customization)
            programs = {
                'dbus-update-activation-environment': (
                    'printf "import:%s\\n" "$WAYLAND_DISPLAY" >> "$TEST_ROOT/events"\n'
                    'printf "vars:%s\\n" "$*" >> "$TEST_ROOT/events"\n'),
                'dbus-send': 'echo reload-bus >> "$TEST_ROOT/events"\n',
                'systemctl': (
                    'echo "$*" >> "$TEST_ROOT/events"\n'
                    'if [[ $* == "--user start ferese-session.target" && $TEST_FAIL == 1 ]]; then exit 1; fi\n'
                    'if [[ ${2:-} == show && ${3:-} == ferese-rime-session@* ]]; then\n'
                    '  if [[ -f $TEST_ROOT/units/$3 ]]; then echo loaded; else echo not-found; fi\n'
                    'elif [[ ${2:-} == stop && ${3:-} == ferese-rime-session@* ]]; then\n'
                    '  rm -f -- "$TEST_ROOT/units/$3"\nfi\n'),
                'ferese-polkit-agent': 'exec sleep 30\n',
                'ferese-shell': 'printf "shell:%s\\n" "$WAYLAND_SOCKET" >> "$TEST_ROOT/events"\nif [[ $TEST_WAIT == 1 ]]; then exec sleep 30; fi\nexit "$TEST_STATUS"\n',
                'ferese-rime': 'exit 0\n',
                'systemd-run': (
                    'for arg; do if [[ $arg == --unit=* ]]; then unit=${arg#--unit=}; fi; done\n'
                    '(set -o noclobber; echo running > "$TEST_ROOT/units/$unit") || exit 1\n'
                    'printf "rime-arg:%s\\n" "$@" >> "$TEST_ROOT/events"\n'
                    'printf "rime-config:%s\\n" "$(cat "$XDG_DATA_HOME/ferese/rime/default.custom.yaml")" >> "$TEST_ROOT/events"\n'
                    'touch "$TEST_ROOT/rime-started"\n'),
            }
            programs.update(programs_override or {})
            for name, body in programs.items():
                path = root / name
                path.write_text('#!/bin/bash\n' + body)
                path.chmod(0o755)
            env = dict(os.environ, PATH=f'{root}:{os.environ["PATH"]}', TEST_ROOT=directory,
                       FERESE_SESSION_MODE='desktop',
                       FERESE_SESSION_IMPORT_ENV=str(int(direct)), WAYLAND_DISPLAY='private-display',
                       FERESE_PUBLIC_WAYLAND_DISPLAY='public-display', WAYLAND_SOCKET='77',
                       FERESE_SOCKET=str(root / 'control.sock'),
                       XDG_DATA_HOME=str(root / 'data'), FERESE_RIME_SHARED_DATA=str(shared),
                       TEST_FAIL=str(int(fail_target)), TEST_STATUS=str(shell_status),
                       TEST_WAIT=str(int(terminate or restart_after_kill)))
            env.pop('DISPLAY', None)
            env.pop('XAUTHORITY', None)
            if display is not None:
                env['DISPLAY'] = display
            if xauthority is not None:
                env['XAUTHORITY'] = xauthority
            env.update(extra_env or {})
            child = subprocess.Popen(['bash', str(helper)], env=env, stdout=subprocess.PIPE,
                                     stderr=subprocess.PIPE, start_new_session=True)
            if restart_after_kill:
                deadline = time.monotonic() + 3
                while not (root / 'rime-started').exists():
                    if time.monotonic() > deadline:
                        os.killpg(child.pid, signal.SIGKILL)
                        self.fail('Native input unit did not start')
                    time.sleep(.01)
                child.kill()
                child.wait(timeout=2)
                # The helper cannot run cleanup after SIGKILL. Dispose of fake
                # shell/agent processes, while its fake systemd service survives.
                os.killpg(child.pid, signal.SIGKILL)
                child.communicate(timeout=2)
                self.assertEqual(len(list(units.glob('*.service'))), 2)
                env.update(TEST_WAIT='0', FERESE_CLIENT_RESTART='1')
                child = subprocess.Popen(['bash', str(helper)], env=env,
                                         stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            if terminate:
                deadline = time.monotonic() + 3
                while not (root / 'events').exists() or 'shell:77' not in (root / 'events').read_text():
                    if time.monotonic() > deadline:
                        child.kill()
                        self.fail('Session shell did not start')
                    time.sleep(.01)
                child.send_signal(signal.SIGTERM)
            child.communicate(timeout=4)
            self.assertEqual(foreign_unit.read_text(), 'foreign service')
            return child.returncode, (root / 'events').read_text().splitlines()

    @staticmethod
    def imported_vars(events):
        for line in events:
            if line.startswith('vars:'):
                return line[len('vars:'):].split(' ')
        raise AssertionError('no activation-environment import recorded')

    def test_direct_session_imports_public_display_before_activation_and_cleans_up(self):
        status, events = self.run_session(shell_status=7)
        self.assertEqual(status, 7)
        self.assertEqual([e for e in events if e.startswith('import:')], ['import:public-display'])
        self.assertEqual([e for e in events if not e.startswith(('import:', 'vars:'))],
                         ['reload-bus', '--user daemon-reload',
                          '--user start ferese-session.target',
                          '--user try-restart xdg-desktop-portal-ferese.service', 'shell:77',
                          '--user stop ferese-session.target'])

    def test_direct_session_imports_its_control_socket(self):
        status, events = self.run_session(extra_env={'FERESE_SOCKET': '/run/user/1000/ferese/control.sock'})
        self.assertEqual(status, 0)
        self.assertIn('FERESE_SOCKET=/run/user/1000/ferese/control.sock', self.imported_vars(events))

    def test_native_rime_is_release_pinned_supervised_and_stopped_with_its_session(self):
        status, events = self.run_session(extra_env={'FERESE_ENABLE_INPUT_METHOD': '1'})
        self.assertEqual(status, 0)
        args = [event.removeprefix('rime-arg:') for event in events if event.startswith('rime-arg:')]
        self.assertIn('--service-type=notify', args)
        for setting in ('Restart=on-failure', 'WatchdogSec=10', 'StartLimitBurst=5',
                        'BindsTo=ferese-session.target'):
            self.assertIn('--property=' + setting, args)
        self.assertIn('WAYLAND_DISPLAY=public-display', args)
        self.assertIn('WAYLAND_SOCKET', args)
        self.assertIn('FERESE_SHELL_CONTROL_SOCKET', args)
        executable = next(arg for arg in args if arg.endswith('/ferese-rime'))
        self.assertNotEqual(executable, '/usr/bin/ferese-rime')
        unit = next(arg.removeprefix('--unit=') for arg in args if arg.startswith('--unit='))
        self.assertIn('--user stop ' + unit, events)
        self.assertEqual(events[-1], '--user stop ferese-session.target')
        self.assertIn('    - schema: luna_pinyin_simp', events)

    def test_native_rime_preserves_existing_user_customization(self):
        status, events = self.run_session(extra_env={'FERESE_ENABLE_INPUT_METHOD': '1'},
                                          user_customization='my-custom-dictionary: preserved\n')
        self.assertEqual(status, 0)
        self.assertIn('rime-config:my-custom-dictionary: preserved', events)

    def test_hard_killed_helper_replaces_only_its_existing_native_rime_unit(self):
        status, events = self.run_session(extra_env={'FERESE_ENABLE_INPUT_METHOD': '1'},
                                          restart_after_kill=True)
        self.assertEqual(status, 0)
        starts = [event.removeprefix('rime-arg:--unit=') for event in events
                  if event.startswith('rime-arg:--unit=')]
        self.assertEqual(len(starts), 2)
        self.assertEqual(starts[0], starts[1])
        stopped = '--user stop ' + starts[0]
        start_events = [index for index, event in enumerate(events) if event == 'rime-arg:--unit=' + starts[0]]
        self.assertIn(stopped, events[start_events[0] + 1:start_events[1]])
        self.assertNotIn('--user stop ferese-rime-session@another-session.service', events)

    def test_unresponsive_old_rime_unit_blocks_replacement_without_stopping_the_desktop(self):
        # Model a timed-out manager request. The helper must not create a second
        # input service while it cannot confirm the state of the first one.
        status, events = self.run_session(extra_env={'FERESE_ENABLE_INPUT_METHOD': '1'},
            programs_override={'systemctl': (
                'echo "$*" >> "$TEST_ROOT/events"\n'
                'if [[ ${2:-} == show ]]; then exit 124; fi\n')})
        self.assertEqual(status, 0)
        self.assertIn('shell:77', events)
        self.assertFalse(any(event.startswith('rime-arg:') for event in events))

    def test_failed_rime_start_leaves_the_desktop_running(self):
        status, events = self.run_session(extra_env={'FERESE_ENABLE_INPUT_METHOD': '1'},
                                          programs_override={'systemd-run': 'exit 7\n'})
        self.assertEqual(status, 0)
        self.assertIn('shell:77', events)
        self.assertFalse(any(event.startswith('--user stop ferese-rime-') for event in events))

    def test_nested_session_never_starts_native_rime(self):
        status, events = self.run_session(direct=False, extra_env={'FERESE_ENABLE_INPUT_METHOD': '1'})
        self.assertEqual(status, 0)
        self.assertEqual(events, ['shell:77'])

    def test_native_input_environment_is_imported_for_activated_applications(self):
        status, events = self.run_session(extra_env={'FERESE_ENABLE_INPUT_METHOD': '1',
            'GTK_IM_MODULE': '', 'QT_IM_MODULE': 'wayland', 'QT_IM_MODULES': 'wayland',
            'IBUS_ADDRESS': '', 'XMODIFIERS': '@im=ibus', 'FERESE_RIME_CONTROL': '/private/rime.sock'})
        self.assertEqual(status, 0)
        imported = self.imported_vars(events)
        for name in ('QT_IM_MODULE', 'QT_IM_MODULES', 'IBUS_ADDRESS', 'FERESE_RIME_CONTROL'):
            self.assertIn(name, imported)
        self.assertIn('GTK_IM_MODULE=', imported)
        self.assertIn('SDL_IM_MODULE=', imported)
        self.assertIn('XMODIFIERS=@im=ibus', imported)

    def test_a_managed_x11_endpoint_is_imported_for_dbus_activators(self):
        status, events = self.run_session(display=':7', xauthority='/run/user/1000/ferese-xauth-abcd')
        self.assertEqual(status, 0)
        imported = self.imported_vars(events)
        self.assertIn('DISPLAY=:7', imported)
        self.assertIn('XAUTHORITY=/run/user/1000/ferese-xauth-abcd', imported)
        self.assertIn('WAYLAND_DISPLAY=public-display', imported)

    def test_a_session_without_x11_imports_empty_values_instead_of_omitting_them(self):
        status, events = self.run_session()
        self.assertEqual(status, 0)
        imported = self.imported_vars(events)
        self.assertIn('DISPLAY=', imported)
        self.assertIn('XAUTHORITY=', imported)

    def test_activation_import_never_leaks_the_private_socket_or_cookies(self):
        status, events = self.run_session(
            display=':7', xauthority='/run/user/1000/ferese-xauth-abcd',
            extra_env={'WAYLAND_SOCKET': '77',
                       'FERESE_SHELL_CONTROL_SOCKET': '/run/user/1000/ferese/shell.sock',
                       'FERESE_PRIVATE_X11_COOKIE': 'deadbeefdeadbeefdeadbeefdeadbeef'})
        self.assertEqual(status, 0)
        imported = self.imported_vars(events)
        for leaked in ('WAYLAND_SOCKET', 'FERESE_SHELL_CONTROL_SOCKET',
                       'FERESE_PRIVATE_X11_COOKIE', 'deadbeef'):
            self.assertFalse([entry for entry in imported if leaked in entry],
                             f'{leaked} must not reach the activation environment')
        self.assertIn('WAYLAND_DISPLAY=public-display', imported)

    def test_nested_session_never_changes_host_activation(self):
        status, events = self.run_session(direct=False)
        self.assertEqual(status, 0)
        self.assertEqual(events, ['shell:77'])

    def test_embedded_policy_blocks_session_services_despite_an_inherited_import_flag(self):
        status, events = self.run_session(direct=True, extra_env={'FERESE_SESSION_MODE': 'embedded'})
        self.assertEqual(status, 0)
        self.assertEqual(events, ['shell:77'])

    def test_failed_target_does_not_prevent_shell_or_stop_unowned_target(self):
        status, events = self.run_session(fail_target=True)
        self.assertEqual(status, 0)
        self.assertIn('shell:77', events)
        self.assertNotIn('--user stop ferese-session.target', events)
        self.assertNotIn('--user try-restart xdg-desktop-portal-ferese.service', events)

    def test_termination_stops_the_session_target(self):
        status, events = self.run_session(terminate=True)
        self.assertEqual(status, 143)
        self.assertEqual(events[-1], '--user stop ferese-session.target')

    def test_agent_failure_recovers_without_restarting_shell(self):
        status, events = self.run_session(programs_override={
            'ferese-polkit-agent': (
                'echo agent >> "$TEST_ROOT/events"\n'
                'if [[ ! -f $TEST_ROOT/agent-failed ]]; then touch "$TEST_ROOT/agent-failed"; exit 7; fi\n'
                'touch "$TEST_ROOT/agent-recovered"\nexec sleep 30\n'),
            'ferese-shell': (
                'echo shell >> "$TEST_ROOT/events"\n'
                'for ((i=0; i<200; i++)); do\n'
                '  [[ -f $TEST_ROOT/agent-recovered ]] && exit 0\n'
                '  sleep .01\ndone\nexit 1\n'),
        })
        self.assertEqual(status, 0)
        self.assertEqual(events.count('agent'), 2)
        self.assertEqual(events.count('shell'), 1)
        self.assertEqual(events[-1], '--user stop ferese-session.target')

    def test_recovered_shell_does_not_repeat_autostart(self):
        programs = {
            'feresectl': 'touch "$TEST_ROOT/autostarted"\necho autostart >> "$TEST_ROOT/events"\n',
            'ferese-shell': (
                'if [[ ${FERESE_CLIENT_RESTART:-0} == 1 ]]; then sleep .1; exit 0; fi\n'
                'for ((i=0; i<100; i++)); do\n'
                '  [[ -f $TEST_ROOT/autostarted ]] && exit 0\n'
                '  sleep .01\ndone\nexit 1\n'),
        }
        for restarted in ('0', '1'):
            status, events = self.run_session(extra_env={'FERESE_CLIENT_RESTART': restarted}, programs_override=programs)
            self.assertEqual(status, 0)
            self.assertEqual(events.count('autostart'), int(restarted == '0'))


class SessionLauncherTest(unittest.TestCase):
    def launch(self, nested=False, enabled=None):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            launcher = root / 'ferese-session'
            shutil.copyfile(SCRIPT.parent / 'ferese-session', launcher)
            for name, body in {'ferese': '/usr/bin/env > "$TEST_ROOT/environment"\n',
                               'systemctl': 'exit 0\n'}.items():
                path = root / name
                path.write_text('#!/bin/bash\n' + body)
                path.chmod(0o755)
            env = dict(os.environ, PATH=f'{root}:{os.environ["PATH"]}', TEST_ROOT=directory,
                HOME=directory, XDG_RUNTIME_DIR=directory, DBUS_SESSION_BUS_ADDRESS='fixture',
                XDG_STATE_HOME=str(root / 'state'), GTK_IM_MODULE='ibus', QT_IM_MODULE='ibus',
                QT_IM_MODULES='wayland;ibus', IBUS_ADDRESS='old-ibus', XMODIFIERS='@im=ibus', SDL_IM_MODULE='ibus')
            env.pop('FERESE_ENABLE_INPUT_METHOD', None)
            if enabled is not None:
                env['FERESE_ENABLE_INPUT_METHOD'] = enabled
            subprocess.run(['bash', str(launcher), *(['--nested'] if nested else [])],
                           env=env, capture_output=True, check=True)
            return dict(line.split('=', 1) for line in (root / 'environment').read_text().splitlines())

    def test_drm_session_overrides_inherited_ibus_only_for_native_input(self):
        env = self.launch()
        self.assertEqual(env['FERESE_ENABLE_INPUT_METHOD'], '1')
        self.assertNotIn('GTK_IM_MODULE', env)
        for variable in ('QT_IM_MODULE', 'QT_IM_MODULES'):
            self.assertEqual(env[variable], 'wayland')
        self.assertEqual(env['IBUS_ADDRESS'], '')
        self.assertEqual(env['XMODIFIERS'], '@im=ibus')
        self.assertNotIn('SDL_IM_MODULE', env)
        self.assertTrue(env['FERESE_RIME_CONTROL'].endswith('/ferese/rime-control.sock'))

    def test_nested_and_explicit_opt_out_preserve_input_configuration(self):
        for options in ({'nested': True}, {'enabled': '0'}):
            with self.subTest(options=options):
                env = self.launch(**options)
                self.assertEqual(env['QT_IM_MODULE'], 'ibus')
                self.assertNotEqual(env.get('FERESE_ENABLE_INPUT_METHOD'), '1')


if __name__ == '__main__':
    unittest.main()