import os
import imp
import sys
import tempfile
import unittest
import subprocess
from tempfile import NamedTemporaryFile

from pargec.structure import (Structure, Field, ArrayField, uint8, FMT_BE_MSB)
from pargec.c_generator import (gen_c_struct_decl, gen_c_serialize_decl,
                                gen_c_deserialize_decl, gen_c_serialize_def,
                                gen_c_deserialize_def,
                                gen_c_constants, generate,
                                build_bit_masks, Byte)

FOO_STRUCT = """typedef struct foo_prot {
    uint8_t field1;
    uint8_t field2;
    uint8_t field3;
    uint8_t buff[8];
    uint8_t field4;
} foo_prot_t;
"""

FOO_DEFINES = """\
enum {
    FOO_PROT_SERIALIZED_N_BYTES = 11,
};
"""

FOO_SERIALIZE_DECL = "void foo_prot_serialize(foo_prot_t* in_struct, uint8_t *out_buff);"
FOO_DESERIALIZE_DECL = "void foo_prot_deserialize(foo_prot_t* out_struct, uint8_t *in_buff);"

FOO_SERIALIZE_DEF = """\
void foo_prot_serialize(foo_prot_t* in_struct, uint8_t *out_buff) {
    out_buff[0] = 0;
    out_buff[0] |= (in_struct->field1 & 0b111111) << 2;
    out_buff[0] |= (in_struct->field2 & 0b1100) >> 2;
    out_buff[1] = 0;
    out_buff[1] |= (in_struct->field2 & 0b11) << 6;
    out_buff[1] |= (in_struct->field3 & 0b111111) << 0;
    out_buff[10] = 0;
    out_buff[10] |= (in_struct->field4 & 0b11111111) << 0;
    for(int i = 0; i < 8; ++i){out_buff[2+i] = in_struct->buff[i];}
}"""

FOO_DESERIALIZE_DEF = """\
void foo_prot_deserialize(foo_prot_t* out_struct, uint8_t *in_buff) {
    out_struct->field1 = 0;
    out_struct->field1 |= (in_buff[0] >> 2) & 0b111111;
    out_struct->field2 = 0;
    out_struct->field2 |= (in_buff[0] & 0b11) << 2;
    out_struct->field2 |= (in_buff[1] >> 6) & 0b11;
    out_struct->field3 = 0;
    out_struct->field3 |= (in_buff[1] >> 0) & 0b111111;
    out_struct->field4 = 0;
    out_struct->field4 |= (in_buff[10] >> 0) & 0b11111111;
    for(int i = 0; i < 8; ++i){out_struct->buff[i] = in_buff[2+i];}
}"""

MULTIBYTE_DESERIALIZE_DEF = """\
void prot_multi_deserialize(prot_multi_t* out_struct, uint8_t *in_buff) {
    out_struct->field1 = 0;
    out_struct->field1 |= (in_buff[0] >> 2) & 0b111111;
    out_struct->field2 = 0;
    out_struct->field2 |= (in_buff[0] & 0b11) << 10;
    out_struct->field2 |= (in_buff[1] & 0b11111111) << 2;
    out_struct->field2 |= (in_buff[2] >> 6) & 0b11;
    out_struct->field3 = 0;
    out_struct->field3 |= (in_buff[2] & 0b111111) << 10;
    out_struct->field3 |= (in_buff[3] & 0b11111111) << 2;
    out_struct->field3 |= (in_buff[4] >> 6) & 0b11;
}"""

MULTIBYTE_SERIALIZE_DEF = """\
void prot_multi_serialize(prot_multi_t* in_struct, uint8_t *out_buff) {
    out_buff[0] = 0;
    out_buff[0] |= (in_struct->field1 & 0b111111) << 2;
    out_buff[0] |= (in_struct->field2 & 0b110000000000) >> 10;
    out_buff[1] = 0;
    out_buff[1] |= (in_struct->field2 & 0b1111111100) >> 2;
    out_buff[2] = 0;
    out_buff[2] |= (in_struct->field2 & 0b11) << 6;
    out_buff[2] |= (in_struct->field3 & 0b1111110000000000) >> 10;
    out_buff[3] = 0;
    out_buff[3] |= (in_struct->field3 & 0b1111111100) >> 2;
    out_buff[4] = 0;
    out_buff[4] |= (in_struct->field3 & 0b11) << 6;
}"""

FOO_ARRAY_STRUCT = """typedef struct foo_prot {
    uint8_t field1;
    uint8_t buff[64];
} foo_prot_t;
"""


class TestCGenerator(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        foo_prot = Structure("foo_prot", [
            Field("field1", uint8, 6, FMT_BE_MSB),
            Field("field2", uint8, 4, FMT_BE_MSB),
            Field("field3", uint8, 6, FMT_BE_MSB),
            ArrayField("buff", uint8, 64, FMT_BE_MSB),
            Field("field4", uint8, 8, FMT_BE_MSB),
            ])
        cls.foo_prot = foo_prot

    def setUp(self):
        self.maxDiff = None

    def test_gen_c_struct_decl(self):
        c_struct = gen_c_struct_decl(self.foo_prot)
        self.assertEqual(c_struct, FOO_STRUCT)

    def test_gen_c_serialize_decl(self):
        self.assertEqual(gen_c_serialize_decl(self.foo_prot), FOO_SERIALIZE_DECL)

    def test_gen_c_deserialize_decl(self):
        self.assertEqual(gen_c_deserialize_decl(self.foo_prot), FOO_DESERIALIZE_DECL)

    def test_gen_c_serialize_def(self):
        self.assertEqual(gen_c_serialize_def(self.foo_prot), FOO_SERIALIZE_DEF)

    def test_gen_c_deserialize_def(self):
        self.maxDiff = None
        self.assertEqual(gen_c_deserialize_def(self.foo_prot), FOO_DESERIALIZE_DEF)

    def test_gen_c_defines(self):
        self.assertEqual(gen_c_constants([self.foo_prot]), FOO_DEFINES)

    def test_bit_masks_inside_byte(self):
        prot = Structure("foo_prot", [
            Field("field1", uint8, 6, FMT_BE_MSB),
            Field("field2", uint8, 4, FMT_BE_MSB),
            Field("field3", uint8, 6, FMT_BE_MSB),
            ])

        res = build_bit_masks(prot)
        masks = [("field1", [(0, 8, 2, 6, 0)]),
                 ("field2", [(0, 2, 0, 4, 2), (1, 8, 6, 2, 0)]),
                 ("field3", [(1, 6, 0, 6, 0)]),
        ]
        self.assertEqual(res, masks)

    def test_bit_masks_multibyte(self):
        prot = Structure("prot_multi", [
            Field("field1", uint8, 6, FMT_BE_MSB),
            Field("field2", uint8, 12, FMT_BE_MSB),
            Field("field3", uint8, 16, FMT_BE_MSB),
            ])

        #byte, byte_first_bit, byte_last_bit, first_bit_to_serialize, last_bit_to_serialize
        res = build_bit_masks(prot)
        masks = [("field1", [(0, 8, 2, 6, 0)]),
                 ("field2", [(0, 2, 0, 12, 10), (1, 8, 0, 10, 2), (2, 8, 6, 2, 0)]),
                 ("field3", [(2, 6, 0, 16, 10), (3, 8, 0, 10, 2), (4, 8, 6, 2, 0)])
        ]
        self.assertEqual(res, masks)
        self.maxDiff = None
        self.assertEqual(gen_c_serialize_def(prot), MULTIBYTE_SERIALIZE_DEF)
        self.assertEqual(gen_c_deserialize_def(prot), MULTIBYTE_DESERIALIZE_DEF)

    def test_byte(self):
        byte = Byte()
        self.assertEqual(byte.process_bits(2), (8, 6, 2))
        self.assertEqual(byte.process_bits(2), (6, 4, 2))
        self.assertEqual(byte.process_bits(2), (4, 2, 2))
        self.assertTrue(byte.bits_available)
        self.assertEqual(byte.process_bits(2), (2, 0, 2))
        self.assertFalse(byte.bits_available)

        byte = Byte()
        self.assertEqual(byte.process_bits(2), (8, 6, 2))
        self.assertEqual(byte.process_bits(4), (6, 2, 4))
        self.assertEqual(byte.process_bits(4), (2, 0, 2))

        byte = Byte()
        self.assertTrue(byte.bits_available)
        self.assertEqual(byte.process_bits(10), (8, 0, 8))
        self.assertFalse(byte.bits_available)

    def test_build(self):
        out_file = NamedTemporaryFile(suffix=".c", delete=False)
        out_file.write("#include <stdint.h>\n".encode())
        out_file.write(gen_c_struct_decl(self.foo_prot).encode())
        out_file.write(gen_c_serialize_decl(self.foo_prot).encode())
        out_file.write(gen_c_serialize_def(self.foo_prot).encode())
        out_file.write("void main(void){}".encode())
        out_file.close()

        tmp_out = "/tmp/a.out"

        subprocess.check_call(["gcc", os.path.abspath(out_file.name), "-o", tmp_out])
        os.remove(tmp_out)
        os.remove(out_file.name)

    def test_online_cffi(self):
        from cffi import FFI

        ffi = FFI()
        ffi.cdef(gen_c_struct_decl(self.foo_prot))
        ffi.cdef(gen_c_serialize_decl(self.foo_prot))

        ffi.set_source("_test_online_cffi",
                       gen_c_struct_decl(self.foo_prot) + "\n" +
                       gen_c_serialize_def(self.foo_prot))

        with tempfile.TemporaryDirectory() as tmpdir:
            sys.path.append(tmpdir)
            ffi.compile(tmpdir=tmpdir)
            from _test_online_cffi import lib, ffi as nffi

        buffer_out = nffi.new("uint8_t[]", 100)
        foo_prot = nffi.new("foo_prot_t *", [0b100001, 0b1111, 0b11, tuple(range(8)), 123])
        lib.foo_prot_serialize(foo_prot, buffer_out)
        self.assertEqual(buffer_out[0], (0b100001 << 2) | 0b11)
        self.assertEqual(buffer_out[1], (0b11 << 6) | 0b11)
        for i in range(8):
            self.assertEqual(buffer_out[2 + i], i)
        self.assertEqual(buffer_out[8 + 2], 123)

    def test_generate_and_integration(self):
        TESTS_PATH = os.path.dirname(os.path.realpath(__file__))
        protocol_file = os.path.join(TESTS_PATH, "struct_defs.py")

        cwd = os.getcwd()
        with tempfile.TemporaryDirectory() as tmpdir:
            os.chdir(tmpdir)
            sys.path.append(tmpdir)
            generate(protocol_file, os.path.join(tmpdir, "foo.h"),
                     os.path.join(tmpdir, "foo.c"),
                     os.path.join(tmpdir, "foo.py"))

            foo = imp.load_source("foo", os.path.join(tmpdir, "foo.py"))

        os.chdir(cwd)
        in_data = {"field1": 0b100001, "field2": 0b1111, "field3": 0b11,
                   "buff": list(range(8)), "field4": 123}
        out_buff = foo.foo_prot_serialize(in_data)
        self.assertEqual(out_buff, bytes([(0b100001 << 2) | 0b11, (0b11 << 6) | 0b11] + list(range(8)) + [123]))
        values = foo.foo_prot_deserialize(out_buff)
        self.assertTrue(all([getattr(values, key) == in_data[key] for key in in_data.keys() if key != "buff"]))
        self.assertEqual(list(values.buff), in_data["buff"])

    def test_multibyte(self):
        prot = Structure("foo_prot", [
            Field("field1", uint8, 6, FMT_BE_MSB),
            ArrayField("buff", uint8, 8 * 64, FMT_BE_MSB),
        ])
        c_struct = gen_c_struct_decl(prot)

        self.assertEqual(c_struct, FOO_ARRAY_STRUCT)



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