from abc import ABC
import argparse
from typing import TypedDict, Optional

from hetu.utils import networking, Certificate
from hetu.utils.btlogging import logging
from hetu import settings
from hetu.config import Config
from hetu.chain_data import NeuronInfo, NeuronInfoLite
from hetu.utils import determine_chain_endpoint_and_network


class HetutensorMixin(ABC):
    network: str
    chain_endpoint: str
    log_verbose: bool

    def __str__(self):
        return f"Network: {self.network}, Chain: {self.chain_endpoint}"

    def __repr__(self):
        return self.__str__()

    def _check_and_log_network_settings(self):
        if (
            self.network == "finney"
            or self.chain_endpoint == settings.FINNEY_ENTRYPOINT
        ) and self.log_verbose:
            logging.info(
                f"You are connecting to {self.network} network with endpoint {self.chain_endpoint}."
            )
            logging.debug(
                "We strongly encourage running a local hetu node whenever possible. "
                "This increases decentralization and resilience of the network."
            )

    @staticmethod  # TODO can this be a class method?
    def config() -> "Config":
        """
        Creates and returns a Hetutensor configuration object.

        Returns:
            config (hetu.config.Config): A Hetutensor configuration object configured with arguments added by
                the `hetu.add_args` method.
        """
        parser = argparse.ArgumentParser()
        HetutensorMixin.add_args(parser)
        return Config(parser)

    @staticmethod
    def setup_config(network: Optional[str], config: "Config"):
        """
        Sets up and returns the configuration for the Hetutensor network and endpoint.

        This method determines the appropriate network and chain endpoint based on the provided network string or
            configuration object. It evaluates the network and endpoint in the following order of precedence:
            1. Provided network string.
            2. Configured chain endpoint in the `config` object.
            3. Configured network in the `config` object.
            4. Default chain endpoint.
            5. Default network.

        Arguments:
            network (Optional[str]): The name of the Hetutensor network. If None, the network and endpoint will be
                determined from the `config` object.
            config (hetu.config.Config): The configuration object containing the network and chain endpoint
                settings.

        Returns:
            tuple: A tuple containing the formatted WebSocket endpoint URL and the evaluated network name.
        """
        if network is None:
            candidates = [
                (
                    config.is_set("hetu.chain_endpoint"),
                    config.hetu.chain_endpoint,
                ),
                (config.is_set("hetu.network"), config.hetu.network),
                (
                    config.hetu.get("chain_endpoint"),
                    config.hetu.chain_endpoint,
                ),
                (config.hetu.get("network"), config.hetu.network),
            ]
            for check, config_network in candidates:
                if check:
                    network = config_network

        evaluated_network, evaluated_endpoint = determine_chain_endpoint_and_network(
            network
        )

        return networking.get_formatted_ws_endpoint_url(
            evaluated_endpoint
        ), evaluated_network

    @classmethod
    def help(cls):
        """Print help to stdout."""
        parser = argparse.ArgumentParser()
        cls.add_args(parser)
        print(cls.__new__.__doc__)
        parser.print_help()

    @classmethod
    def add_args(cls, parser: "argparse.ArgumentParser", prefix: Optional[str] = None):
        """
        Adds command-line arguments to the provided ArgumentParser for configuring the Hetutensor settings.

        Arguments:
            parser (argparse.ArgumentParser): The ArgumentParser object to which the Hetutensor arguments will be added.
            prefix (Optional[str]): An optional prefix for the argument names. If provided, the prefix is prepended to
                each argument name.

        Arguments added:
            --hetu.network: The Hetutensor network flag. Possible values are 'finney', 'test', 'archive', and
                'local'. Overrides the chain endpoint if set.
            --hetu.chain_endpoint: The Hetutensor chain endpoint flag. If set, it overrides the network flag.
            --hetu._mock: If true, uses a mocked connection to the chain.

        Example:
            parser = argparse.ArgumentParser()
            Hetutensor.add_args(parser)
        """
        prefix_str = "" if prefix is None else f"{prefix}."
        try:
            default_network = settings.DEFAULTS.hetu.network
            default_chain_endpoint = settings.DEFAULTS.hetu.chain_endpoint

            parser.add_argument(
                f"--{prefix_str}hetu.network",
                default=default_network,
                type=str,
                help="""The hetu network flag. The likely choices are:
                                        -- finney (main network)
                                        -- test (test network)
                                        -- archive (archive network +300 blocks)
                                        -- local (local running network)
                                    If this option is set it overloads hetu.chain_endpoint with
                                    an entry point node from that network.
                                    """,
            )
            parser.add_argument(
                f"--{prefix_str}hetu.chain_endpoint",
                default=default_chain_endpoint,
                type=str,
                help="""The hetu endpoint flag. If set, overrides the --network flag.""",
            )
            parser.add_argument(
                f"--{prefix_str}hetu._mock",
                default=False,
                type=bool,
                help="""If true, uses a mocked connection to the chain.""",
            )

        except argparse.ArgumentError:
            # re-parsing arguments.
            pass


class AxonServeCallParams:
    def __init__(
        self,
        version: int,
        ip: int,
        port: int,
        ip_type: int,
        netuid: int,
        hotkey: str,
        coldkey: str,
        protocol: int,
        placeholder1: int,
        placeholder2: int,
        certificate: Optional[Certificate],
    ):
        self.version = version
        self.ip = ip
        self.port = port
        self.ip_type = ip_type
        self.netuid = netuid
        self.hotkey = hotkey
        self.coldkey = coldkey
        self.protocol = protocol
        self.placeholder1 = placeholder1
        self.placeholder2 = placeholder2
        self.certificate = certificate

    def __eq__(self, other):
        if isinstance(other, self.__class__):
            return all(
                getattr(self, attr) == getattr(other, attr) for attr in self.__dict__
            )
        elif isinstance(other, dict):
            return all(getattr(self, attr) == other.get(attr) for attr in self.__dict__)
        elif isinstance(other, (NeuronInfo, NeuronInfoLite)):
            return all(
                [
                    self.version == other.axon_info.version,
                    self.ip == networking.ip_to_int(other.axon_info.ip),
                    self.port == other.axon_info.port,
                    self.ip_type == other.axon_info.ip_type,
                    self.netuid == other.netuid,
                    self.hotkey == other.hotkey,
                    self.coldkey == other.coldkey,
                    self.protocol == other.axon_info.protocol,
                    self.placeholder1 == other.axon_info.placeholder1,
                    self.placeholder2 == other.axon_info.placeholder2,
                ]
            )
        else:
            raise NotImplementedError(
                f"AxonServeCallParams equality not implemented for {type(other)}"
            )

    def copy(self) -> "AxonServeCallParams":
        return self.__class__(
            self.version,
            self.ip,
            self.port,
            self.ip_type,
            self.netuid,
            self.hotkey,
            self.coldkey,
            self.protocol,
            self.placeholder1,
            self.placeholder2,
            self.certificate,
        )

    def dict(self) -> dict:
        """
        Returns a dict representation of this object. If `self.certificate` is `None`,
        it is not included in this.
        """
        d = {
            "version": self.version,
            "ip": self.ip,
            "port": self.port,
            "ip_type": self.ip_type,
            "netuid": self.netuid,
            "hotkey": self.hotkey,
            "coldkey": self.coldkey,
            "protocol": self.protocol,
            "placeholder1": self.placeholder1,
            "placeholder2": self.placeholder2,
        }
        if self.certificate is not None:
            d["certificate"] = self.certificate
        return d


class PrometheusServeCallParams(TypedDict):
    """Prometheus serve chain call parameters."""

    version: int
    ip: int
    port: int
    ip_type: int
    netuid: int


class ParamWithTypes(TypedDict):
    name: str  # Name of the parameter.
    type: str  # ScaleType string of the parameter.
