# -*- coding: utf-8 -*-
# ------------------------------------------------------------------------------
#
#   Copyright 2018-2021 Fetch.AI Limited
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# ------------------------------------------------------------------------------

"""Client functionality."""

import json
import math
import time
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Tuple

import certifi
import grpc
from dateutil.parser import isoparse
from google.protobuf.timestamp_pb2 import Timestamp

from web3 import Web3
from kiipy.aerial.client.bank import create_bank_send_msg
from kiipy.aerial.client.distribution import create_withdraw_delegator_reward
from kiipy.aerial.client.staking import (
    ValidatorStatus,
    create_delegate_msg,
    create_redelegate_msg,
    create_undelegate_msg,
)
from kiipy.aerial.client.utils import (
    ensure_timedelta,
    get_paginated,
    prepare_and_broadcast_basic_transaction,
)
from kiipy.aerial.config import NetworkConfig
from kiipy.aerial.exceptions import NotFoundError, QueryTimeoutError
from kiipy.aerial.gas import GasStrategy, SimulationGasStrategy
from kiipy.aerial.tx import Transaction, TxState
from kiipy.aerial.tx_helpers import MessageLog, SubmittedTx, TxResponse
from kiipy.aerial.urls import Protocol, parse_url
from kiipy.aerial.wallet import Wallet
from kiipy.auth.rest_client import AuthRestClient
from kiipy.bank.rest_client import BankRestClient
from kiipy.common.rest_client import RestClient
from kiipy.cosmwasm.rest_client import CosmWasmRestClient
from kiipy.crypto.address import Address
from kiipy.crypto.hashfuncs import sha256
from kiipy.distribution.rest_client import DistributionRestClient
from kiipy.params.rest_client import ParamsRestClient
from kiipy.protos.cosmos.auth.v1beta1.auth_pb2 import BaseAccount
from kiipy.protos.cosmos.auth.v1beta1.query_pb2 import QueryAccountRequest
from kiipy.protos.cosmos.auth.v1beta1.query_pb2_grpc import QueryStub as AuthGrpcClient
from kiipy.protos.cosmos.bank.v1beta1.query_pb2 import (
    QueryAllBalancesRequest,
    QueryBalanceRequest,
)
from kiipy.protos.cosmos.bank.v1beta1.query_pb2_grpc import QueryStub as BankGrpcClient
from kiipy.protos.cosmos.base.tendermint.v1beta1.query_pb2 import (
    GetBlockByHeightRequest,
    GetLatestBlockRequest,
)
from kiipy.protos.cosmos.base.tendermint.v1beta1.query_pb2_grpc import (
    ServiceStub as TendermintQueryGrpcClient,
)
from kiipy.protos.cosmos.crypto.ed25519.keys_pb2 import (  # noqa # pylint: disable=unused-import
    PubKey,
)
from kiipy.protos.cosmos.distribution.v1beta1.query_pb2 import (
    QueryDelegationRewardsRequest,
)
from kiipy.protos.cosmos.distribution.v1beta1.query_pb2_grpc import (
    QueryStub as DistributionGrpcClient,
)
from kiipy.protos.cosmos.params.v1beta1.query_pb2 import QueryParamsRequest
from kiipy.protos.cosmos.params.v1beta1.query_pb2_grpc import (
    QueryStub as QueryParamsGrpcClient,
)
from kiipy.protos.cosmos.staking.v1beta1.query_pb2 import (
    QueryDelegatorDelegationsRequest,
    QueryDelegatorUnbondingDelegationsRequest,
    QueryValidatorsRequest,
)
from kiipy.protos.cosmos.staking.v1beta1.query_pb2_grpc import (
    QueryStub as StakingGrpcClient,
)
from kiipy.protos.cosmos.tx.v1beta1.service_pb2 import (
    BroadcastMode,
    BroadcastTxRequest,
    GetTxRequest,
    SimulateRequest,
)
from kiipy.protos.cosmos.tx.v1beta1.service_pb2_grpc import ServiceStub as TxGrpcClient
from kiipy.protos.cosmwasm.wasm.v1.query_pb2_grpc import QueryStub as CosmWasmGrpcClient
from kiipy.staking.rest_client import StakingRestClient
from kiipy.tendermint.rest_client import (
    CosmosBaseTendermintRestClient as TendermintRestClient,
)
from kiipy.tx.rest_client import TxRestClient


DEFAULT_BANK_CONTRACT = "0x4381dC2aB14285160c808659aEe005D51255adD7"
DEFAULT_QUERY_TIMEOUT_SECS = 15
DEFAULT_QUERY_INTERVAL_SECS = 2
DEFAULT_TX_GAS_LIMIT = 2000000
COSMOS_SDK_DEC_COIN_PRECISION = (
    10**18
)  # TODO: Revisit this, based on discussion with Matt, this should be 10^6


@dataclass
class Account:
    """Account."""

    address: Address
    number: int
    sequence: int


@dataclass
class StakingPosition:
    """Staking positions."""

    validator: Address
    amount: int
    reward: int


@dataclass
class UnbondingPositions:
    """Unbonding positions."""

    validator: Address
    amount: int


@dataclass
class Validator:
    """Validator."""

    address: Address  # the operators address
    tokens: int  # The total amount of tokens for the validator
    moniker: str
    status: ValidatorStatus


@dataclass
class Coin:
    """Coins."""

    amount: int
    denom: str


@dataclass
class StakingSummary:
    """Get the staking summary."""

    current_positions: List[StakingPosition]
    unbonding_positions: List[UnbondingPositions]

    @property
    def total_staked(self) -> int:
        """Get the total staked amount."""
        return sum(map(lambda p: p.amount, self.current_positions))

    @property
    def total_rewards(self) -> int:
        """Get the total rewards."""
        return sum(map(lambda p: p.reward, self.current_positions))

    @property
    def total_unbonding(self) -> int:
        """total unbonding."""
        return sum(map(lambda p: p.amount, self.unbonding_positions))


@dataclass
class Block:
    """Block."""

    height: int
    time: datetime
    chain_id: str
    tx_hashes: List[str]

    @staticmethod
    def from_proto(block: Any) -> "Block":
        """Parse the block.

        :param block: block as Any
        :return: parsed block as Block
        """
        return Block(
            height=int(block.header.height),
            time=Block._parse_timestamp(block.header.time),
            tx_hashes=[sha256(tx).hex().upper() for tx in block.data.txs],
            chain_id=block.header.chain_id,
        )

    @staticmethod
    def _parse_timestamp(timestamp: Timestamp):
        """Parse the timestamp.

        :param timestamp: timestamp
        :return: parsed timestamp
        """
        return datetime.fromtimestamp(timestamp.seconds, tz=timezone.utc) + timedelta(
            microseconds=timestamp.nanos // 1000
        )


class LedgerClient:
    """Ledger client."""

    def __init__(
        self,
        cfg: NetworkConfig,
        query_interval_secs: int = DEFAULT_QUERY_INTERVAL_SECS,
        query_timeout_secs: int = DEFAULT_QUERY_TIMEOUT_SECS,
    ):
        """Init ledger client.

        :param cfg: Network configurations
        :param query_interval_secs: int. optional interval int seconds
        :param query_timeout_secs: int. optional interval int seconds
        """
        self._query_interval_secs = query_interval_secs
        self._query_timeout_secs = query_timeout_secs
        cfg.validate()
        self._network_config = cfg
        self._gas_strategy: GasStrategy = SimulationGasStrategy(self)

        parsed_url = parse_url(cfg.url)

        self.w3 = Web3(Web3.HTTPProvider(parsed_url.rest_url))
        
        # if parsed_url.protocol == Protocol.GRPC:
        #     if parsed_url.secure:
        #         with open(certifi.where(), "rb") as f:
        #             trusted_certs = f.read()
        #         credentials = grpc.ssl_channel_credentials(
        #             root_certificates=trusted_certs
        #         )
        #         grpc_client = grpc.secure_channel(parsed_url.host_and_port, credentials)
        #     else:
        #         grpc_client = grpc.insecure_channel(parsed_url.host_and_port)

        #     self.wasm = CosmWasmGrpcClient(grpc_client)
        #     self.auth = AuthGrpcClient(grpc_client)
        #     self.txs = TxGrpcClient(grpc_client)
        #     self.bank = BankGrpcClient(grpc_client)
        #     self.staking = StakingGrpcClient(grpc_client)
        #     self.distribution = DistributionGrpcClient(grpc_client)
        #     self.params = QueryParamsGrpcClient(grpc_client)
        #     self.tendermint = TendermintQueryGrpcClient(grpc_client)
        # else:
        #     rest_client = RestClient(parsed_url.rest_url)

        #     self.wasm = CosmWasmRestClient(rest_client)  # type: ignore
        #     self.auth = AuthRestClient(rest_client)  # type: ignore
        #     self.txs = TxRestClient(rest_client)  # type: ignore
        #     self.bank = BankRestClient(rest_client)  # type: ignore
        #     self.staking = StakingRestClient(rest_client)  # type: ignore
        #     self.distribution = DistributionRestClient(rest_client)  # type: ignore
        #     self.params = ParamsRestClient(rest_client)  # type: ignore
        #     self.tendermint = TendermintRestClient(rest_client)  # type: ignore

    @property
    def network_config(self) -> NetworkConfig:
        """Get the network config.

        :return: network config
        """
        return self._network_config

    @property
    def gas_strategy(self) -> GasStrategy:
        """Get gas strategy.

        :return: gas strategy
        """
        return self._gas_strategy

    @gas_strategy.setter
    def gas_strategy(self, strategy: GasStrategy):
        """Set gas strategy.

        :param strategy: strategy
        :raises RuntimeError: Invalid strategy must implement GasStrategy interface
        """
        if not isinstance(strategy, GasStrategy):
            raise RuntimeError("Invalid strategy must implement GasStrategy interface")
        self._gas_strategy = strategy

    def get_balance(self, address: Address) -> Any:
        """Query wallet balance

        :param address: address
        :return: wallet balance
        """

        balance = 0
        if self.w3.is_connected():
            balance = self.w3.eth.get_balance(address)
        return balance
        
    def query_account(self, address: Address) -> Account:
        """Query account.

        :param address: address
        :raises RuntimeError: Unexpected account type returned from query
        :return: account details
        """
        request = QueryAccountRequest(address=str(address))
        response = self.auth.Account(request)

        account = BaseAccount()
        if not response.account.Is(BaseAccount.DESCRIPTOR):
            raise RuntimeError("Unexpected account type returned from query")
        response.account.Unpack(account)

        return Account(
            address=address,
            number=account.account_number,
            sequence=account.sequence,
        )

    def query_params(self, subspace: str, key: str) -> Any:
        """Query Prams.

        :param subspace: subspace
        :param key: key
        :return: Query params
        """
        req = QueryParamsRequest(subspace=subspace, key=key)
        resp = self.params.Params(req)
        return json.loads(resp.param.value)

    def query_bank_balance(self, address: Address) -> int:
        """Query wallet balance.

        :param address: address
        :return: wallet balance
        """
        denom = denom or self.network_config.fee_denomination

        req = QueryBalanceRequest(
            address=str(address),
            denom=denom,
        )

        resp = self.bank.Balance(req)
        assert resp.balance.denom == denom  # sanity check

        return int(float(resp.balance.amount))

    def send_tokens(
        self,
        destination: Address,
        amount: int,
        denom: str,
        sender: Wallet,
        memo: Optional[str] = None,
        gas_limit: Optional[int] = DEFAULT_TX_GAS_LIMIT,
    ) -> SubmittedTx:
        """Send tokens.

        :param destination: destination address
        :param amount: amount
        :param denom: denom
        :param sender: sender
        :param memo: memo, defaults to None
        :param gas_limit: gas limit, defaults to None
        :return: prepare and broadcast the transaction and transaction details
        """
        # build up the store transaction
        tx = Transaction()
        tx.add_message(
            create_bank_send_msg(sender.address(), destination, amount, denom)
        )

        return prepare_and_broadcast_basic_transaction(
            self, tx, sender, gas_limit=gas_limit, memo=memo
        )

    def query_validators(
        self, status: Optional[ValidatorStatus] = None
    ) -> List[Validator]:
        """Query validators.

        :param status: validator status, defaults to None
        :return: List of validators
        """
        filtered_status = status or ValidatorStatus.BONDED

        req = QueryValidatorsRequest()
        if filtered_status != ValidatorStatus.UNSPECIFIED:
            req.status = filtered_status.value

        resp = self.staking.Validators(req)

        validators: List[Validator] = []
        for validator in resp.validators:
            validators.append(
                Validator(
                    address=Address(validator.operator_address),
                    tokens=int(float(validator.tokens)),
                    moniker=str(validator.description.moniker),
                    status=ValidatorStatus.from_proto(validator.status),
                )
            )
        return validators

    def query_staking_summary(self, address: Address) -> StakingSummary:
        """Query staking summary.

        :param address: address
        :return: staking summary
        """
        current_positions: List[StakingPosition] = []

        req = QueryDelegatorDelegationsRequest(delegator_addr=str(address))

        for resp in get_paginated(
            req, self.staking.DelegatorDelegations, per_page_limit=1
        ):
            for item in resp.delegation_responses:
                req = QueryDelegationRewardsRequest(
                    delegator_address=str(address),
                    validator_address=str(item.delegation.validator_address),
                )
                rewards_resp = self.distribution.DelegationRewards(req)

                stake_reward = 0
                for reward in rewards_resp.rewards:
                    if reward.denom == self.network_config.staking_denomination:
                        stake_reward = (
                            int(float(reward.amount)) // COSMOS_SDK_DEC_COIN_PRECISION
                        )
                        break

                current_positions.append(
                    StakingPosition(
                        validator=Address(item.delegation.validator_address),
                        amount=int(float(item.balance.amount)),
                        reward=stake_reward,
                    )
                )

        unbonding_summary: Dict[str, int] = {}
        req = QueryDelegatorUnbondingDelegationsRequest(delegator_addr=str(address))

        for resp in get_paginated(req, self.staking.DelegatorUnbondingDelegations):
            for item in resp.unbonding_responses:
                validator = str(item.validator_address)
                total_unbonding = unbonding_summary.get(validator, 0)

                for entry in item.entries:
                    total_unbonding += int(float(entry.balance))

                unbonding_summary[validator] = total_unbonding

        # build the final list of unbonding positions
        unbonding_positions: List[UnbondingPositions] = []
        for validator, total_unbonding in unbonding_summary.items():
            unbonding_positions.append(
                UnbondingPositions(
                    validator=Address(validator),
                    amount=total_unbonding,
                )
            )

        return StakingSummary(
            current_positions=current_positions, unbonding_positions=unbonding_positions
        )

    def delegate_tokens(
        self,
        validator: Address,
        amount: int,
        sender: Wallet,
        memo: Optional[str] = None,
        gas_limit: Optional[int] = DEFAULT_TX_GAS_LIMIT,
    ) -> SubmittedTx:
        """Delegate tokens.

        :param validator: validator address
        :param amount: amount
        :param sender: sender
        :param memo: memo, defaults to None
        :param gas_limit: gas limit, defaults to None
        :return: prepare and broadcast the transaction and transaction details
        """
        tx = Transaction()
        tx.add_message(
            create_delegate_msg(
                sender.address(),
                validator,
                amount,
                self.network_config.staking_denomination,
            )
        )

        return prepare_and_broadcast_basic_transaction(
            self, tx, sender, gas_limit=gas_limit, memo=memo
        )

    def redelegate_tokens(
        self,
        current_validator: Address,
        next_validator: Address,
        amount: int,
        sender: Wallet,
        memo: Optional[str] = None,
        gas_limit: Optional[int] = DEFAULT_TX_GAS_LIMIT,
    ) -> SubmittedTx:
        """Redelegate tokens.

        :param current_validator: current validator address
        :param next_validator: next validator address
        :param amount: amount
        :param sender: sender
        :param memo: memo, defaults to None
        :param gas_limit: gas limit, defaults to None
        :return: prepare and broadcast the transaction and transaction details
        """
        tx = Transaction()
        tx.add_message(
            create_redelegate_msg(
                sender.address(),
                current_validator,
                next_validator,
                amount,
                self.network_config.staking_denomination,
            )
        )

        return prepare_and_broadcast_basic_transaction(
            self, tx, sender, gas_limit=gas_limit, memo=memo
        )

    def undelegate_tokens(
        self,
        validator: Address,
        amount: int,
        sender: Wallet,
        memo: Optional[str] = None,
        gas_limit: Optional[int] = DEFAULT_TX_GAS_LIMIT,
    ) -> SubmittedTx:
        """Undelegate tokens.

        :param validator: validator
        :param amount: amount
        :param sender: sender
        :param memo: memo, defaults to None
        :param gas_limit: gas limit, defaults to None
        :return: prepare and broadcast the transaction and transaction details
        """
        tx = Transaction()
        tx.add_message(
            create_undelegate_msg(
                sender.address(),
                validator,
                amount,
                self.network_config.staking_denomination,
            )
        )

        return prepare_and_broadcast_basic_transaction(
            self, tx, sender, gas_limit=gas_limit, memo=memo
        )

    def claim_rewards(
        self,
        validator: Address,
        sender: Wallet,
        memo: Optional[str] = None,
        gas_limit: Optional[int] = DEFAULT_TX_GAS_LIMIT,
    ) -> SubmittedTx:
        """claim rewards.

        :param validator: validator
        :param sender: sender
        :param memo: memo, defaults to None
        :param gas_limit: gas limit, defaults to None
        :return: prepare and broadcast the transaction and transaction details
        """
        tx = Transaction()
        tx.add_message(create_withdraw_delegator_reward(sender.address(), validator))

        return prepare_and_broadcast_basic_transaction(
            self, tx, sender, gas_limit=gas_limit, memo=memo
        )

    def estimate_gas_for_tx(self, tx: Transaction) -> int:
        """Estimate gas for transaction.

        :param tx: transaction
        :return: Estimated gas for transaction
        """
        return self._gas_strategy.estimate_gas(tx)

    def estimate_fee_from_gas(self, gas_limit: int) -> str:
        """Estimate fee from gas.

        :param gas_limit: gas limit
        :return: Estimated fee for transaction
        """
        fee = math.ceil(gas_limit * self.network_config.fee_minimum_gas_price)
        return f"{fee}{self.network_config.fee_denomination}"

    def estimate_gas_and_fee_for_tx(self, tx: Transaction) -> Tuple[int, str]:
        """Estimate gas and fee for transaction.

        :param tx: transaction
        :return: estimate gas, fee for transaction
        """
        gas_estimate = self.estimate_gas_for_tx(tx)
        fee = self.estimate_fee_from_gas(gas_estimate)
        return gas_estimate, fee

    def wait_for_query_tx(
        self,
        tx_hash: str,
        timeout: Optional[timedelta] = None,
        poll_period: Optional[timedelta] = None,
    ) -> TxResponse:
        """Wait for query transaction.

        :param tx_hash: transaction hash
        :param timeout: timeout, defaults to None
        :param poll_period: poll_period, defaults to None

        :raises QueryTimeoutError: timeout

        :return: transaction response
        """
        timeout = (
            ensure_timedelta(timeout)
            if timeout
            else timedelta(seconds=self._query_timeout_secs)
        )
        poll_period = (
            ensure_timedelta(poll_period)
            if poll_period
            else timedelta(seconds=self._query_interval_secs)
        )

        start = datetime.now()
        while True:
            try:
                return self.query_tx(tx_hash)
            except NotFoundError:
                pass

            delta = datetime.now() - start
            if delta >= timeout:
                raise QueryTimeoutError()

            time.sleep(poll_period.total_seconds())

    def query_tx(self, tx_hash: str) -> TxResponse:
        """query transaction.

        :param tx_hash: transaction hash
        :raises NotFoundError: Tx details not found
        :raises grpc.RpcError: RPC connection issue
        :return: query response
        """
        req = GetTxRequest(hash=tx_hash)
        try:
            resp = self.txs.GetTx(req)
        except grpc.RpcError as e:
            details = e.details()
            if "not found" in details:
                raise NotFoundError() from e
            raise
        except RuntimeError as e:
            details = str(e)
            if "tx" in details and "not found" in details:
                raise NotFoundError() from e
            raise

        return self._parse_tx_response(resp.tx_response)

    @staticmethod
    def _parse_tx_response(tx_response: Any) -> TxResponse:
        # parse the transaction logs
        logs = []
        for log_data in tx_response.logs:
            events = {}
            for event in log_data.events:
                events[event.type] = {a.key: a.value for a in event.attributes}
            logs.append(
                MessageLog(
                    index=int(log_data.msg_index), log=log_data.msg_index, events=events
                )
            )

        # parse the transaction events
        events = {}
        for event in tx_response.events:
            event_data = events.get(event.type, {})
            for attribute in event.attributes:
                event_data[attribute.key] = attribute.value
            events[event.type] = event_data

        timestamp = None
        if tx_response.timestamp:
            timestamp = isoparse(tx_response.timestamp)

        return TxResponse(
            hash=str(tx_response.txhash),
            height=int(tx_response.height),
            code=int(tx_response.code),
            gas_wanted=int(tx_response.gas_wanted),
            gas_used=int(tx_response.gas_used),
            raw_log=str(tx_response.raw_log),
            logs=logs,
            events=events,
            timestamp=timestamp,
        )

    def simulate_tx(self, tx: Transaction) -> int:
        """simulate transaction.

        :param tx: transaction
        :raises RuntimeError: Unable to simulate non final transaction
        :return: gas used in transaction
        """
        if tx.state != TxState.Final:
            raise RuntimeError("Unable to simulate non final transaction")

        req = SimulateRequest(tx=tx.tx)
        resp = self.txs.Simulate(req)

        return int(resp.gas_info.gas_used)

    def broadcast_tx(self, tx: Transaction) -> SubmittedTx:
        """Broadcast transaction.

        :param tx: transaction
        :return: Submitted transaction
        """
        # create the broadcast request
        broadcast_req = BroadcastTxRequest(
            tx_bytes=tx.tx.SerializeToString(), mode=BroadcastMode.BROADCAST_MODE_SYNC
        )

        # broadcast the transaction
        resp = self.txs.BroadcastTx(broadcast_req)
        tx_digest = resp.tx_response.txhash

        # check that the response is successful
        initial_tx_response = self._parse_tx_response(resp.tx_response)
        initial_tx_response.ensure_successful()

        return SubmittedTx(self, tx_digest)

    def query_latest_block(self) -> Block:
        """Query the latest block.

        :return: latest block
        """
        req = GetLatestBlockRequest()
        resp = self.tendermint.GetLatestBlock(req)
        return Block.from_proto(resp.block)

    def query_block(self, height: int) -> Block:
        """Query the block.

        :param height: block height
        :return: block
        """
        req = GetBlockByHeightRequest(height=height)
        resp = self.tendermint.GetBlockByHeight(req)
        return Block.from_proto(resp.block)

    def query_height(self) -> int:
        """Query the latest block height.

        :return: latest block height
        """
        return self.query_latest_block().height

    def query_chain_id(self) -> str:
        """Query the chain id.

        :return: chain id
        """
        return self.query_latest_block().chain_id
