#!/usr/bin/env python3
"""Cancel only fake find workers; no real cache deletions or Perforce calls."""
import os
from pathlib import Path
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import unittest

HELPER = Path(__file__).resolve().parents[1] / 'Server/Unix/p4/common/bin/sdp_cache_cleanup.sh'
BASH = os.environ.get('SDP_TEST_BASH', shutil.which('bash'))


class CacheCancellation(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory(prefix='sdp-cancel-test-')
        self.addCleanup(self.tmp.cleanup)
        self.root = Path(self.tmp.name)
        (self.root / 'bin').mkdir()
        fake = self.root / 'bin/find'
        fake.write_text('#!' + sys.executable + '\n' + '''import os, signal, sys, time
from pathlib import Path
root = Path(os.environ['TEST_ROOT'])
if os.environ.get('IGNORE_TERM'):
    signal.signal(signal.SIGTERM, signal.SIG_IGN)
sys.stdout.write('...'); sys.stdout.flush()
(root / ('started.' + str(os.getpid()))).write_text(str(os.getpid()))
if not os.environ.get('FINISH_NOW'):
    while not (root/'proceed').exists(): time.sleep(.01)
(root / ('continued.' + str(os.getpid()))).write_text('unexpected after cancellation')
sys.exit(int(os.environ.get('FIND_EXIT', '0')))
''')
        fake.chmod(0o755)
        self.script = self.root / 'run.sh'
        self.script.write_text('''#!/bin/bash
source "$HELPER"
trap 'echo original-term >> "$TEST_ROOT/traps"; exit 143' TERM
trap 'echo original-int >> "$TEST_ROOT/traps"; exit 130' INT
trap 'echo original-hup >> "$TEST_ROOT/traps"; exit 129' HUP
trap 'echo "exit:$? removed:$SDP_CACHE_REMOVED" >> "$TEST_ROOT/traps"' EXIT
sdp_add_field() { echo "$1=$2" >> "$TEST_ROOT/fields"; }
sdp_set_error() { echo "$1" >> "$TEST_ROOT/error"; }
sdp_cache_cleanup_dirs 7 no "$TEST_ROOT/depot1" "$TEST_ROOT/depot2" "$TEST_ROOT/depot3"
result=$?
echo "result:$result removed:$SDP_CACHE_REMOVED" >> "$TEST_ROOT/traps"
if [[ ${CHECK_RESTORED:-0} == 1 ]]; then kill -TERM $$; fi
exit "$result"
''')
        self.env = dict(os.environ, TEST_ROOT=str(self.root), HELPER=str(HELPER), P4TMP=str(self.root),
                        PATH=str(self.root/'bin')+':'+os.environ['PATH'], SDP_CACHE_CLEAN_WORKERS='2')

    def test_signals_stop_actual_find_children_preserve_counts_and_caller_traps(self):
        for sig, code, name, ignore in [(signal.SIGTERM,143,'term',False), (signal.SIGINT,130,'int',True), (signal.SIGHUP,129,'hup',False)]:
            with self.subTest(signal=sig):
                for path in self.root.glob('started.*'): path.unlink()
                for name_file in ['traps','fields','error','proceed']:
                    (self.root/name_file).unlink(missing_ok=True)
                env = dict(self.env)
                if ignore: env['IGNORE_TERM'] = '1'
                process = subprocess.Popen([BASH,str(self.script)],env=env,stdout=subprocess.PIPE,stderr=subprocess.PIPE,text=True)
                try:
                    deadline = time.monotonic()+5
                    while len(list(self.root.glob('started.*'))) < 2 and time.monotonic()<deadline:
                        time.sleep(.01)
                    pids = [int(path.read_text()) for path in self.root.glob('started.*')]
                    self.assertEqual(len(pids),2)
                    process.send_signal(sig)
                    out,err = process.communicate(timeout=5)
                    self.assertEqual(process.returncode,code,out+err)
                    self.assertIn('original-'+name,(self.root/'traps').read_text())
                    self.assertIn('removed:6',(self.root/'traps').read_text())
                    self.assertIn('files_removed=6',(self.root/'fields').read_text())
                    self.assertIn('interrupted',(self.root/'error').read_text())
                    self.assertEqual(list(self.root.glob('cache-clean.*')),[])
                    for pid in pids:
                        with self.assertRaises(ProcessLookupError): os.kill(pid,0)
                    (self.root/'proceed').touch()
                    self.assertEqual(list(self.root.glob('continued.*')),[])
                finally:
                    if process.poll() is None: process.kill(); process.communicate()

    def test_success_restores_original_signal_and_exit_handlers(self):
        result = subprocess.run([BASH,str(self.script)],env=dict(self.env,FINISH_NOW='1',CHECK_RESTORED='1'),capture_output=True,text=True,timeout=5)
        self.assertEqual(result.returncode,143,result.stderr)
        text=(self.root/'traps').read_text()
        self.assertIn('result:0 removed:9',text)
        self.assertIn('original-term',text)
        self.assertIn('exit:143',text)
        self.assertEqual(list(self.root.glob('cache-clean.*')),[])

    def test_worker_failure_counts_partial_work_and_reaps_remaining_workers(self):
        result = subprocess.run([BASH,str(self.script)],env=dict(self.env,FINISH_NOW='1',FIND_EXIT='7'),capture_output=True,text=True,timeout=5)
        self.assertEqual(result.returncode,1,result.stderr)
        self.assertIn('result:1 removed:9',(self.root/'traps').read_text())
        self.assertEqual(len(list(self.root.glob('continued.*'))),3)
        self.assertEqual(list(self.root.glob('cache-clean.*')),[])


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