pub use OmniBridge::Event as OmniEvents;
use starknet::{ClassHash, ContractAddress};
pub use crate::bridge_types::{
    DeployToken, FinTransfer, InitTransfer, LogMetadata, MetadataPayload, Signature,
    TransferMessagePayload,
};

#[starknet::interface]
pub trait IOmniBridge<TContractState> {
    fn log_metadata(ref self: TContractState, token: ContractAddress);
    fn deploy_token(ref self: TContractState, signature: Signature, payload: MetadataPayload);
    fn fin_transfer(
        ref self: TContractState, signature: Signature, payload: TransferMessagePayload,
    );
    fn init_transfer(
        ref self: TContractState,
        token_address: ContractAddress,
        amount: u128,
        fee: u128,
        native_fee: u128,
        recipient: ByteArray,
        message: ByteArray,
    );
    fn upgrade_token(
        ref self: TContractState, token_address: ContractAddress, new_class_hash: ClassHash,
    );
    fn set_pause_flags(ref self: TContractState, flags: u8);
    fn pause_all(ref self: TContractState);
    fn get_token_address(self: @TContractState, token_id: ByteArray) -> ContractAddress;
    fn is_bridge_token(self: @TContractState, token_address: ContractAddress) -> bool;
    fn is_transfer_finalised(self: @TContractState, nonce: u64) -> bool;
}

#[starknet::contract]
mod OmniBridge {
    use core::keccak::compute_keccak_byte_array;
    use core::num::traits::Zero;
    use openzeppelin::access::accesscontrol::AccessControlComponent;
    use openzeppelin::introspection::src5::SRC5Component;
    use openzeppelin::token::erc20::interface::{IERC20Dispatcher, IERC20DispatcherTrait};
    use openzeppelin::upgrades::interface::{
        IUpgradeable, IUpgradeableDispatcher, IUpgradeableDispatcherTrait,
    };
    use openzeppelin::upgrades::upgradeable::UpgradeableComponent;
    use starknet::eth_signature::verify_eth_signature;
    use starknet::event::EventEmitter;
    use starknet::secp256_trait::signature_from_vrs;
    use starknet::storage::{
        Map, StorageMapReadAccess, StorageMapWriteAccess, StoragePointerReadAccess,
        StoragePointerWriteAccess,
    };
    use starknet::syscalls::deploy_syscall;
    use starknet::{
        ClassHash, ContractAddress, EthAddress, SyscallResultTrait, get_caller_address,
        get_contract_address, syscalls,
    };
    use crate::bridge_types::{
        DeployToken, FinTransfer, InitTransfer, LogMetadata, MetadataPayload, MetadataPayloadTrait,
        PauseStateChanged, Signature, TransferMessagePayload, TransferMessagePayloadTrait,
    };
    use crate::utils;
    use crate::utils::reverse_u256_bytes;

    // Role constants
    const DEFAULT_ADMIN_ROLE: felt252 = 0;
    const PAUSER_ROLE: felt252 = selector!("PAUSER_ROLE");

    // Pause flag constants
    const PAUSE_INIT_TRANSFER: u8 = 0x01; // 0001
    const PAUSE_FIN_TRANSFER: u8 = 0x02; // 0010
    const PAUSE_DEPLOY_TOKEN: u8 = 0x04; // 0100
    const PAUSE_ALL: u8 = 0xFF; // 1111

    component!(path: AccessControlComponent, storage: accesscontrol, event: AccessControlEvent);
    component!(path: SRC5Component, storage: src5, event: SRC5Event);
    component!(path: UpgradeableComponent, storage: upgradeable, event: UpgradeableEvent);

    #[abi(embed_v0)]
    impl AccessControlMixinImpl =
        AccessControlComponent::AccessControlMixinImpl<ContractState>;
    impl AccessControlInternalImpl = AccessControlComponent::InternalImpl<ContractState>;
    impl UpgradeableInternalImpl = UpgradeableComponent::InternalImpl<ContractState>;

    #[event]
    #[derive(Drop, starknet::Event)]
    pub enum Event {
        LogMetadata: LogMetadata,
        DeployToken: DeployToken,
        InitTransfer: InitTransfer,
        FinTransfer: FinTransfer,
        PauseStateChanged: PauseStateChanged,
        #[flat]
        AccessControlEvent: AccessControlComponent::Event,
        #[flat]
        SRC5Event: SRC5Component::Event,
        #[flat]
        UpgradeableEvent: UpgradeableComponent::Event,
    }

    // Used nonces
    #[storage]
    struct Storage {
        #[substorage(v0)]
        accesscontrol: AccessControlComponent::Storage,
        #[substorage(v0)]
        src5: SRC5Component::Storage,
        #[substorage(v0)]
        upgradeable: UpgradeableComponent::Storage,
        pause_flags: u8,
        bridge_token_class_hash: ClassHash,
        current_origin_nonce: u64,
        // Bitmap: slot = nonce / 251, bit = nonce % 251
        completed_transfers: Map<u64, felt252>,
        starknet_to_near_token: Map<ContractAddress, ByteArray>,
        // Can't use ByteArray as a key. Using hash instead
        near_to_starknet_token: Map<u256, ContractAddress>,
        omni_bridge_chain_id: u8,
        omni_bridge_derived_address: EthAddress,
        strk_token_address: ContractAddress,
    }

    #[constructor]
    fn constructor(
        ref self: ContractState,
        omni_bridge_derived_address: EthAddress,
        omni_bridge_chain_id: u8,
        token_class_hash: ClassHash,
        default_admin: ContractAddress,
        strk_token_address: ContractAddress,
    ) {
        self.omni_bridge_derived_address.write(omni_bridge_derived_address);
        self.omni_bridge_chain_id.write(omni_bridge_chain_id);
        self.bridge_token_class_hash.write(token_class_hash);
        self.strk_token_address.write(strk_token_address);
        self.pause_flags.write(0);

        self.accesscontrol.initializer();
        self.accesscontrol._grant_role(DEFAULT_ADMIN_ROLE, default_admin);
    }

    #[abi(embed_v0)]
    #[feature("safe_dispatcher")]
    impl OmniBridgeImpl of super::IOmniBridge<ContractState> {
        fn log_metadata(ref self: ContractState, token: ContractAddress) {
            // There are two possible metadata standards in use.
            // 1. Old style: name and symbol are felt252 values.
            // 2. New style: name and symbol are ByteArray values (ERC20 ABI).
            // We are using low-level contract calls to determine the type.

            let call_data: Array<felt252> = array![];
            let mut res = syscalls::call_contract_syscall(
                token, selector!("name"), call_data.span(),
            )
                .unwrap_syscall();

            let name = if res.len() == 1 {
                // Old standard (felt252)
                let name = OptionTrait::expect(
                    Serde::<felt252>::deserialize(ref res), 'Could not deserialize name',
                );
                utils::felt252_to_string(name)
            } else {
                // New standard (ByteArray)
                OptionTrait::expect(
                    Serde::<ByteArray>::deserialize(ref res), 'Could not deserialize name',
                )
            };

            let mut res = syscalls::call_contract_syscall(
                token, selector!("symbol"), call_data.span(),
            )
                .unwrap_syscall();

            let symbol = if res.len() == 1 {
                // Old standard (felt252)
                let symbol = OptionTrait::expect(
                    Serde::<felt252>::deserialize(ref res), 'Could not deserialize symbol',
                );
                utils::felt252_to_string(symbol)
            } else {
                // New standard (ByteArray)
                OptionTrait::expect(
                    Serde::<ByteArray>::deserialize(ref res), 'Could not deserialize symbol',
                )
            };

            let decimals = {
                let mut res = syscalls::call_contract_syscall(
                    token, selector!("decimals"), call_data.span(),
                )
                    .unwrap_syscall();

                let decimals = OptionTrait::expect(
                    Serde::<u8>::deserialize(ref res), 'Could not deserialize decimals',
                );
                decimals
            };

            self.emit(Event::LogMetadata(LogMetadata { address: token, name, symbol, decimals }))
        }

        fn deploy_token(ref self: ContractState, signature: Signature, payload: MetadataPayload) {
            assert(!_is_paused(@self, PAUSE_DEPLOY_TOKEN), 'ERR_DEPLOY_TOKEN_PAUSED');

            _verify_borsh_signature(ref self, @payload.to_borsh(), signature);

            let token_id_hash = compute_keccak_byte_array(@payload.token);
            let existing_token = self.near_to_starknet_token.read(token_id_hash);
            assert(existing_token.is_zero(), 'ERR_TOKEN_ALREADY_DEPLOYED');

            let decimals = _normalizeDecimals(payload.decimals);

            let mut constructor_calldata: Array<felt252> = array![];
            (payload.name.clone(), payload.symbol.clone(), decimals)
                .serialize(ref constructor_calldata);

            // Use the low part of the u256 hash to ensure it fits in felt252
            let salt: felt252 = token_id_hash.low.into();
            let (contract_address, _) = deploy_syscall(
                self.bridge_token_class_hash.read(), salt, constructor_calldata.span(), false,
            )
                .unwrap_syscall();

            self.starknet_to_near_token.write(contract_address, payload.token.clone());
            self.near_to_starknet_token.write(token_id_hash, contract_address);

            self
                .emit(
                    Event::DeployToken(
                        DeployToken {
                            token_address: contract_address,
                            near_token_id: payload.token,
                            name: payload.name,
                            symbol: payload.symbol,
                            decimals,
                            origin_decimals: payload.decimals,
                        },
                    ),
                )
        }

        fn fin_transfer(
            ref self: ContractState, signature: Signature, payload: TransferMessagePayload,
        ) {
            assert(!_is_paused(@self, PAUSE_FIN_TRANSFER), 'ERR_FIN_TRANSFER_PAUSED');

            assert(
                !self.is_transfer_finalised(payload.destination_nonce), 'ERR_NONCE_ALREADY_USED',
            );
            _set_transfer_finalised(ref self, payload.destination_nonce);

            _verify_borsh_signature(
                ref self, @payload.to_borsh(self.omni_bridge_chain_id.read()), signature,
            );

            if self.is_bridge_token(payload.token_address) {
                IBridgeTokenDispatcher { contract_address: payload.token_address }
                    .mint(payload.recipient, payload.amount.into());
            } else {
                let success = IERC20Dispatcher { contract_address: payload.token_address }
                    .transfer(payload.recipient, payload.amount.into());
                assert(success, 'ERR_TRANSFER_FAILED');
            }

            self
                .emit(
                    Event::FinTransfer(
                        FinTransfer {
                            origin_chain: payload.origin_chain,
                            origin_nonce: payload.origin_nonce,
                            token_address: payload.token_address,
                            amount: payload.amount,
                            recipient: payload.recipient,
                            fee_recipient: payload.fee_recipient,
                            message: payload.message,
                        },
                    ),
                )
        }

        fn init_transfer(
            ref self: ContractState,
            token_address: ContractAddress,
            amount: u128,
            fee: u128,
            native_fee: u128,
            recipient: ByteArray,
            message: ByteArray,
        ) {
            assert(!_is_paused(@self, PAUSE_INIT_TRANSFER), 'ERR_INIT_TRANSFER_PAUSED');

            assert(amount > 0, 'ERR_ZERO_AMOUNT');
            assert(fee < amount, 'ERR_INVALID_FEE');

            let origin_nonce = self.current_origin_nonce.read() + 1;
            self.current_origin_nonce.write(origin_nonce);

            let caller = get_caller_address();

            if self.is_bridge_token(token_address) {
                IBridgeTokenDispatcher { contract_address: token_address }
                    .burn(caller, amount.into());
            } else {
                let success = IERC20Dispatcher { contract_address: token_address }
                    .transfer_from(caller, get_contract_address(), amount.into());
                assert(success, 'ERR_TRANSFER_FROM_FAILED');
            }

            if native_fee > 0 {
                let native_token = self.strk_token_address.read();
                let success = IERC20Dispatcher { contract_address: native_token }
                    .transfer_from(caller, get_contract_address(), native_fee.into());
                assert(success, 'ERR_FEE_TRANSFER_FAILED');
            }

            self
                .emit(
                    Event::InitTransfer(
                        InitTransfer {
                            sender: caller,
                            token_address,
                            origin_nonce,
                            amount,
                            fee,
                            native_fee,
                            recipient,
                            message,
                        },
                    ),
                )
        }

        fn upgrade_token(
            ref self: ContractState, token_address: ContractAddress, new_class_hash: ClassHash,
        ) {
            self.accesscontrol.assert_only_role(DEFAULT_ADMIN_ROLE);
            assert(self.is_bridge_token(token_address), 'ERR_NOT_BRIDGE_TOKEN');

            let upgradeable = IUpgradeableDispatcher { contract_address: token_address };
            upgradeable.upgrade(new_class_hash);
        }

        fn set_pause_flags(ref self: ContractState, flags: u8) {
            self.accesscontrol.assert_only_role(DEFAULT_ADMIN_ROLE);
            let old_flags = self.pause_flags.read();
            self.pause_flags.write(flags);

            self
                .emit(
                    Event::PauseStateChanged(
                        PauseStateChanged {
                            old_flags, new_flags: flags, admin: get_caller_address(),
                        },
                    ),
                );
        }

        fn pause_all(ref self: ContractState) {
            self.accesscontrol.assert_only_role(PAUSER_ROLE);
            let old_flags = self.pause_flags.read();
            self.pause_flags.write(PAUSE_ALL);

            self
                .emit(
                    Event::PauseStateChanged(
                        PauseStateChanged {
                            old_flags, new_flags: PAUSE_ALL, admin: get_caller_address(),
                        },
                    ),
                );
        }

        fn get_token_address(self: @ContractState, token_id: ByteArray) -> ContractAddress {
            let token_id_hash = compute_keccak_byte_array(@token_id);
            self.near_to_starknet_token.read(token_id_hash)
        }

        fn is_bridge_token(self: @ContractState, token_address: ContractAddress) -> bool {
            self.starknet_to_near_token.read(token_address).len() > 0
        }

        fn is_transfer_finalised(self: @ContractState, nonce: u64) -> bool {
            let (slot, bit) = _nonce_slot_and_bit(nonce);
            let bitmap: u256 = self.completed_transfers.read(slot).into();
            bitmap & bit != 0
        }
    }

    #[abi(embed_v0)]
    impl UpgradeableImpl of IUpgradeable<ContractState> {
        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
            self.accesscontrol.assert_only_role(DEFAULT_ADMIN_ROLE);
            self.upgradeable.upgrade(new_class_hash);
        }
    }

    // Helper functions
    fn _verify_borsh_signature(
        ref self: ContractState, borsh_bytes: @ByteArray, signature: Signature,
    ) {
        let message_hash_le = compute_keccak_byte_array(borsh_bytes);
        let message_hash = reverse_u256_bytes(message_hash_le);

        let sig = signature_from_vrs(signature.v, signature.r, signature.s);
        verify_eth_signature(message_hash, sig, self.omni_bridge_derived_address.read());
    }

    fn _is_paused(self: @ContractState, flag: u8) -> bool {
        let flags = self.pause_flags.read();
        (flags & flag) != 0
    }

    #[starknet::interface]
    trait IBridgeToken<TContractState> {
        fn mint(ref self: TContractState, recipient: ContractAddress, amount: u256);
        fn burn(ref self: TContractState, account: ContractAddress, amount: u256);
    }

    fn _nonce_slot_and_bit(nonce: u64) -> (u64, u256) {
        let slot = nonce / 251;
        let bit: u256 = _pow2((nonce % 251).into());
        (slot, bit)
    }

    fn _set_transfer_finalised(ref self: ContractState, nonce: u64) {
        let (slot, bit) = _nonce_slot_and_bit(nonce);
        let bitmap: u256 = self.completed_transfers.read(slot).into();
        self.completed_transfers.write(slot, (bitmap | bit).try_into().unwrap());
    }

    fn _pow2(mut exp: u128) -> u256 {
        let mut result: u256 = 1;
        while exp > 0 {
            result *= 2;
            exp -= 1;
        }
        result
    }

    fn _normalizeDecimals(decimals: u8) -> u8 {
        let maxAllowedDecimals: u8 = 18;
        if (decimals > maxAllowedDecimals) {
            return maxAllowedDecimals;
        }
        return decimals;
    }

    #[cfg(test)]
    mod tests {
        use super::{OmniBridgeImpl, _set_transfer_finalised};

        fn setup() -> super::ContractState {
            super::contract_state_for_testing()
        }

        #[test]
        fn test_set_transfer_finalised_and_check_range() {
            let mut state = setup();
            let max_nonce: u64 = 0xFFFFFFFFFFFFFFFF;
            let nonces = [
                0, 1, 40, 42, 43, 44, 45, 125, 250, 251, 255, 256, 257, 502, 1000, max_nonce - 2,
                max_nonce - 1, max_nonce,
            ]
                .span();

            // Verify all nonces are initially unset
            let mut i: usize = 0;
            while i < nonces.len() {
                assert!(!OmniBridgeImpl::is_transfer_finalised(@state, *nonces[i]));
                i += 1;
            }

            // Set all nonces
            let mut i: usize = 0;
            while i < nonces.len() {
                _set_transfer_finalised(ref state, *nonces[i]);
                i += 1;
            }

            // Verify all nonces are now set
            let mut i: usize = 0;
            while i < nonces.len() {
                assert!(OmniBridgeImpl::is_transfer_finalised(@state, *nonces[i]));
                i += 1;
            }

            // Verify unset nonces remain unset
            let unset = [2, 41, 124, 249, 252, 254, 999, max_nonce - 3].span();
            let mut i: usize = 0;
            while i < unset.len() {
                assert!(!OmniBridgeImpl::is_transfer_finalised(@state, *unset[i]));
                i += 1;
            };
        }

        #[test]
        fn test_set_transfer_finalised_idempotent() {
            let mut state = setup();
            _set_transfer_finalised(ref state, 42);
            _set_transfer_finalised(ref state, 42);
            assert!(OmniBridgeImpl::is_transfer_finalised(@state, 42));
        }
    }
}