import pytest

from blspy import G2Element

from stor.clvm.spend_sim import SpendSim, SimClient
from stor.types.blockchain_format.sized_bytes import bytes32
from stor.types.blockchain_format.program import Program
from stor.types.spend_bundle import SpendBundle
from stor.types.coin_spend import CoinSpend


class TestSpendSim:
    @pytest.mark.asyncio
    async def test_farming(self):
        try:
            sim = await SpendSim.create()
            for i in range(0, 5):
                await sim.farm_block()

            assert len(sim.blocks) == 5
            assert sim.blocks[-1].height == 4
            assert sim.block_records[0].reward_claims_incorporated[0].amount == 18375000000000000000
        finally:
            await sim.close()

    @pytest.mark.asyncio
    async def test_rewind(self):
        try:
            sim = await SpendSim.create()
            for i in range(0, 5):
                await sim.farm_block()

            save_height = sim.get_height()
            await sim.farm_block()
            await sim.rewind(save_height)

            assert len(sim.blocks) == 5
            assert sim.blocks[-1].height == 4
        finally:
            await sim.close()


class TestSimClient:
    @pytest.mark.asyncio
    async def test_all_endpoints(self):
        try:
            sim = await SpendSim.create()
            sim_client = SimClient(sim)
            for i in range(0, 5):
                await sim.farm_block()
            await sim.farm_block(bytes32([0] * 32))
            await sim.farm_block(bytes32([1] * 32))
            for i in range(0, 5):
                await sim.farm_block()

            # get_coin_records_by_puzzle_hash
            coin_records = await sim_client.get_coin_records_by_puzzle_hash(bytes32([0] * 32))
            coin_record_name = coin_records[0].coin.name()
            assert len(coin_records) == 2

            coin_records = await sim_client.get_coin_records_by_puzzle_hash(
                bytes32([0] * 32), start_height=0, end_height=2
            )
            assert len(coin_records) == 0

            # get_coin_records_by_puzzle_hashes
            coin_records = await sim_client.get_coin_records_by_puzzle_hashes([bytes32([0] * 32), bytes32([1] * 32)])
            assert len(coin_records) == 4

            coin_records = await sim_client.get_coin_records_by_puzzle_hashes(
                [bytes32([0] * 32), bytes32([1] * 32)], start_height=0, end_height=2
            )
            assert len(coin_records) == 0

            # get_coin_record_by_name
            assert await sim_client.get_coin_record_by_name(coin_record_name)

            # get_block_records
            block_records = await sim_client.get_block_records(0, 5)
            assert len(block_records) == 5

            # get_block_record_by_height
            block_record = await sim_client.get_block_record_by_height(0)
            assert block_record
            assert block_record == block_records[0]

            # get_block_record
            same_block_record = await sim_client.get_block_record(block_record.header_hash)
            assert same_block_record == block_record

            # get_block
            full_block = await sim_client.get_block(block_record.header_hash)
            assert full_block.transactions_generator is None

            # get_all_block
            full_blocks = await sim_client.get_all_block(0, 5)
            assert full_blocks[0] == full_block

            # push_tx
            puzzle_hash = bytes.fromhex(
                "9dcf97a184f32623d11a73124ceb99a5709b083721e878a16d78f596718ba7b2"
            )  # Program.to(1)
            await sim.farm_block(puzzle_hash)
            spendable_coin = await sim_client.get_coin_records_by_puzzle_hash(puzzle_hash)
            spendable_coin = spendable_coin[0].coin
            bundle = SpendBundle(
                [
                    CoinSpend(
                        spendable_coin,
                        Program.to(1),
                        Program.to([[51, puzzle_hash, 1]]),
                    )
                ],
                G2Element(),
            )
            result, error = await sim_client.push_tx(bundle)

            # get_all_mempool_tx_ids
            mempool_items = await sim_client.get_all_mempool_tx_ids()
            assert len(mempool_items) == 1

            # get_mempool_item_by_tx_id
            mempool_item = await sim_client.get_mempool_item_by_tx_id(mempool_items[0])
            assert mempool_item

            # get_all_mempool_items
            mempool_items = await sim_client.get_all_mempool_items()
            assert len(mempool_items) == 1

            # get_additions_and_removals
            await sim.farm_block()
            latest_block = sim.block_records[-1]
            additions, removals = await sim_client.get_additions_and_removals(latest_block.header_hash)
            assert additions
            assert removals

            # get_puzzle_and_solution
            coin_solution = await sim_client.get_puzzle_and_solution(spendable_coin.name(), latest_block.height)
            assert coin_solution
        finally:
            await sim.close()
