from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass

from chia_rs import G1Element, G2Element, TransactionsInfo, serialized_length
from chia_rs.sized_bytes import bytes32
from chia_rs.sized_ints import uint32
from chiabip158 import PyBIP158

from chia.types.blockchain_format.coin import Coin
from chia.types.blockchain_format.serialized_program import SerializedProgram


def skip_list(buf: memoryview, skip_item: Callable[[memoryview], memoryview]) -> memoryview:
    n = int.from_bytes(buf[:4], "big", signed=False)
    buf = buf[4:]
    for _ in range(n):
        buf = skip_item(buf)
    return buf


def skip_bytes(buf: memoryview) -> memoryview:
    n = int.from_bytes(buf[:4], "big", signed=False)
    buf = buf[4:]
    assert n >= 0
    return buf[n:]


def skip_optional(buf: memoryview, skip_item: Callable[[memoryview], memoryview]) -> memoryview:
    if buf[0] == 0:
        return buf[1:]
    assert buf[0] == 1
    return skip_item(buf[1:])


def skip_bytes32(buf: memoryview) -> memoryview:
    return buf[32:]


def skip_uint32(buf: memoryview) -> memoryview:
    return buf[4:]


def skip_uint64(buf: memoryview) -> memoryview:
    return buf[8:]


def skip_uint128(buf: memoryview) -> memoryview:
    return buf[16:]


def skip_uint8(buf: memoryview) -> memoryview:
    return buf[1:]


def skip_bool(buf: memoryview) -> memoryview:
    assert buf[0] in {0, 1}
    return buf[1:]


# def skip_class_group_element(buf: memoryview) -> memoryview:
#    return buf[100:]  # bytes100


def skip_vdf_info(buf: memoryview) -> memoryview:
    #    buf = skip_bytes32(buf)
    #    buf = skip_uint64(buf)
    #    return skip_class_group_element(buf)
    return buf[32 + 8 + 100 :]


def skip_vdf_proof(buf: memoryview) -> memoryview:
    buf = skip_uint8(buf)  # witness_type
    buf = skip_bytes(buf)  # witness
    return skip_bool(buf)  # normalized_to_identity


def skip_challenge_chain_sub_slot(buf: memoryview) -> memoryview:
    buf = skip_vdf_info(buf)
    buf = skip_optional(buf, skip_bytes32)  # infused challenge chain sub skit hash
    buf = skip_optional(buf, skip_bytes32)  # subepoch_summary_hash
    buf = skip_optional(buf, skip_uint64)  # new_sub_slot_iters
    return skip_optional(buf, skip_uint64)  # new_difficulty


def skip_infused_challenge_chain(buf: memoryview) -> memoryview:
    return skip_vdf_info(buf)  # infused_challenge_chain_end_of_slot_vdf


def skip_reward_chain_sub_slot(buf: memoryview) -> memoryview:
    buf = skip_vdf_info(buf)  # end_of_slot_vdf
    buf = skip_bytes32(buf)  # challenge_chain_sub_slot_hash
    buf = skip_optional(buf, skip_bytes32)  # infused_challenge_chain_sub_slot_hash
    return skip_uint8(buf)


def skip_sub_slot_proofs(buf: memoryview) -> memoryview:
    buf = skip_vdf_proof(buf)  # challenge_chain_slot_proof
    buf = skip_optional(buf, skip_vdf_proof)  # infused_challenge_chain_slot_proof
    return skip_vdf_proof(buf)  # reward_chain_slot_proof


def skip_end_of_sub_slot_bundle(buf: memoryview) -> memoryview:
    buf = skip_challenge_chain_sub_slot(buf)
    buf = skip_optional(buf, skip_infused_challenge_chain)
    buf = skip_reward_chain_sub_slot(buf)
    return skip_sub_slot_proofs(buf)


def skip_g1_element(buf: memoryview) -> memoryview:
    return buf[G1Element.SIZE :]


def skip_g2_element(buf: memoryview) -> memoryview:
    return buf[G2Element.SIZE :]


def skip_proof_of_space(buf: memoryview) -> memoryview:
    buf = skip_bytes32(buf)  # challenge
    buf = skip_optional(buf, skip_g1_element)  # pool_public_key
    buf = skip_optional(buf, skip_bytes32)  # pool_contract_puzzle_hash
    buf = skip_g1_element(buf)  # plot_public_key
    buf = skip_uint8(buf)  # size
    return skip_bytes(buf)  # proof


def skip_reward_chain_block(buf: memoryview) -> memoryview:
    buf = skip_uint128(buf)  # weight
    buf = skip_uint32(buf)  # height
    buf = skip_uint128(buf)  # total_iters
    buf = skip_uint8(buf)  # signage_point_index
    buf = skip_bytes32(buf)  # pos_ss_cc_challenge_hash

    buf = skip_proof_of_space(buf)  # proof_of_space
    buf = skip_optional(buf, skip_vdf_info)  # challenge_chain_sp_vdf
    buf = skip_g2_element(buf)  # challenge_chain_sp_signature
    buf = skip_vdf_info(buf)  # challenge_chain_ip_vdf
    buf = skip_optional(buf, skip_vdf_info)  # reward_chain_sp_vdf
    buf = skip_g2_element(buf)  # reward_chain_sp_signature
    buf = skip_vdf_info(buf)  # reward_chain_ip_vdf

    discriminant = buf[0]
    buf = buf[1:]
    has_vdf_info = (discriminant & 1) != 0
    has_mmr_root = (discriminant & 2) != 0
    if has_vdf_info:
        buf = skip_vdf_info(buf)  # infused_challenge_chain_ip_vdf
    if has_mmr_root:
        buf = skip_bytes32(buf)  # header_mmr_root
    return skip_bool(buf)  # is_transaction_block


def skip_pool_target(buf: memoryview) -> memoryview:
    # buf = skip_bytes32(buf)  # puzzle_hash
    # return skip_uint32(buf)  # max_height
    return buf[32 + 4 :]


def skip_foliage_block_data(buf: memoryview) -> memoryview:
    buf = skip_bytes32(buf)  # unfinished_reward_block_hash
    buf = skip_pool_target(buf)  # pool_target
    buf = skip_optional(buf, skip_g2_element)  # pool_signature
    buf = skip_bytes32(buf)  # farmer_reward_puzzle_hash
    return skip_bytes32(buf)  # extension_data


def skip_foliage(buf: memoryview) -> tuple[memoryview, bool]:
    buf = skip_bytes32(buf)  # prev_block_hash
    buf = skip_bytes32(buf)  # reward_block_hash
    buf = skip_foliage_block_data(buf)  # foliage_block_data
    buf = skip_g2_element(buf)  # foliage_block_data_signature
    is_tx_block = buf[0] != 0
    buf = skip_optional(buf, skip_bytes32)  # foliage_transaction_block_hash
    return (skip_optional(buf, skip_g2_element), is_tx_block)  # foliage_transaction_block_signature


def prev_hash_from_foliage(buf: memoryview) -> tuple[memoryview, bytes32]:
    prev_hash = buf[:32]  # prev_block_hash
    buf = skip_bytes32(buf)  # prev_block_hash
    buf = skip_bytes32(buf)  # reward_block_hash
    buf = skip_foliage_block_data(buf)  # foliage_block_data
    buf = skip_g2_element(buf)  # foliage_block_data_signature
    buf = skip_optional(buf, skip_bytes32)  # foliage_transaction_block_hash
    return skip_optional(buf, skip_g2_element), bytes32(prev_hash)  # foliage_transaction_block_signature


def skip_foliage_transaction_block(buf: memoryview) -> memoryview:
    # buf = skip_bytes32(buf)  # prev_transaction_block_hash
    # buf = skip_uint64(buf)  # timestamp
    # buf = skip_bytes32(buf)  # filter_hash
    # buf = skip_bytes32(buf)  # additions_root
    # buf = skip_bytes32(buf)  # removals_root
    # return skip_bytes32(buf)  # transactions_info_hash
    return buf[32 + 8 + 32 + 32 + 32 + 32 :]


def skip_coin(buf: memoryview) -> memoryview:
    # buf = skip_bytes32(buf)  # parent_coin_info
    # buf = skip_bytes32(buf)  # puzzle_hash
    # return skip_uint64(buf)  # amount
    return buf[32 + 32 + 8 :]


def skip_transactions_info(buf: memoryview) -> memoryview:
    # buf = skip_bytes32(buf)  # generator_root
    # buf = skip_bytes32(buf)  # generator_refs_root
    # buf = skip_g2_element(buf)  # aggregated_signature
    # buf = skip_uint64(buf)  # fees
    # buf = skip_uint64(buf)  # cost
    buf = buf[32 + 32 + G2Element.SIZE + 8 + 8 :]
    return skip_list(buf, skip_coin)


def generator_from_block(buf: memoryview) -> bytes | None:
    buf = skip_list(buf, skip_end_of_sub_slot_bundle)  # finished_sub_slots
    buf = skip_reward_chain_block(buf)  # reward_chain_block
    buf = skip_optional(buf, skip_vdf_proof)  # challenge_chain_sp_proof
    buf = skip_vdf_proof(buf)  # challenge_chain_ip_proof
    buf = skip_optional(buf, skip_vdf_proof)  # reward_chain_sp_proof
    buf = skip_vdf_proof(buf)  # reward_chain_ip_proof
    buf = skip_optional(buf, skip_vdf_proof)  # infused_challenge_chain_ip_proof
    buf, _ = skip_foliage(buf)  # foliage
    buf = skip_optional(buf, skip_foliage_transaction_block)  # foliage_transaction_block
    buf = skip_optional(buf, skip_transactions_info)  # transactions_info

    # this is the transactions_generator optional
    if buf[0] == 0:
        return None

    buf = buf[1:]
    length = serialized_length(buf)
    return bytes(buf[:length])


# this implements the BlockInfo protocol
@dataclass(frozen=True)
class GeneratorBlockInfo:
    prev_header_hash: bytes32
    transactions_generator: SerializedProgram | None
    transactions_generator_ref_list: list[uint32]


def block_info_from_block(buf: memoryview) -> GeneratorBlockInfo:
    buf = skip_list(buf, skip_end_of_sub_slot_bundle)  # finished_sub_slots
    buf = skip_reward_chain_block(buf)  # reward_chain_block
    buf = skip_optional(buf, skip_vdf_proof)  # challenge_chain_sp_proof
    buf = skip_vdf_proof(buf)  # challenge_chain_ip_proof
    buf = skip_optional(buf, skip_vdf_proof)  # reward_chain_sp_proof
    buf = skip_vdf_proof(buf)  # reward_chain_ip_proof
    buf = skip_optional(buf, skip_vdf_proof)  # infused_challenge_chain_ip_proof
    buf, prev_hash = prev_hash_from_foliage(buf)  # foliage
    buf = skip_optional(buf, skip_foliage_transaction_block)  # foliage_transaction_block
    buf = skip_optional(buf, skip_transactions_info)  # transactions_info

    # this is the transactions_generator optional
    generator = None
    if buf[0] != 0:
        buf = buf[1:]
        length = serialized_length(buf)
        generator = SerializedProgram.from_bytes(bytes(buf[:length]))
        buf = buf[length:]
    else:
        buf = buf[1:]

    refs_length = uint32.from_bytes(buf[:4])
    buf = buf[4:]

    refs = []
    for i in range(refs_length):
        refs.append(uint32.from_bytes(buf[:4]))
        buf = buf[4:]

    return GeneratorBlockInfo(prev_hash, generator, refs)


def header_block_from_block(
    buf: memoryview, request_filter: bool = True, tx_addition_coins: list[Coin] = [], removal_names: list[bytes32] = []
) -> bytes:
    """
    if request_filter is True it will generate the BIP158 filter based on the
    additions and removals passed into the function
    """
    buf2 = buf[:]
    buf2 = skip_list(buf2, skip_end_of_sub_slot_bundle)  # finished_sub_slots
    buf2 = skip_reward_chain_block(buf2)  # reward_chain_block
    buf2 = skip_optional(buf2, skip_vdf_proof)  # challenge_chain_sp_proof
    buf2 = skip_vdf_proof(buf2)  # challenge_chain_ip_proof
    buf2 = skip_optional(buf2, skip_vdf_proof)  # reward_chain_sp_proof
    buf2 = skip_vdf_proof(buf2)  # reward_chain_ip_proof
    buf2 = skip_optional(buf2, skip_vdf_proof)  # infused_challenge_chain_ip_proof
    buf2, is_transaction_block = skip_foliage(buf2)  # foliage
    buf2 = skip_optional(buf2, skip_foliage_transaction_block)  # foliage_transaction_block

    transactions_info: TransactionsInfo | None = None
    transactions_info_optional: bytes = bytes([0])
    encoded_filter = b"\x00"

    # this is the transactions_info optional
    if buf2[0] == 0:
        transactions_info_optional = bytes([0])
    else:
        transactions_info_optional = bytes([1])
        transactions_info, _advance = TransactionsInfo.parse_rust(buf2[1:])

    if request_filter:
        byte_array_tx: list[bytearray] = []
        if is_transaction_block and transactions_info:
            addition_coins = tx_addition_coins + list(transactions_info.reward_claims_incorporated)
            for coin in addition_coins:
                byte_array_tx.append(bytearray(coin.puzzle_hash))
            for name in removal_names:
                byte_array_tx.append(bytearray(name))

        bip158: PyBIP158 = PyBIP158(byte_array_tx)
        encoded_filter = bytes(bip158.GetEncoded())

    # Takes everything up to but not including transactions info
    header_block: bytes = bytes(buf[: (len(buf) - len(buf2))])
    # Transactions filter, potentially with added / removal coins
    header_block += (len(encoded_filter)).to_bytes(4, "big") + encoded_filter
    # Add transactions info
    header_block += transactions_info_optional
    if transactions_info is not None:
        header_block += bytes(transactions_info)

    return header_block


def get_height_and_tx_status_from_block(buf: memoryview) -> tuple[uint32, bool]:
    """
    Returns the height of the block and whether it's a transaction block or not.
    We're considering the block as a transaction block if transactions_info is not None.
    """
    reward_chain_block_buf = skip_list(buf, skip_end_of_sub_slot_bundle)  # finished_sub_slots
    # Get the block height from reward_chain_block_buf
    height_buf = skip_uint128(reward_chain_block_buf)  # weight
    height = uint32.from_bytes(height_buf[:4])
    # Let's continue down to transactions_info
    buf = skip_reward_chain_block(reward_chain_block_buf)  # reward_chain_block
    buf = skip_optional(buf, skip_vdf_proof)  # challenge_chain_sp_proof
    buf = skip_vdf_proof(buf)  # challenge_chain_ip_proof
    buf = skip_optional(buf, skip_vdf_proof)  # reward_chain_sp_proof
    buf = skip_vdf_proof(buf)  # reward_chain_ip_proof
    buf = skip_optional(buf, skip_vdf_proof)  # infused_challenge_chain_ip_proof
    buf, is_tx_block = skip_foliage(buf)  # foliage
    buf = skip_optional(buf, skip_foliage_transaction_block)  # foliage_transaction_block
    return height, is_tx_block