import unittest
import contextlib
import multiprocessing
from concurrent.futures import ProcessPoolExecutor

from vapoursynth import core
import vapoursynth as vs


class StubPolicy(vs.EnvironmentPolicy):

    def __init__(self) -> None:
        self._current = None
        self._api = None

    def on_policy_registered(self, special_api):
        self._api = special_api
        self._current = None

    def on_policy_cleared(self):
        assert self._api is not None
        self._current = None

    def get_current_environment(self):
        return self._current

    def set_environment(self, environment):
        self._current = environment


@contextlib.contextmanager
def _with_policy():
    pol = StubPolicy()
    vs.register_policy(pol)
    try:
        yield pol
    finally:
        pol._api.unregister_policy()


def _wrap_with(func):
    func(EnvironmentTest())


def subprocess_runner(func):
    def _wrapper(self):
        with ProcessPoolExecutor(max_workers=1, mp_context=multiprocessing.get_context("spawn")) as executor:
            executor.submit(_wrap_with, func).result()
    return _wrapper

def t_environment_use_unsets_environment_on_exit(self):
    with _with_policy() as pol:
        env = pol._api.create_environment()
        wrapped = pol._api.wrap_environment(env)

        with self.assertRaises(RuntimeError):
            vs.get_current_environment()

        with wrapped.use():
            print(vs.get_current_environment())

        with self.assertRaises(RuntimeError):
            print(vs.get_current_environment())


def t_environment_use_restores_environment_on_exit(self):
    with _with_policy() as pol:
        env1 = pol._api.create_environment()
        wrapped1 = pol._api.wrap_environment(env1)

        env2 = pol._api.create_environment()
        wrapped2 = pol._api.wrap_environment(env2)

        with wrapped1.use():
            ce1 = vs.get_current_environment()

            with wrapped2.use():
                self.assertNotEqual(ce1, vs.get_current_environment())

            self.assertEqual(ce1, vs.get_current_environment())


class EnvironmentTest(unittest.TestCase):

    test_environment_use_unsets_environment_on_exit = subprocess_runner(t_environment_use_unsets_environment_on_exit)
    test_environment_use_restores_environment_on_exit = subprocess_runner(t_environment_use_restores_environment_on_exit)


@@ -829,88 +829,262 @@
 )%0A%0A%0A
-def _wrap_with(func):%0A    func(EnvironmentTest())%0A%0A%0Adef subprocess_runner(func):
+test_functions = %7B%7D%0Acounter = 0%0A%0Adef _wrap_with(ident):%0A    func = test_functions%5Bident%5D%0A    return func(EnvironmentTest())%0A%0A%0Adef subprocess_runner(func):%0A    global counter%0A    my_counter = counter%0A    counter += 1%0A    test_functions%5Bmy_counter%5D = func%0A
 %0A   
@@ -1254,20 +1254,26 @@
 p_with, 
-func
+my_counter
 ).result
@@ -1271,24 +1271,25 @@
 r).result()%0A
+%0A
     return _
@@ -1297,20 +1297,93 @@
 rapper%0A%0A
+%0Aclass EnvironmentTest(unittest.TestCase):%0A    @subprocess_runner%0A    
 def 
+tes
 t_enviro
@@ -1418,32 +1418,36 @@
 _on_exit(self):%0A
+    
     with _with_p
@@ -1470,16 +1470,20 @@
         
+ 
+   
 env = po
@@ -1514,24 +1514,28 @@
 t()%0A        
+  
+  
 wrapped = po
@@ -1563,16 +1563,20 @@
 nt(env)%0A
+    
 %0A       
@@ -1568,32 +1568,36 @@
 v)%0A    %0A        
+    
 with self.assert
@@ -1630,16 +1630,20 @@
         
+  
+  
 vs.get_c
@@ -1655,32 +1655,36 @@
 t_environment()%0A
+    
 %0A        with wr
@@ -1668,32 +1668,36 @@
 ()%0A    %0A        
+    
 with wrapped.use
@@ -1704,37 +1704,52 @@
 ():%0A            
+ 
-print
+   self.assertEqual
 (vs.get_current_
@@ -1753,35 +1753,52 @@
 nt_environment()
-)%0A%0A
+, wrapped)%0A    %0A    
         with sel
@@ -1843,14 +1843,12 @@
     
-print(
+    
 vs.g
@@ -1871,24 +1871,61 @@
 onment()
-)%0A%0A%0A
+%0A    %0A    %0A    @subprocess_runner%0A    
 def 
+tes
 t_enviro
@@ -1966,24 +1966,28 @@
 exit(self):%0A
+    
     with _wi
@@ -2014,16 +2014,20 @@
         
+    
 env1 = p
@@ -2047,32 +2047,36 @@
 e_environment()%0A
+    
         wrapped1
@@ -2110,16 +2110,20 @@
 t(env1)%0A
+    
 %0A       
@@ -2123,16 +2123,20 @@
         
+    
 env2 = p
@@ -2156,32 +2156,36 @@
 e_environment()%0A
+    
         wrapped2
@@ -2219,17 +2219,25 @@
 t(env2)%0A
-%0A
+    %0A    
         
@@ -2249,32 +2249,36 @@
 wrapped1.use():%0A
+    
             ce1 
@@ -2304,25 +2304,33 @@
 vironment()%0A
-%0A
+    %0A    
             
@@ -2342,32 +2342,36 @@
 wrapped2.use():%0A
+    
                 
@@ -2425,17 +2425,25 @@
 ment())%0A
-%0A
+    %0A    
         
@@ -2502,290 +2502,4 @@
 ())%0A
-%0A%0Aclass EnvironmentTest(unittest.TestCase):%0A%0A    test_environment_use_unsets_environment_on_exit = subprocess_runner(t_environment_use_unsets_environment_on_exit)%0A    test_environment_use_restores_environment_on_exit = subprocess_runner(t_environment_use_restores_environment_on_exit)%0A%0A

Refactor environment tests.