"""
Create a transition tool for the given fork.
"""

import argparse
import fnmatch
import json
import os
from contextlib import AbstractContextManager
from typing import Any, Final, TextIO, Tuple, Type, TypeVar

from ethereum_rlp import rlp
from ethereum_types.numeric import U64, U256, Uint
from typing_extensions import override

from ethereum import trace
from ethereum.exceptions import EthereumException, InvalidBlock
from ethereum.fork_criteria import ByBlockNumber, ByTimestamp, Unscheduled
from ethereum_spec_tools.forks import Hardfork, TemporaryHardfork

from ..loaders.fixture_loader import Load
from ..utils import (
    FatalError,
    find_fork,
    get_stream_logger,
    parse_hex_or_int,
)
from .env import Env
from .evm_trace.count import CountTracer
from .evm_trace.eip3155 import Eip3155Tracer
from .evm_trace.group import GroupTracer
from .t8n_types import Alloc, Result, Txs

T = TypeVar("T")


def t8n_arguments(subparsers: argparse._SubParsersAction) -> None:
    """
    Adds the arguments for the t8n tool subparser.
    """
    t8n_parser = subparsers.add_parser("t8n", help="This is the t8n tool.")

    t8n_parser.add_argument(
        "--input.alloc", dest="input_alloc", type=str, default="alloc.json"
    )
    t8n_parser.add_argument(
        "--input.env", dest="input_env", type=str, default="env.json"
    )
    t8n_parser.add_argument(
        "--input.txs", dest="input_txs", type=str, default="txs.json"
    )
    t8n_parser.add_argument(
        "--input.blobParams",
        dest="blob_parameters",
        type=str,
        default=None,
    )
    t8n_parser.add_argument(
        "--output.alloc", dest="output_alloc", type=str, default="alloc.json"
    )
    t8n_parser.add_argument(
        "--output.basedir", dest="output_basedir", type=str, default="."
    )
    t8n_parser.add_argument("--output.body", dest="output_body", type=str)
    t8n_parser.add_argument(
        "--output.result",
        dest="output_result",
        type=str,
        default="result.json",
    )
    t8n_parser.add_argument(
        "--state.chainid", dest="state_chainid", type=int, default=1
    )
    t8n_parser.add_argument(
        "--state.fork", dest="state_fork", type=str, default="Frontier"
    )
    t8n_parser.add_argument(
        "--state.reward", dest="state_reward", type=int, default=None
    )
    t8n_parser.add_argument("--trace", action="store_true")
    t8n_parser.add_argument("--trace.memory", action="store_true")
    t8n_parser.add_argument("--trace.nomemory", action="store_true")
    t8n_parser.add_argument("--trace.noreturndata", action="store_true")
    t8n_parser.add_argument("--trace.nostack", action="store_true")
    t8n_parser.add_argument("--trace.returndata", action="store_true")

    t8n_parser.add_argument("--opcode.count", dest="opcode_count", type=str)

    t8n_parser.add_argument("--state-test", action="store_true")


class ForkCache(AbstractContextManager):
    """
    Stores references to temporary hardforks and cleans them up when exited.
    """

    _cache: Final[dict[Tuple[object, ...], TemporaryHardfork]]

    def __init__(self) -> None:
        self._cache = {}

    @override
    def __exit__(self, *args: object, **kwargs: object) -> None:
        for fork in self._cache.values():
            fork.__exit__(*args, **kwargs)
        self._cache.clear()

    def get(
        self,
        template: Hardfork,
        fork_criteria: ByBlockNumber | ByTimestamp | Unscheduled | None = None,
        target_blob_gas_per_block: U64 | None = None,
        gas_per_blob: U64 | None = None,
        min_blob_gasprice: Uint | None = None,
        blob_base_fee_update_fraction: Uint | None = None,
        max_blob_gas_per_block: U64 | None = None,
        blob_schedule_target: U64 | None = None,
        blob_schedule_max: U64 | None = None,
    ) -> Hardfork:
        """
        Search the cache for a matching hardfork, or create one if it doesn't
        exist.
        """
        cache_key = (
            template.short_name,
            fork_criteria,
            target_blob_gas_per_block,
            gas_per_blob,
            min_blob_gasprice,
            blob_base_fee_update_fraction,
            max_blob_gas_per_block,
            blob_schedule_target,
            blob_schedule_max,
        )
        if all(x is None for x in cache_key[1:]):
            return template

        try:
            return self._cache[cache_key]
        except KeyError:
            pass

        clone = Hardfork.clone(
            template=template,
            fork_criteria=fork_criteria,
            target_blob_gas_per_block=target_blob_gas_per_block,
            gas_per_blob=gas_per_blob,
            min_blob_gasprice=min_blob_gasprice,
            blob_base_fee_update_fraction=blob_base_fee_update_fraction,
            max_blob_gas_per_block=max_blob_gas_per_block,
            blob_schedule_target=blob_schedule_target,
            blob_schedule_max=blob_schedule_max,
        )
        self._cache[cache_key] = clone
        return clone


class T8N(Load):
    """The class that carries out the transition."""

    tracers: Final[GroupTracer | None]

    def __init__(
        self,
        options: Any,
        out_file: TextIO,
        in_file: TextIO,
        cache: ForkCache,
    ) -> None:
        self.out_file = out_file
        self.in_file = in_file
        self.options = options
        forks = Hardfork.discover()

        if "stdin" in (
            options.input_env,
            options.input_alloc,
            options.input_txs,
            options.blob_parameters,
        ):
            stdin = json.load(in_file)
        else:
            stdin = None

        fork_module, self.fork_block = find_fork(forks, self.options, stdin)

        fork_criteria = None
        if self.fork_block is not None and self.fork_block != 0:
            # I can't find where `self.fork_block` is even used, and the vast
            # majority of the time it's zero anyway. Not changing the fork
            # criteria doesn't seem to break the tests, but changing it
            # introduces cloning overhead, so... pretend it didn't happen.
            fork_criteria = ByBlockNumber(self.fork_block)

        target_blobs_per_block = None
        max_blobs_per_block = None
        base_fee_update_fraction = None

        blob_parameters = None
        if options.blob_parameters == "stdin":
            assert stdin is not None
            blob_parameters = stdin["blobParams"]
        elif options.blob_parameters is not None:
            with open(options.blob_parameters, "r") as f:
                blob_parameters = json.load(f)

        if blob_parameters is not None:
            target_blobs_per_block = parse_hex_or_int(
                blob_parameters["target"],
                U64,
            )
            max_blobs_per_block = parse_hex_or_int(
                blob_parameters["max"],
                U64,
            )
            base_fee_update_fraction = parse_hex_or_int(
                blob_parameters["baseFeeUpdateFraction"],
                Uint,
            )

        fork = cache.get(
            fork_module,
            fork_criteria,
            blob_schedule_target=target_blobs_per_block,
            blob_schedule_max=max_blobs_per_block,
            blob_base_fee_update_fraction=base_fee_update_fraction,
        )

        tracers = GroupTracer()

        if self.options.trace:
            trace_memory = getattr(self.options, "trace.memory", False)
            trace_stack = not getattr(self.options, "trace.nostack", False)
            trace_return_data = getattr(self.options, "trace.returndata")
            tracers.add(
                Eip3155Tracer(
                    trace_memory=trace_memory,
                    trace_stack=trace_stack,
                    trace_return_data=trace_return_data,
                    output_basedir=self.options.output_basedir,
                )
            )

        if self.options.opcode_count is not None:
            tracers.add(CountTracer())

        maybe_tracers: GroupTracer | None
        if tracers.tracers:
            trace.set_evm_trace(tracers)
            maybe_tracers = tracers
        else:
            maybe_tracers = None

        self.tracers = maybe_tracers

        self.logger = get_stream_logger("T8N")

        super().__init__(
            self.options.state_fork,
            fork,
        )

        self.chain_id = parse_hex_or_int(self.options.state_chainid, U64)
        self.alloc = Alloc(self, stdin)
        self.env = Env(self, stdin)
        self.txs = Txs(self, stdin)
        self.result = Result(
            self.env.block_difficulty, self.env.base_fee_per_gas
        )

    def _tracer(self, type_: Type[T]) -> T:
        group = self.tracers
        if group is None:
            raise Exception("no tracer configured")
        found = next((x for x in group.tracers if isinstance(x, type_)), None)
        if found is None:
            raise Exception(f"no tracer of type `{type_}` found")
        return found

    def block_environment(self) -> Any:
        """
        Create the environment for the transaction. The keyword
        arguments are adjusted according to the fork.
        """
        kw_arguments = {
            "block_hashes": self.env.block_hashes,
            "coinbase": self.env.coinbase,
            "number": self.env.block_number,
            "time": self.env.block_timestamp,
            "state": self.alloc.state,
            "block_gas_limit": self.env.block_gas_limit,
            "chain_id": self.chain_id,
        }

        block_environment = self.fork.BlockEnvironment

        if self.fork.has_calculate_base_fee_per_gas:
            kw_arguments["base_fee_per_gas"] = self.env.base_fee_per_gas

        if self.fork.hardfork.consensus.is_pos():
            kw_arguments["prev_randao"] = self.env.prev_randao
        else:
            kw_arguments["difficulty"] = self.env.block_difficulty

        if self.fork.has_beacon_roots_address:
            kw_arguments["parent_beacon_block_root"] = (
                self.env.parent_beacon_block_root
            )
            kw_arguments["excess_blob_gas"] = self.env.excess_blob_gas

        return block_environment(**kw_arguments)

    def backup_state(self) -> None:
        """Back up the state in order to restore in case of an error."""
        state = self.alloc.state
        self.alloc.state_backup = (
            self.fork.copy_trie(state._main_trie),
            {
                k: self.fork.copy_trie(t)
                for (k, t) in state._storage_tries.items()
            },
        )

    def restore_state(self) -> None:
        """Restore the state from the backup."""
        state = self.alloc.state
        state._main_trie, state._storage_tries = self.alloc.state_backup

    def pay_block_rewards(self, block_reward: U256, block_env: Any) -> None:
        """Apply the block rewards to the block coinbase."""
        ommer_count = U256(len(self.env.ommers))
        miner_reward = block_reward + (
            ommer_count * (block_reward // U256(32))
        )
        self.fork.create_ether(
            block_env.state, block_env.coinbase, miner_reward
        )

        for ommer in self.env.ommers:
            # Ommer age with respect to the current block.
            ommer_age = U256(block_env.number - ommer.number)
            ommer_miner_reward = (
                (U256(8) - ommer_age) * block_reward
            ) // U256(8)
            self.fork.create_ether(
                block_env.state, ommer.coinbase, ommer_miner_reward
            )

    def run_state_test(self) -> Any:
        """
        Apply a single transaction on pre-state. No system operations
        are performed.
        """
        block_env = self.block_environment()
        block_output = self.fork.BlockOutput()
        self.backup_state()
        if len(self.txs.transactions) > 0:
            tx = self.txs.transactions[0]
            try:
                self.fork.process_transaction(
                    block_env=block_env,
                    block_output=block_output,
                    tx=tx,
                    index=Uint(0),
                )
            except EthereumException as e:
                self.txs.rejected_txs[0] = f"Failed transaction: {e!r}"
                self.restore_state()
                self.logger.warning(f"Transaction {0} failed: {str(e)}")

        self.result.update(self, block_env, block_output)
        self.result.rejected = self.txs.rejected_txs

    def _run_blockchain_test(self, block_env: Any, block_output: Any) -> None:
        if self.fork.has_compute_requests_hash:
            self.fork.process_unchecked_system_transaction(
                block_env=block_env,
                target_address=self.fork.HISTORY_STORAGE_ADDRESS,
                data=block_env.block_hashes[-1],  # The parent hash
            )

        if self.fork.has_beacon_roots_address:
            self.fork.process_unchecked_system_transaction(
                block_env=block_env,
                target_address=self.fork.BEACON_ROOTS_ADDRESS,
                data=block_env.parent_beacon_block_root,
            )

        for i, tx in zip(
            self.txs.successfully_parsed,
            self.txs.transactions,
            strict=True,
        ):
            self.backup_state()
            try:
                self.fork.process_transaction(
                    block_env, block_output, tx, Uint(i)
                )
            except EthereumException as e:
                self.txs.rejected_txs[i] = f"Failed transaction: {e!r}"
                self.restore_state()
                self.logger.warning(f"Transaction {i} failed: {e!r}")

        if not self.fork.proof_of_stake:
            if self.options.state_reward is None:
                self.pay_block_rewards(self.fork.BLOCK_REWARD, block_env)
            elif self.options.state_reward != -1:
                self.pay_block_rewards(
                    U256(self.options.state_reward), block_env
                )

        if self.fork.has_withdrawal:
            self.fork.process_withdrawals(
                block_env, block_output, self.env.withdrawals
            )

        if self.fork.has_compute_requests_hash:
            self.fork.process_general_purpose_requests(block_env, block_output)

    def run_blockchain_test(self) -> None:
        """
        Apply a block on the pre-state. Also includes system operations.
        """
        block_env = self.block_environment()
        block_output = self.fork.BlockOutput()

        try:
            self._run_blockchain_test(block_env, block_output)
        except InvalidBlock as e:
            self.result.block_exception = f"{e}"

        self.result.update(self, block_env, block_output)
        self.result.rejected = self.txs.rejected_txs

    def run(self) -> int:
        """Run the transition and provide the relevant outputs."""
        # Clear files that may have been created in a previous
        # run of the t8n tool.
        # Define the specific files and pattern to delete
        files_to_delete = [
            self.options.output_result,
            self.options.output_alloc,
            self.options.output_body,
        ]
        pattern_to_delete = "trace-*.jsonl"

        # Iterate through the directory
        for file in os.listdir(self.options.output_basedir):
            file_path = os.path.join(self.options.output_basedir, file)

            # Check if the file matches the specific names or the pattern
            if file in files_to_delete or fnmatch.fnmatch(
                file, pattern_to_delete
            ):
                os.remove(file_path)

        try:
            if self.options.state_test:
                self.run_state_test()
            else:
                self.run_blockchain_test()
        except FatalError as e:
            self.logger.error(str(e))
            return 1

        json_state = self.alloc.to_json()
        json_result = self.result.to_json()

        json_output: dict[str, object] = {}

        if self.options.output_body == "stdout":
            txs_rlp = "0x" + rlp.encode(self.txs.all_txs).hex()
            json_output["body"] = txs_rlp
        elif self.options.output_body is not None:
            txs_rlp_path = os.path.join(
                self.options.output_basedir,
                self.options.output_body,
            )
            txs_rlp = "0x" + rlp.encode(self.txs.all_txs).hex()
            with open(txs_rlp_path, "w") as f:
                json.dump(txs_rlp, f)
            self.logger.info(f"Wrote transaction rlp to {txs_rlp_path}")

        if self.options.output_alloc == "stdout":
            json_output["alloc"] = json_state
        else:
            alloc_output_path = os.path.join(
                self.options.output_basedir,
                self.options.output_alloc,
            )
            with open(alloc_output_path, "w") as f:
                json.dump(json_state, f, indent=4)
            self.logger.info(f"Wrote alloc to {alloc_output_path}")

        if self.options.output_result == "stdout":
            json_output["result"] = json_result
        else:
            result_output_path = os.path.join(
                self.options.output_basedir,
                self.options.output_result,
            )
            with open(result_output_path, "w") as f:
                json.dump(json_result, f, indent=4)
            self.logger.info(f"Wrote result to {result_output_path}")

        if self.options.opcode_count == "stdout":
            opcode_count_results = self._tracer(CountTracer).results()
            json_output["opcodeCount"] = opcode_count_results
        elif self.options.opcode_count is not None:
            opcode_count_results = self._tracer(CountTracer).results()
            result_output_path = os.path.join(
                self.options.output_basedir,
                self.options.opcode_count,
            )
            with open(result_output_path, "w") as f:
                json.dump(opcode_count_results, f, indent=4)
            self.logger.info(f"Wrote opcode counts to {result_output_path}")

        if json_output:
            json.dump(json_output, self.out_file, indent=4)

        return 0
