import asyncio
import time
import unittest
from typing import Optional

from quarkchain.cluster.miner import DoubleSHA256, Miner, MiningWork, validate_seal
from quarkchain.config import ConsensusType
from quarkchain.core import RootBlock, RootBlockHeader
from quarkchain.p2p import ecies
from quarkchain.utils import sha3_256


class TestMiner(unittest.TestCase):
    def setUp(self):
        super().setUp()

        def miner_gen(consensus, create_func, add_func, **kwargs):
            m = Miner(
                consensus, create_func, add_func, self.get_mining_params, **kwargs
            )
            m.enabled = True
            return m

        self.miner_gen = miner_gen
        self.added_blocks = []

    @staticmethod
    def get_mining_params(rounds: Optional[int] = None):
        # guarantee target time is hit
        ret = {"target_block_time": 0.0, "is_test": True}
        if rounds is not None:
            ret["rounds"] = rounds
        return ret

    def test_mine_new_block_normal_case(self):
        async def create(retry=True):
            if len(self.added_blocks) >= 5:
                return None  # stop the game
            return RootBlock(
                RootBlockHeader(create_time=int(time.time())),
                tracking_data="{}".encode("utf-8"),
            )

        async def add(block):
            nonlocal miner
            self.added_blocks.append(block)

        for consensus in (
            ConsensusType.POW_SIMULATE,
            ConsensusType.POW_ETHASH,
            ConsensusType.POW_SHA3SHA3,
        ):
            miner = self.miner_gen(consensus, create, add)
            # should generate 5 blocks and then end
            loop = asyncio.get_event_loop()
            loop.run_until_complete(miner._mine_new_block_async())
            self.assertEqual(len(self.added_blocks), 5)

    def test_simulate_mine_handle_block_exception(self):
        i = 0

        async def create(retry=True):
            nonlocal i
            if i >= 5:
                return None
            return RootBlock(
                RootBlockHeader(create_time=int(time.time())),
                tracking_data="{}".encode("utf-8"),
            )

        async def add(block):
            nonlocal i, miner
            try:
                if i % 2 == 0:
                    raise Exception("(╯°□°）╯︵ ┻━┻")
                else:
                    self.added_blocks.append(block)
            finally:
                i += 1

        miner = self.miner_gen(ConsensusType.POW_SIMULATE, create, add)
        # only 2 blocks can be added
        loop = asyncio.get_event_loop()
        loop.run_until_complete(miner._mine_new_block_async())
        self.assertEqual(len(self.added_blocks), 2)

    def test_sha3sha3(self):
        miner = self.miner_gen(ConsensusType.POW_SHA3SHA3, None, None)
        block = RootBlock(
            RootBlockHeader(create_time=42, difficulty=5),
            tracking_data="{}".encode("utf-8"),
        )
        work = MiningWork(block.header.get_hash_for_mining(), 42, 5)
        # only process one block, which is passed in. `None` means termination right after
        miner.input_q.put((None, {}))
        miner.mine_loop(
            work,
            {"consensus_type": ConsensusType.POW_SHA3SHA3},
            miner.input_q,
            miner.output_q,
        )
        mined_res = miner.output_q.get()
        block.header.nonce = mined_res.nonce
        validate_seal(block.header, ConsensusType.POW_SHA3SHA3)

    def test_qkchash(self):
        miner = self.miner_gen(ConsensusType.POW_QKCHASH, None, None)
        block = RootBlock(
            RootBlockHeader(create_time=42, difficulty=5),
            tracking_data="{}".encode("utf-8"),
        )
        work = MiningWork(block.header.get_hash_for_mining(), 42, 5)
        # only process one block, which is passed in. `None` means termination right after
        miner.input_q.put((None, {}))
        miner.mine_loop(
            work,
            {"consensus_type": ConsensusType.POW_QKCHASH},
            miner.input_q,
            miner.output_q,
        )
        mined_res = miner.output_q.get()
        block.header.nonce = mined_res.nonce
        block.header.mixhash = mined_res.mixhash
        validate_seal(block.header, ConsensusType.POW_QKCHASH)

    def test_only_remote(self):
        async def go():
            miner = self.miner_gen(ConsensusType.POW_SHA3SHA3, None, None)
            with self.assertRaises(ValueError):
                await miner.get_work()
            with self.assertRaises(ValueError):
                await miner.submit_work(b"", 42, b"")

        loop = asyncio.get_event_loop()
        loop.run_until_complete(go())

    def test_get_work(self):
        now = 42

        async def create(retry=True):
            nonlocal now
            return RootBlock(RootBlockHeader(create_time=now, extra_data=b"{}"))

        miner = self.miner_gen(ConsensusType.POW_SHA3SHA3, create, None, remote=True)

        async def go():
            nonlocal now
            # no current work, will generate a new one
            work = await miner.get_work(now=now)
            self.assertEqual(len(work), 3)
            self.assertEqual(len(miner.work_map), 1)
            h = list(miner.work_map.keys())[0]
            self.assertEqual(work.hash, h)
            # cache hit
            now += 1
            work = await miner.get_work(now=now)
            self.assertEqual(work.hash, h)
            self.assertEqual(len(miner.work_map), 1)
            # new work if interval passed
            now += 10
            work = await miner.get_work(now=now)
            self.assertEqual(len(miner.work_map), 2)
            self.assertNotEqual(work.hash, h)
            # work map cleaned up if too much time passed
            now += 100
            await miner.get_work(now=now)
            self.assertEqual(len(miner.work_map), 1)  # only new work itself

        loop = asyncio.get_event_loop()
        loop.run_until_complete(go())

    def test_submit_work(self):
        now = 42
        block = RootBlock(
            RootBlockHeader(create_time=42, extra_data=b"{}", difficulty=5)
        )

        async def create(retry=True):
            return block

        async def add(block_to_add):
            self.added_blocks.append(block_to_add)

        miner = self.miner_gen(ConsensusType.POW_SHA3SHA3, create, add, remote=True)

        async def go():
            work = await miner.get_work(now=now)
            self.assertEqual(work.height, 0)
            self.assertEqual(work.difficulty, 5)
            # submitted block doesn't exist
            res = await miner.submit_work(b"lolwut", 0, sha3_256(b""))
            self.assertFalse(res)

            solver = DoubleSHA256(work)
            sol = solver.mine(100, 200).nonce
            self.assertGreater(sol, 100)  # ensure non-solution is tried
            non_sol = sol - 1
            # invalid pow proof
            res = await miner.submit_work(work.hash, non_sol, sha3_256(b""))
            self.assertFalse(res)
            # valid submission, also check internal state afterwards
            res = await miner.submit_work(work.hash, sol, sha3_256(b""))
            self.assertTrue(res)
            self.assertEqual(miner.work_map, {})
            self.assertEqual(len(self.added_blocks), 1)
            self.assertIsNone(miner.current_work)

        loop = asyncio.get_event_loop()
        loop.run_until_complete(go())

    def test_submit_work_with_guardian(self):
        now = 42
        block = RootBlock(
            RootBlockHeader(create_time=42, extra_data=b"{}", difficulty=1000)
        )

        async def create(retry=True):
            return block

        async def add(_):
            pass

        miner = self.miner_gen(
            ConsensusType.POW_SHA3SHA3,
            create,
            add,
            remote=True,
            # fake pk, will succeed in test but fail in real world when
            # adding the block to the root chain
            guardian_private_key=ecies.generate_privkey(),
        )

        async def go():
            for i in range(42, 100):
                work = await miner.get_work(now=now)
                self.assertEqual(work.height, 0)

                # guardian: diff 1000 -> 1, any number should work
                res = await miner.submit_work(work.hash, i, sha3_256(b""))
                self.assertTrue(res)

        loop = asyncio.get_event_loop()
        loop.run_until_complete(go())

    def test_validate_seal_with_adjusted_diff(self):
        diff = 1000
        block = RootBlock(
            RootBlockHeader(create_time=42, difficulty=diff),
            tracking_data="{}".encode("utf-8"),
        )
        block.header.nonce = 0
        with self.assertRaises(ValueError):
            validate_seal(block.header, ConsensusType.POW_SHA3SHA3)

        # significantly lowering the diff should pass
        validate_seal(block.header, ConsensusType.POW_SHA3SHA3, adjusted_diff=1)

[((83, 15, 83, 39), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((104, 8, 104, 63), 'quarkchain.cluster.miner.validate_seal', 'validate_seal', ({(104, 22, 104, 34): 'block.header', (104, 36, 104, 62): 'ConsensusType.POW_SHA3SHA3'}, {}), '(block.header, ConsensusType.POW_SHA3SHA3)', False, 'from quarkchain.cluster.miner import DoubleSHA256, Miner, MiningWork, validate_seal\n'), ((124, 8, 124, 62), 'quarkchain.cluster.miner.validate_seal', 'validate_seal', ({(124, 22, 124, 34): 'block.header', (124, 36, 124, 61): 'ConsensusType.POW_QKCHASH'}, {}), '(block.header, ConsensusType.POW_QKCHASH)', False, 'from quarkchain.cluster.miner import DoubleSHA256, Miner, MiningWork, validate_seal\n'), ((134, 15, 134, 39), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((169, 15, 169, 39), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((208, 15, 208, 39), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((242, 15, 242, 39), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((256, 8, 256, 80), 'quarkchain.cluster.miner.validate_seal', 'validate_seal', (), '', False, 'from quarkchain.cluster.miner import DoubleSHA256, Miner, MiningWork, validate_seal\n'), ((18, 16, 20, 13), 'quarkchain.cluster.miner.Miner', 'Miner', ({(19, 16, 19, 25): 'consensus', (19, 27, 19, 38): 'create_func', (19, 40, 19, 48): 'add_func', (19, 50, 19, 72): 'self.get_mining_params'}, {}), '(consensus, create_func, add_func, self.get_mining_params, **kwargs)', False, 'from quarkchain.cluster.miner import DoubleSHA256, Miner, MiningWork, validate_seal\n'), ((55, 19, 55, 43), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ({}, {}), '()', False, 'import asyncio\n'), ((90, 12, 90, 57), 'quarkchain.core.RootBlockHeader', 'RootBlockHeader', (), '', False, 'from quarkchain.core import RootBlock, RootBlockHeader\n'), ((109, 12, 109, 57), 'quarkchain.core.RootBlockHeader', 'RootBlockHeader', (), '', False, 'from quarkchain.core import RootBlock, RootBlockHeader\n'), ((175, 12, 175, 75), 'quarkchain.core.RootBlockHeader', 'RootBlockHeader', (), '', False, 'from quarkchain.core import RootBlock, RootBlockHeader\n'), ((194, 21, 194, 39), 'quarkchain.cluster.miner.DoubleSHA256', 'DoubleSHA256', ({(194, 34, 194, 38): 'work'}, {}), '(work)', False, 'from quarkchain.cluster.miner import DoubleSHA256, Miner, MiningWork, validate_seal\n'), ((214, 12, 214, 78), 'quarkchain.core.RootBlockHeader', 'RootBlockHeader', (), '', False, 'from quarkchain.core import RootBlock, RootBlockHeader\n'), ((248, 12, 248, 60), 'quarkchain.core.RootBlockHeader', 'RootBlockHeader', (), '', False, 'from quarkchain.core import RootBlock, RootBlockHeader\n'), ((253, 12, 253, 67), 'quarkchain.cluster.miner.validate_seal', 'validate_seal', ({(253, 26, 253, 38): 'block.header', (253, 40, 253, 66): 'ConsensusType.POW_SHA3SHA3'}, {}), '(block.header, ConsensusType.POW_SHA3SHA3)', False, 'from quarkchain.cluster.miner import DoubleSHA256, Miner, MiningWork, validate_seal\n'), ((142, 29, 142, 79), 'quarkchain.core.RootBlockHeader', 'RootBlockHeader', (), '', False, 'from quarkchain.core import RootBlock, RootBlockHeader\n'), ((230, 33, 230, 57), 'quarkchain.p2p.ecies.generate_privkey', 'ecies.generate_privkey', ({}, {}), '()', False, 'from quarkchain.p2p import ecies\n'), ((191, 56, 191, 69), 'quarkchain.utils.sha3_256', 'sha3_256', ({(191, 65, 191, 68): "b''"}, {}), "(b'')", False, 'from quarkchain.utils import sha3_256\n'), ((199, 62, 199, 75), 'quarkchain.utils.sha3_256', 'sha3_256', ({(199, 71, 199, 74): "b''"}, {}), "(b'')", False, 'from quarkchain.utils import sha3_256\n'), ((202, 58, 202, 71), 'quarkchain.utils.sha3_256', 'sha3_256', ({(202, 67, 202, 70): "b''"}, {}), "(b'')", False, 'from quarkchain.utils import sha3_256\n'), ((239, 60, 239, 73), 'quarkchain.utils.sha3_256', 'sha3_256', ({(239, 69, 239, 72): "b''"}, {}), "(b'')", False, 'from quarkchain.utils import sha3_256\n'), ((40, 48, 40, 59), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((67, 48, 67, 59), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n')]