import json
from typing import Dict, TextIO, Union

from .etheno import EthenoPlugin
from .utils import format_hex_address
from .client import JSONRPCError

# source: https://ethereum.stackexchange.com/a/83855
import rlp
from eth_typing import HexStr
from eth_utils import keccak, to_bytes
from rlp.sedes import Binary, big_endian_int, binary
from web3 import Web3
from web3.auto import w3


class Transaction(rlp.Serializable):
    fields = [
        ("nonce", big_endian_int),
        ("gas_price", big_endian_int),
        ("gas", big_endian_int),
        ("to", Binary.fixed_length(20, allow_empty=True)),
        ("value", big_endian_int),
        ("data", binary),
        ("v", big_endian_int),
        ("r", big_endian_int),
        ("s", big_endian_int),
    ]


def hex_to_bytes(data: str) -> bytes:
    return to_bytes(hexstr=HexStr(data))


def decode_raw_tx(raw_tx: str):
    tx_bytes = hex_to_bytes(raw_tx)
    tx = rlp.decode(tx_bytes, Transaction)
    hash_tx = Web3.toHex(keccak(tx_bytes))
    from_ = w3.eth.account.recover_transaction(raw_tx)
    to = w3.toChecksumAddress(tx.to) if tx.to else None
    data = w3.toHex(tx.data)
    r = hex(tx.r)
    s = hex(tx.s)
    chain_id = (tx.v - 35) // 2 if tx.v % 2 else (tx.v - 36) // 2
    return {
        "txHash": hash_tx,
        "from": from_,
        "to": to,
        "nonce": tx.nonce,
        "gas": hex(tx.gas),
        "gasPrice": hex(tx.gas_price),
        "value": hex(tx.value),
        "data": data,
        "chainId": chain_id,
        "r": r,
        "s": s,
        "v": tx.v,
    }


class JSONExporter:
    def __init__(self, out_stream: Union[str, TextIO]):
        self._was_path = isinstance(out_stream, str)
        if self._was_path:
            self.output = open(out_stream, "w", encoding="utf8")
        else:
            self.output = out_stream
        self.output.write("[")
        self._count = 0
        self._finalized = False

    def finalize(self):
        if self._finalized:
            return
        if self._count:
            self.output.write("\n")
        self.output.write("]")
        self.output.flush()
        if self._was_path:
            self.output.close()
        self._finalized = True

    def write_entry(self, entry):
        if self._finalized:
            return
        if self._count > 0:
            self.output.write(",")
        self._count += 1
        self.output.write("\n")
        json.dump(entry, self.output)
        self.output.flush()


class JSONRPCExportPlugin(EthenoPlugin):
    def __init__(self, out_stream: Union[str, TextIO]):
        self._exporter = JSONExporter(out_stream)

    def after_post(self, post_data, client_results):
        self._exporter.write_entry([post_data, client_results])

    def finalize(self):
        self._exporter.finalize()
        if hasattr(self._exporter.output, "name"):
            self.logger.info(
                f"Raw JSON RPC messages dumped to {self._exporter.output.name}"
            )


class EventSummaryPlugin(EthenoPlugin):
    def __init__(self):
        self._transactions: Dict[
            int, Dict[str, object]
        ] = {}  # Maps transaction hashes to their eth_sendTransaction arguments

    def handle_contract_created(
        self,
        creator_address: str,
        contract_address: str,
        gas_used: str,
        gas_price: str,
        data: str,
        value: str,
    ):
        self.logger.info(
            f"Contract created at {contract_address} with {(len(data)-2)//2} bytes of data by account {creator_address} for {gas_used} gas with a gas price of {gas_price}"
        )

    def handle_function_call(
        self,
        from_address: str,
        to_address: str,
        gas_used: str,
        gas_price: str,
        data: str,
        value: str,
    ):
        self.logger.info(
            f"Function call with {value} wei from {from_address} to {to_address} with {(len(data)-2)//2} bytes of data for {gas_used} gas with a gas price of {gas_price}"
        )

    def handle_unlogged_transactions(self):
        unlogged_transactions = dict(
            filter(lambda txn: txn[1]["is_logged"] == False, self._transactions.items())
        )
        for (tx_hash, txn) in unlogged_transactions.items():
            post_data = self._etheno.get_transaction_receipt_request(tx_hash)
            self._etheno.post(post_data)

    def after_post(self, post_data, result):
        if len(result):
            result = result[0]
        if "method" not in post_data:
            return
        # Fixes bug that occurs when a JSONRPCError is attempted to be logged
        if isinstance(result, JSONRPCError):
            self.logger.info(
                f"Received a JSON RPC Error when logging transaction...skipping event logging"
            )
            return

        elif (
            post_data["method"] == "eth_sendTransaction"
            or post_data["method"] == "eth_sendRawTransaction"
        ) and "result" in result:
            transaction_hash = result["result"]
            # Add a boolean to check at shutdown whether everything has been logged.
            if post_data["method"] == "eth_sendRawTransaction":
                self._transactions[transaction_hash] = {
                    "transaction": decode_raw_tx(post_data["params"][0]),
                    "is_logged": False,
                }
            else:
                self._transactions[transaction_hash] = {
                    "transaction": post_data["params"][0],
                    "is_logged": False,
                }
        elif post_data["method"] == "evm_mine":
            self.handle_increase_block_number()
        elif post_data["method"] == "evm_increaseTime":
            self.handle_increase_block_timestamp(post_data["params"][0])
        elif post_data["method"] == "eth_getTransactionReceipt":
            transaction_hash = post_data["params"][0]
            if transaction_hash not in self._transactions:
                self.logger.error(
                    f'Received transaction receipt {result} for unknown transaction hash {post_data["params"][0]}'
                )
                return
            (original_transaction, is_logged) = (
                self._transactions[transaction_hash]["transaction"],
                self._transactions[transaction_hash]["is_logged"],
            )
            # Check if it was logged already
            if is_logged:
                self.logger.debug(
                    f"Transaction hash {transaction_hash} has already been logged. This should not happen."
                )
                return
            if (
                "value" not in original_transaction
                or original_transaction["value"] is None
            ):
                value = "0x0"
            else:
                value = original_transaction["value"]
            if "to" not in result["result"] or result["result"]["to"] is None:
                # this transaction is creating a contract:
                # TODO: key errors are likely here...need to figure out a better way to do error handling
                # TODO: log a warning about non-zero ether values
                contract_address = result["result"]["contractAddress"]
                self.handle_contract_created(
                    original_transaction["from"],
                    contract_address,
                    result["result"]["gasUsed"],
                    result["result"]["effectiveGasPrice"],
                    original_transaction["data"],
                    value,
                )
            else:
                self.handle_function_call(
                    original_transaction["from"],
                    original_transaction["to"],
                    result["result"]["gasUsed"],
                    result["result"]["effectiveGasPrice"],
                    original_transaction["data"]
                    if "data" in original_transaction
                    else "0x",
                    value,
                )
            # Transaction has been logged successfully
            self._transactions[transaction_hash]["is_logged"] = True


class EventSummaryExportPlugin(EventSummaryPlugin):
    def __init__(self, out_stream: Union[str, TextIO]):
        super().__init__()
        self._exporter = JSONExporter(out_stream)

    def run(self):
        for address in self.etheno.accounts:
            self._exporter.write_entry(
                {"event": "AccountCreated", "address": format_hex_address(address)}
            )
        super().run()

    def handle_increase_block_number(self):
        self._exporter.write_entry(
            {"event": "BlockMined", "number_increment": "1", "timestamp_increment": "0"}
        )

    def handle_increase_block_timestamp(self, number: str):
        self._exporter.write_entry(
            {
                "event": "BlockMined",
                "number_increment": "0",
                "timestamp_increment": str(number),
            }
        )

    def handle_contract_created(
        self,
        creator_address: str,
        contract_address: str,
        gas_used: str,
        gas_price: str,
        data: str,
        value: str,
    ):
        self._exporter.write_entry(
            {
                "event": "ContractCreated",
                "from": creator_address,
                "contract_address": contract_address,
                "gas_used": gas_used,
                "gas_price": gas_price,
                "data": data,
                "value": value,
            }
        )
        super().handle_contract_created(
            creator_address, contract_address, gas_used, gas_price, data, value
        )

    def handle_function_call(
        self,
        from_address: str,
        to_address: str,
        gas_used: str,
        gas_price: str,
        data: str,
        value: str,
    ):
        self._exporter.write_entry(
            {
                "event": "FunctionCall",
                "from": from_address,
                "to": to_address,
                "gas_used": gas_used,
                "gas_price": gas_price,
                "data": data,
                "value": value,
            }
        )
        super().handle_function_call(
            from_address, to_address, gas_used, gas_price, data, value
        )

    def finalize(self):
        super().handle_unlogged_transactions()
        self._exporter.finalize()
        if hasattr(self._exporter.output, "name"):
            self.logger.info(
                f"Event summary JSON saved to {self._exporter.output.name}"
            )
