import os
import string
import time
import unittest
from typing import Any

import nacl.signing

from rainbowadn.chain.blockchain import BlockChainFactory
from rainbowadn.chain.chaincollectioninterface import ChainCollectionInterface
from rainbowadn.chain.reduction.reductionchainmetafactory import ReductionChainMetaFactory
from rainbowadn.data.atomic.integer import Integer
from rainbowadn.data.atomic.plain import Plain
from rainbowadn.data.collection.trees.binary.activebinarytree import ActiveBinaryTree
from rainbowadn.data.collection.trees.binary.avl import AVLBTBP
from rainbowadn.data.collection.trees.comparison.comparator import Replace
from rainbowadn.data.collection.trees.comparison.plaincomparator import PlainComparator
from rainbowadn.encryption.encrypted import Encrypted
from rainbowadn.hashing.hashpoint import HashPoint
from rainbowadn.hashing.nullability.notnull import NotNull
from rainbowadn.hashing.resolvermetaorigin import ResolverMetaOrigin
from rainbowadn.testing.dictresolver import DictResolver
from rainbowadn.v13.algo import MINT_CONST
from rainbowadn.v13.bankchain import BankChain
from rainbowadn.v13.subject import Subject
from rainbowadn.v13.transaction import CoinData, Transaction
from rainbowadn.wrisbt.wrisbtchainprotocol import WrisbtChainProtocol
from rainbowadn.wrisbt.wrisbtparametres import WrisbtParametres
from rainbowadn.wrisbt.wrisbtroot import WrisbtRoot


class TestAll(unittest.TestCase):
    def test_bankchain(self):
        with self.subTest('setup'):
            dr = DictResolver()
        with self.subTest('create empty'):
            bank = BankChain.empty(ReductionChainMetaFactory())
        with self.subTest('prepare transactions'):
            key_0 = nacl.signing.SigningKey.generate()
            transaction_0 = Transaction.make(
                [],
                [CoinData.of(Subject(key_0.verify_key), 1_000_000)],
                []
            )
            coin_0, coin_1 = transaction_0.coins(MINT_CONST, NotNull(HashPoint.of(Subject(key_0.verify_key))))
        with self.subTest('add transactions'):
            bank = bank.adds(
                [
                    transaction_0,
                    Transaction.make(
                        [coin_1],
                        [CoinData.of(Subject(nacl.signing.SigningKey.generate().verify_key), 10_000)],
                        [key_0]
                    ),
                ]
            )
        with self.subTest('add empty'):
            bank = bank.adds(
                []
            )
            print(bank)
        with self.subTest('verify'):
            assert bank.verify()
        with self.subTest('recover'):
            dr.save(HashPoint.of(bank.reference))
            bank = BankChain.from_reference(
                ReductionChainMetaFactory(), ResolverMetaOrigin(dr).migrate(HashPoint.of(bank.reference)).resolve()
            )
            print(bank)
        with self.subTest('verify'):
            assert bank.verify()

    def test_wrisbt(self):
        with self.subTest('setup'):
            stoptime = time.process_time()

            def measure(message: str) -> float:
                nonlocal stoptime
                now = time.process_time()
                delta = now - stoptime
                print(message, delta)
                stoptime = now
                return delta

            dr = DictResolver()
            n = 10000
        with self.subTest('create empty'):
            btree = WrisbtRoot.empty(WrisbtParametres(1, 5))
            measure('init')
        with self.subTest('add keys', n=n):
            for _ in range(n):
                key = os.urandom(5)
                assert not btree.contains(key)
                btree = btree.add(key)
                assert btree.contains(key)
            measure('add')
        with self.subTest('save'):
            dr.save(HashPoint.of(btree))
            measure('save')
        with self.subTest('resolve and iterate'):
            btree = ResolverMetaOrigin(dr).migrate(HashPoint.of(btree)).resolve()
            assert len(btree.keys()) == n
            print(btree.height)
            measure('resolve and iterate')
        with self.subTest('resolve and add', n=n):
            for _ in range(n):
                key = os.urandom(5)
                assert not btree.contains(key)
                btree = btree.add(key)
                assert btree.contains(key)
            print(btree.height)
            measure('resolve and add')

    def test_wrisbt_index(self):
        with self.subTest('create empty'):
            chain: ChainCollectionInterface[Any, Plain, WrisbtRoot] = BlockChainFactory(
                WrisbtChainProtocol(Plain.factory(), 2).loose()
            ).empty().loose()
        with self.subTest('fill'):
            for _ in range(1000):
                chain = chain.add(HashPoint.of(Plain(os.urandom(16))))
        with self.subTest('check'):
            assert chain
        with self.subTest('measure height'):
            reference = chain.actual_state().reference
            assert not reference.null()
            print(reference.resolve().resolve().height)

    def test_avl(self):
        tree: ActiveBinaryTree[Plain, Integer] = ActiveBinaryTree.empty(
            AVLBTBP(PlainComparator(Replace())), Plain.factory()
        )
        for i in range(26):
            tree = tree.add(HashPoint.of(Plain(bytes([ord('A') + i]))))
        print(tree.reference.str(0))

    def test_encryption(self):
        with self.subTest('setup'):
            key = b'a' * 32
            dr = DictResolver()
        with self.subTest('create empty'):
            tree: ActiveBinaryTree[Plain, Integer] = ActiveBinaryTree.empty(
                AVLBTBP(PlainComparator(Replace())), Plain.factory()
            )
        with self.subTest('fill'):
            for char in string.ascii_uppercase:
                tree = tree.add(HashPoint.of(Plain(char.encode())))
            print(tree.reference.str(0))
        with self.subTest('encrypt'):
            target = tree.reference
            target = Encrypted.encrypt(target, key).decrypted
            print(Encrypted.ecc)
            tree = tree.create(target)
            print(tree.reference.str(0))
        with self.subTest('alter'):
            tree = tree.add(HashPoint.of(Plain(b'NEWKEY')))
            tree = tree.remove(HashPoint.of(Plain(b'Q')))
            print(tree.reference.str(0))
        with self.subTest('encrypt and migrate'):
            target = tree.reference
            eeed = Encrypted.encrypt(target, key)
            print(Encrypted.ecc)
            dr.save(HashPoint.of(eeed))
            print(ResolverMetaOrigin(dr).migrate(HashPoint.of(eeed)).resolve().decrypted.str(0))
        with self.subTest('re-encrypt'):
            new_key = b'b' * 32
            target = eeed.decrypted
            Encrypted.encrypt(target, new_key)
            print(Encrypted.ecc)