import unittest
import os.path
import zipfile
import time
import io
import imp
import marshal
import sys


TEST_INIT = """\
from .ztst import get_co_filename
"""

TEST_SRC = """\
import sys
def get_co_filename():
    return sys._getframe().f_code.co_filename
"""

TEMP_ZIP_PY = os.path.join(os.path.abspath(os.path.dirname(__file__)), "cofn_py.zip")
TEMP_ZIP_PYC = os.path.join(os.path.abspath(os.path.dirname(__file__)), "cofn_pyc.zip")
TEMP_ZIP_BOTH = os.path.join(os.path.abspath(os.path.dirname(__file__)), "cofn_both.zip")

NOW = time.time()
NOW_Z = time.localtime(NOW)[:6]


class ZippedModuleCoFilenameTestCase(unittest.TestCase):
    def doTest(self, zpath):
        sys.path.insert(0, zpath)
        try:
            if 'mycofn' in sys.modules:
                del sys.modules['mycofn']
            if 'mycofn.ztst' in sys.modules:
                del sys.modules['mycofn.ztst']
            import mycofn
            cofn = mycofn.get_co_filename()
            self.assertTrue(cofn.startswith(zpath))
        finally:
            del sys.path[0]
    
    def testPy(self):
        self.doTest(TEMP_ZIP_PY)
 
    def testPyc(self):
        self.doTest(TEMP_ZIP_PYC)
        self.assertTrue(True)

    def testBoth(self):
        self.doTest(TEMP_ZIP_BOTH)
        self.assertTrue(True)


def compile_to_pyc(source, codename):
    def wr_long(f, x):
        # write a 32-bit int to a file in little-endian order.
        f.write(bytes([x         & 0xff,
                       (x >> 8)  & 0xff,
                       (x >> 16) & 0xff,
                       (x >> 24) & 0xff]))
    with io.BytesIO() as out:
        ast = compile(source, codename, 'exec', optimize=1)
        timestamp = int(NOW)
        size = len(source)
        out.write(b'\0\0\0\0')
        wr_long(out, timestamp)
        wr_long(out, size)
        marshal.dump(ast, out)
        out.flush()
        out.seek(0, 0)
        out.write(imp.get_magic())
        return out.getvalue()


def write_zip_py(z, zname, data):
    zinfo = zipfile.ZipInfo(zname, NOW_Z)
    z.writestr(zinfo, data)


def write_zip_pyc(z, zname, data):
    co_name = zname.replace('.pyc', '.py')
    pyc_data = compile_to_pyc(data, co_name)
    zinfo = zipfile.ZipInfo(zname, NOW_Z)
    z.writestr(zinfo, pyc_data)


def make_zip(zpath, keep_src, keep_code):
    if os.path.exists(zpath):
        os.remove(zpath)
    with zipfile.ZipFile(zpath, "w", zipfile.ZIP_DEFLATED) as z:
        if keep_src:
            write_zip_py(z, 'mycofn/__init__.py', TEST_INIT)
            write_zip_py(z, 'mycofn/ztst.py', TEST_SRC)
        if keep_code:
            write_zip_pyc(z, 'mycofn/__init__.pyc', TEST_INIT)
            write_zip_pyc(z, 'mycofn/ztst.pyc', TEST_SRC)


def test_main():
    make_zip(TEMP_ZIP_PY, keep_src=True, keep_code=False)
    make_zip(TEMP_ZIP_PYC, keep_src=False, keep_code=True)
    make_zip(TEMP_ZIP_BOTH, keep_src=True, keep_code=True)

    unittest.main()


if __name__ == "__main__":
    test_main()