from __future__ import annotations
from functools import cached_property
import json
import os
import time
from typing import Any, Dict, List
from datetime import datetime
import polars as pl
from web3.constants import ADDRESS_ZERO
from web3.contract import Contract
import numpy as np

from web3cat.fetcher.erc20_metas import ERC20Meta
from web3cat.fetcher.events import Event
from web3cat.data.core import DataCore

RESOLVER_MAPPING = {"weth": "eth", "wbtc": "btc"}


class ChainlinkUSDData(DataCore):
    """
    Chainlink data for a specific token.
    """

    _token: str

    def __init__(
        self, token: str, start: int | datetime, end: int | datetime, **kwargs
    ):
        super().__init__(start, end, **kwargs)
        self._token = token

    @cached_property
    def meta(self) -> ERC20Meta:
        """
        Metadata for tokens (like name, symbol and decimals)
        """
        return self._erc20_metas_service.get(self._token)

    @cached_property
    def oracle_proxy_contract(self) -> Contract:
        """
        A reference to Chainlink oracle proxy contract
        """
        current_folder = os.path.realpath(os.path.dirname(__file__))
        with open(
            f"{current_folder}/oracle_proxy.abi.json", "r", encoding="utf-8"
        ) as f:
            abi = json.load(f)
        oracle_address = self._resolve_chainlink_address(self.meta.symbol.lower())
        if oracle_address is None:
            raise LookupError(
                f"Chainlink oracle for token `{self.meta.symbol.lower()}` "
                f"on chain with id `{self._blocks_service.chain_id}` is not found"
            )

        return self.w3.eth.contract(
            address=self.w3.to_checksum_address(oracle_address), abi=abi
        )

    @cached_property
    def oracle_aggregator_contract(self) -> Contract:
        """
        A reference to Chainlink oracle aggregator contract
        """

        current_folder = os.path.realpath(os.path.dirname(__file__))
        with open(
            f"{current_folder}/oracle_aggregator.abi.json", "r", encoding="utf-8"
        ) as f:
            abi = json.load(f)
        oracle_address = self._calls_service.get_call(
            self.oracle_proxy_contract.functions.aggregator(),
            self.to_block_number,
        ).response
        return self.w3.eth.contract(
            address=self.w3.to_checksum_address(oracle_address),
            abi=abi,
        )

    @cached_property
    def oracle_decimals(self) -> int:
        """
        Chainlink oracle decimals
        """
        return int(
            self._calls_service.get_call(
                self.oracle_proxy_contract.functions.decimals(),
                self.to_block_number,
            ).response
        )

    @cached_property
    def initial_price(self) -> np.float64:
        """
        Oracle price at from_block
        """
        price = self._calls_service.get_call(
            self.oracle_aggregator_contract.functions.latestRoundData(),
            self.from_block_number,
        ).response[1]
        return price / 10**self.oracle_decimals

    @cached_property
    def index(self) -> Dict[str, Any]:
        """
        Oracles index
        """
        current_folder = os.path.realpath(os.path.dirname(__file__))
        with open(f"{current_folder}/oracles.json", "r", encoding="utf-8") as f:
            return json.load(f)

    def prices(self, timepoints: List[int | datetime]) -> pl.DataFrame:
        """
        Chainlink price feed.

        Arguments:
            timepoints: A list of timepoints for prices
        """

        blocks = self._resolve_timepoints(sorted(timepoints))
        updates = self.updates[["block_number", "price"]].to_dicts()
        i = 0
        price_list = []
        while len(updates) > 0 and blocks[i].number < updates[0]["block_number"]:
            price_list.append(self.initial_price)
            i += 1
            if i == len(blocks):
                break
        j = 0
        for b in blocks[i:]:
            while j < len(updates) and updates[j]["block_number"] <= b.number:
                j += 1
            # now ts < updates[j]
            idx = max([j - 1, 0])
            price_list.append(updates[idx]["price"])
        out = [
            {
                "timestamp": b.timestamp,
                "date": datetime.fromtimestamp(b.timestamp),
                "block_number": b.number,
                "price": price,
            }
            for b, price in zip(blocks, price_list)
        ]
        return pl.DataFrame(
            out,
            {
                "timestamp": pl.UInt64,
                "date": pl.Datetime,
                "block_number": pl.UInt64,
                "price": pl.Float64,
            },
        )

    @cached_property
    def updates(self) -> pl.DataFrame:
        """
        A list of oracle updates
        """
        events: List[Event] = self._events_service.get_events(
            self.oracle_aggregator_contract.events.AnswerUpdated,
            self.from_block_number,
            self.to_block_number,
        )

        block_numbers = [e.block_number for e in events]
        blocks = self._blocks_service.get_blocks(block_numbers)
        timestamps = [b.timestamp for b in blocks]
        ts_index = dict(zip(block_numbers, timestamps))
        factor = 10**self.oracle_decimals
        return pl.DataFrame(
            [self._event_to_row(e, ts_index[e.block_number], factor) for e in events],
            {
                "timestamp": pl.UInt64,
                "date": pl.Datetime,
                "block_number": pl.UInt64,
                "transaction_hash": pl.Utf8,
                "log_index": pl.UInt64,
                "round": pl.UInt64,
                "updated_at": pl.UInt64,
                "price": pl.Float64,
            },
        ).sort(pl.col("timestamp"))

    def _event_to_row(self, ev: Event, ts: int, val_factor: int) -> Dict[str, Any]:
        return {
            "timestamp": ts,
            "date": datetime.fromtimestamp(ts),
            "block_number": ev.block_number,
            "transaction_hash": ev.transaction_hash,
            "log_index": ev.log_index,
            "round": ev.args["roundId"],
            "updated_at": ev.args["updatedAt"],
            "price": ev.args["current"] / val_factor,
        }

    def _resolve_chainlink_address(self, token: str) -> str | None:
        cid = str(self._blocks_service.chain_id)
        if not cid in self.index:
            return None
        oracles = self.index[cid]
        if token in oracles:
            return oracles[token]["address"]
        token = RESOLVER_MAPPING.get(token, token)
        if not token in oracles:
            return None
        return oracles[token]["address"]

    def _resolve_timestamps(self, timestamps: List[int | datetime]) -> List[int]:
        resolved = []
        for ts in timestamps:
            # resolve datetimes to timestamps
            if isinstance(ts, datetime):
                resolved.append(int(time.mktime(ts.timetuple())))
            else:
                resolved.append(ts)
        return resolved


class ChainlinkData(DataCore):
    """
    Chainlink price feed for tokens.

    Arguments:
        tokens: a list of ERC20 tokens
        start: start of a dataset
        end: end of a dataset
        kwargs: arguments for fetchers, see :class:`fetcher.core.Core`
    """

    _datas: Dict[str, ChainlinkUSDData]
    _kwargs: Dict[str, Any]

    def __init__(
        self,
        tokens: List[str],
        start: int | datetime,
        end: int | datetime,
        **kwargs,
    ):
        super().__init__(start, end, **kwargs)
        self._datas = {}
        self._kwargs = kwargs
        for token in tokens:
            self.add_token(token)

    def add_token(
        self,
        token: str,
        start: int | datetime | None = None,
        end: int | datetime | None = None,
        proxy: str | None = None,
    ):
        """
        Add token to prices dataset.

        Arguments:
            token: token to add
            start: start of the data
            end: end of the data
        """
        if token == ADDRESS_ZERO or token.upper() == "USD":
            return
        meta = self._erc20_metas_service.get(token)
        if proxy is None:
            self._datas[meta.address] = ChainlinkUSDData(
                token, start or self.start, end or self.end, **self._kwargs
            )
        else:
            proxy_meta = self._erc20_metas_service.get(token)
            self._datas[meta.address] = self._datas[proxy_meta.address]

    def get_data(self, token: str) -> ChainlinkUSDData:
        """
        Get chainlink USD data for token.

        Agruments:
            token: ERC20 token

        Returns:
            Chainlink USD data for token
        """

        meta = self._erc20_metas_service.get(token)
        if not meta.address in self._datas:
            self.add_token(token)
        return self._datas[meta.address]

    def prices(
        self, token0: str, token1: str, timepoints: List[int | datetime]
    ) -> pl.DataFrame:
        """
        Chainlink prices for pair of tokens.
        The token0 = "WETH", token1 = "USDC" would mean the regular WETH/USDC
        (that is how much dollars we give for 1 ETH)

        Dataframe with balances for addresses over time.

        Args:
            token0: The first token
            token1: The second token
            timepoints: A list of timepoints (see :class:`ERC20Data`)

        Returns:
            A Dataframe with fields

        +----------------------+----------------------------+------------------------------------------------------------------------------+
        | Field                | Type                       | Description                                                                  |
        +======================+============================+==============================================================================+
        | ``timestamp``        | :attr:`numpy.int64`        | Timestamp for the snapshot of the balance                                    |
        +----------------------+----------------------------+------------------------------------------------------------------------------+
        | ``date``             | :class:`datetime.datetime` | Date for the timestamp                                                       |
        +----------------------+----------------------------+------------------------------------------------------------------------------+
        | ``block_number``     | :class:`int`               | Number of the block                                                          |
        +----------------------+----------------------------+------------------------------------------------------------------------------+
        | ``price``            | :attr:`numpy.float64`      | Relative price                                                               |
        +----------------------+----------------------------+------------------------------------------------------------------------------+

        """
        blocks = self._resolve_timepoints(timepoints)
        token0_is_usd = token0.upper() == "USD" or token0 == ADDRESS_ZERO
        token1_is_usd = token1.upper() == "USD" or token1 == ADDRESS_ZERO

        if token0_is_usd and token1_is_usd:
            prices = [1.0] * len(timepoints)
        elif token0_is_usd:
            data = self.get_data(token1)
            prices = [1.0 / p for p in data.prices(timepoints)["price"]]
        elif token1_is_usd:
            data = self.get_data(token0)
            prices = data.prices(timepoints)["price"]
        else:
            token0_data = self.get_data(token0)
            token1_data = self.get_data(token1)
            block_numbers = [b.number for b in blocks]
            prices0 = token0_data.prices(block_numbers)["price"].to_list()
            prices1 = token1_data.prices(block_numbers)["price"].to_list()
            prices = [p0 / p1 for p0, p1 in zip(prices0, prices1)]

        out = [
            {
                "timestamp": b.timestamp,
                "date": datetime.fromtimestamp(b.timestamp),
                "block_number": b.number,
                "price": p,
            }
            for b, p in zip(blocks, prices)
        ]
        df = pl.DataFrame(
            out,
            {
                "timestamp": pl.UInt64,
                "date": pl.Datetime,
                "block_number": pl.UInt64,
                "price": pl.Float64,
            },
        )
        return df
