# Functions to help with compatibility between BCH and NEX

from test_framework.script import OP_TRUE, CScript
from test_framework.environment import on_bch, on_nex
from test_framework.nex.script import p2st_locking_script, p2st_unlocking_script


def lock(script: CScript):
    if on_bch():
        return script
    if on_nex():
        return p2st_locking_script(script)
    raise NotImplementedError()


def unlock(script: CScript):
    if on_bch():
        return CScript([OP_TRUE])
    if on_nex():
        return p2st_unlocking_script(script)
    raise NotImplementedError()
