# 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.

import asyncio
import base64
import datetime
import functools
import hashlib
import json
import logging
import os
import os.path
import platform
import threading
import typing
from urllib.request import urlopen

import dateutil.parser
import websockets

logger = logging.getLogger(__name__)


# We need to use this syntax of creating TypeDict because “from” is a reserved keyword in Python.
TransferTx = typing.TypedDict(
    'TransferTx',
    {
        'type': int,
        'from': bytes,
        'to': bytes,
        'value': int,
        'message': bytes,
        'hash': str,
        'height': int,
        'fee': int,
        'raw': bytes,
        'index': int,
        'block_time': datetime.datetime,
    },
)


def _b64_to_bytes(b64: str) -> bytes:
    return base64.decodebytes(b64.encode('ascii'))


def _b64_to_str(b64: str) -> str:
    return _b64_to_bytes(b64).decode('ascii')


def _iter_event_attrs(attrs) -> typing.Iterator[typing.Tuple[str, typing.Optional[str]]]:
    for attr in attrs:
        key = _b64_to_str(attr['key'])
        value = None if (maybe_value_b64 := attr['value']) is None else _b64_to_str(maybe_value_b64)
        yield key, value


def _events_to_dict(events):
    d = {}
    for event in events:
        assert event['type'] == 'tx'
        for key, value in _iter_event_attrs(event['attributes']):
            if key in {'type', 'value', 'height', 'fee', 'index'}:
                value_prime = int(value)
            elif key in {'from', 'to', 'message'}:
                value_prime = _b64_to_bytes(value or '')
            else:
                value_prime = value
            d[key] = value_prime
    return d


def get_local_data_dir():
    system = platform.system()
    if system == 'Windows':
        return os.environ['LocalAppData']
    elif system == 'Darwin':
        return os.path.join(
            os.environ['HOME'],
            'Library',
            'Application Support',
        )
    else:
        return os.getenv(
            'XDG_DATA_HOME',
            os.path.join(os.environ['HOME'], '.local', 'share'),
        )


JSONRequestID = int


class TendermintServerError(Exception):
    pass


class JSONRPCClient:
    def __init__(self, ws):
        self._ws = ws
        self._counter = 0
        self._counter_lock = threading.Lock()

    def _get_new_id(self):
        self._counter_lock.acquire()
        self._counter += 1
        try:
            return self._counter
        finally:
            self._counter_lock.release()

    async def req(self, method, params) -> JSONRequestID:
        id_ = self._get_new_id()
        body = json.dumps(
            {
                'jsonrpc': '2.0',
                'id': id_,
                'method': method,
                'params': params,
            })
        await self._ws.send(body)
        return id_

    async def recv(self) -> dict:
        msg = json.loads(await self._ws.recv())
        if msg.get('error'):
            error_desc = msg["error"]["data"]
            logger.error(f'Received server error: “{error_desc}”')
            raise TendermintServerError(error_desc)
        return msg


Height = typing.NewType('Height', int)
TxHash = typing.NewType('TxHash', str)
Base64String = typing.NewType('Base64String', str)
Address = typing.NewType('Address', typing.Union[Base64String, bytes])


def _maybe_enforce_b64(maybe_address: typing.Optional[Address]) -> typing.Optional[Base64String]:
    if isinstance(maybe_address, str):
        return maybe_address
    elif isinstance(maybe_address, bytes):
        return base64.encodebytes(maybe_address).decode("ascii").strip()


class ErcoinReactor:
    retry_timeout = 10

    def __init__(
            self,
            *,
            node,
            to_address: typing.Optional[Address] = None,
            from_address: typing.Optional[Address] = None,
            ssl=True,
            port=26657,
    ):
        """Initialize reactor.

        to_address and from_address cannot be given simultaneously.
        """

        if to_address and from_address:
            raise ValueError('At most one address can be given.')

        self._node = node
        self._to_address_b64 = _maybe_enforce_b64(to_address)
        self._from_address_b64 = _maybe_enforce_b64(from_address)
        self._port = port
        self._ssl = ssl

        self._tx_buf: typing.List[TransferTx] = []
        self._seen_tx_hashes: typing.Set[str] = set()
        self._catched_up = False

    @functools.cached_property
    def _ws_endpoint(self):
        protocol = 'wss' if self._ssl else 'ws'
        return f'{protocol}://{self._node}:{self._port}/websocket'

    @functools.cached_property
    def _http_endpoint(self):
        protocol = "https" if self._ssl else "http"
        return f"{protocol}://{self._node}:{self._port}"

    @functools.cached_property
    def _subscription_query(self):
        return self._join_queries([
            "tm.event='Tx'",
            self._base_tx_query,
        ])

    @functools.cached_property
    def _catchup_query(self):
        return self._join_queries([
            self._base_tx_query,
            f'tx.height >= {self._last_height}',
        ])

    @functools.cached_property
    def _base_tx_query(self):
        if self._to_address_b64:
            return f"tx.to='{self._to_address_b64}'"
        elif self._from_address_b64:
            return f"tx.from='{self._from_address_b64}'"
        return ''

    @staticmethod
    def _join_queries(queries):
        return ' AND '.join(query for query in queries if query)

    @functools.cached_property
    def _state_filepath(self):
        h = hashlib.new('sha224')
        h.update(self._base_tx_query.encode('ascii'))
        basename = h.digest().hex()

        return os.path.join(
            get_local_data_dir(),
            'ern_reactor',
            self.get_namespace(),
            f'{basename}.json',
        )

    def _load_state(self) -> typing.Optional[dict]:
        state_filename = self._state_filepath
        if os.path.isfile(state_filename):
            with open(state_filename) as f:
                return json.load(f)

    def _dump_state(self, state):
        tmp_filepath = self._state_filepath + '.tmp'
        dirpath = os.path.dirname(self._state_filepath)
        if not os.path.isdir(dirpath):
            os.makedirs(dirpath)
        with open(tmp_filepath, 'w') as f:
            json.dump(state, f)
        os.replace(tmp_filepath, self._state_filepath)

    def _load_sync_status(self) -> None:
        if (state := self._load_state()) is None:
            self._last_height = 0
            self._last_index = 0
        else:
            self._last_height = state['height']
            self._last_index = state['index']

    def _save_sync_status(self) -> None:
        state = self._load_state() or {}
        state['height'] = self._last_height
        state['index'] = self._last_index
        self._dump_state(state)

    async def _start_catchup(self):
        self._last_catchup_page = 0
        return await self._continue_catchup()

    async def _continue_catchup(self) -> JSONRequestID:
        self._last_catchup_page += 1
        return await self._rpc_client.req(
            'tx_search',
            {
                'query': self._catchup_query,
                'page': str(self._last_catchup_page),
                'per_page': '100',
            },
        )

    async def _handle_connection(self, ws):
        self._rpc_client = JSONRPCClient(ws)
        subscribe_id = await self._rpc_client.req('subscribe', {'query': self._subscription_query})
        subscribe_resp = await self._rpc_client.recv()
        assert subscribe_resp.get('result') == {}
        catchup_id = await self._start_catchup()
        while True:
            msg = await self._rpc_client.recv()
            if (msg_id := msg['id']) == catchup_id:
                self._catched_up = int(msg['result']['total_count']) <= self._last_catchup_page * 100
                if not self._catched_up:
                    catchup_id = await self._continue_catchup()
                await self._handle_txs(self._search_result_to_tx(res) for res in msg['result']['txs'])
                if self._catched_up:
                    await self._handle_txs(self._tx_buf)
            elif msg_id == subscribe_id:
                assert msg['result']['data']['type'] == 'tendermint/event/Tx'
                await self._handle_tx(self._subscription_msg_to_tx(msg))

    async def start(self) -> typing.NoReturn:
        self._load_sync_status()
        while True:
            try:
                async with websockets.connect(self._ws_endpoint) as ws:
                    await self._handle_connection(ws)
            except (websockets.exceptions.ConnectionClosedError, ConnectionError, TendermintServerError):
                logger.warning(f'Connection error, sleeping {self.retry_timeout}s before retrying…')
                await asyncio.sleep(self.retry_timeout)

    @functools.lru_cache(maxsize=1)
    def _fetch_block_time(self, height: int):
        # It would be more efficient to reuse the WebSocket connection, but it would be more problematic in terms of caching and message ordering, so for now we just do a HTTP call.
        url = f"{self._http_endpoint}/block?height={height}"
        with urlopen(url) as f:
            resp = json.load(f)
            timestamp_str = resp["result"]["block"]["header"]["time"]
            return dateutil.parser.isoparse(timestamp_str)

    def _subscription_msg_to_tx(self, tx_msg) -> TransferTx:
        tx_result = tx_msg['result']['data']['value']['TxResult']
        tx = _events_to_dict(tx_result['result']['events'])
        tx['raw'] = _b64_to_bytes(tx_result['tx'])
        tx['index'] = tx_result['index']
        tx['height'] = int(tx_result['height'])
        tx['block_time'] = self._fetch_block_time(tx['height'])
        for key, value in tx_msg['result']['events'].items():
            if key == 'tx.hash':
                tx['hash'] = value[0]
                break
        return TransferTx(**tx)

    def _search_result_to_tx(self, search_result) -> TransferTx:
        tx = _events_to_dict(search_result['tx_result']['events'])
        tx['raw'] = _b64_to_bytes(search_result['tx'])
        tx['index'] = search_result['index']
        tx['hash'] = search_result['hash']
        tx['height'] = int(search_result['height'])
        tx['block_time'] = self._fetch_block_time(tx['height'])
        return TransferTx(**tx)

    async def _handle_tx(self, tx: TransferTx):
        index = tx['index']
        height = tx['height']
        if (height, index) <= (self._last_height, self._last_index):
            logger.debug(f'Omitting already processed transaction {tx["hash"]} at height {height} and index {index}.')
            return
        await self.process_tx(tx)
        self._last_height = tx['height']
        self._last_index = tx['index']
        self._save_sync_status()

    async def _handle_txs(self, txs):
        for tx in txs:
            await self._handle_tx(tx)

    def get_namespace(self) -> str:
        """Return namespace used to distinguish different instances of the reactor."""
        raise NotImplementedError()

    async def process_tx(self, tx: TransferTx) -> None:
        raise NotImplementedError()
