# piker: trading gear for hackers
# Copyright (C) Tyler Goodlet (in stewardship for pikers)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

'''
Real-time and historical data feed endpoints.

'''
from contextlib import (
    asynccontextmanager as acm,
    aclosing,
)
from datetime import datetime
from typing import (
    AsyncGenerator,
    Callable,
    Optional,
)
import time

import numpy as np
import pendulum
from trio_typing import TaskStatus
import trio

from piker.accounting._mktinfo import (
    MktPair,
)
from piker.brokers import (
    open_cached_client,
)
from piker.brokers._util import (
    BrokerError,
    DataThrottle,
    DataUnavailable,
)
from piker.types import Struct
from piker.data.validate import FeedInit
from piker.data._web_bs import open_autorecon_ws, NoBsWs
from .api import (
    log,
)
from .symbols import get_mkt_info


class OHLC(Struct, frozen=True):
    '''
    Description of the flattened OHLC quote format.

    For schema details see:
        https://docs.kraken.com/websockets/#message-ohlc

    '''
    chan_id: int  # internal kraken id
    chan_name: str  # eg. ohlc-1  (name-interval)
    pair: str  # fx pair

    # unpacked from array
    time: float  # Begin time of interval, in seconds since epoch
    etime: float  # End time of interval, in seconds since epoch
    open: float  # Open price of interval
    high: float  # High price within interval
    low: float  # Low price within interval
    close: float  # Close price of interval
    vwap: float  # Volume weighted average price within interval
    volume: float  # Accumulated volume **within interval**
    count: int  # Number of trades within interval


async def stream_messages(
    ws: NoBsWs,
):
    '''
    Message stream parser and heartbeat handler.

    Deliver ws subscription messages as well as handle heartbeat logic
    though a single async generator.

    '''
    last_hb: float = 0

    async for msg in ws:
        match msg:
            case {'event': 'heartbeat'}:
                now = time.time()
                delay = now - last_hb
                last_hb = now

                # XXX: why tf is this not printing without --tl flag?
                log.debug(f"Heartbeat after {delay}")
                # print(f"Heartbeat after {delay}")

                continue

            case _:
                # passthrough sub msgs
                yield msg


async def process_data_feed_msgs(
    ws: NoBsWs,
):
    '''
    Parse and pack data feed messages.

    '''
    async with aclosing(stream_messages(ws)) as ws_stream:
        async for msg in ws_stream:
            match msg:
                case {
                    'errorMessage': errmsg
                }:
                    raise BrokerError(errmsg)

                case {
                    'event': 'subscriptionStatus',
                } as sub:
                    log.info(
                        'WS subscription is active:\n'
                        f'{sub}'
                    )
                    continue

                case [
                    chan_id,
                    *payload_array,
                    chan_name,
                    pair
                ]:
                    if 'ohlc' in chan_name:
                        array: list = payload_array[0]
                        ohlc = OHLC(
                            chan_id,
                            chan_name,
                            pair,
                            *map(float, array[:-1]),
                            count=array[-1],
                        )
                        yield 'ohlc', ohlc.copy()

                    elif 'spread' in chan_name:

                        bid, ask, ts, bsize, asize = map(
                            float, payload_array[0])

                        # TODO: really makes you think IB has a horrible API...
                        quote = {
                            'symbol': pair.replace('/', ''),
                            'ticks': [
                                {'type': 'bid', 'price': bid, 'size': bsize},
                                {'type': 'bsize', 'price': bid, 'size': bsize},

                                {'type': 'ask', 'price': ask, 'size': asize},
                                {'type': 'asize', 'price': ask, 'size': asize},
                            ],
                        }
                        yield 'l1', quote

                    # elif 'book' in msg[-2]:
                    #     chan_id, *payload_array, chan_name, pair = msg
                    #     print(msg)

                case {
                    'connectionID': conid,
                    'event': 'systemStatus',
                    'status': 'online',
                    'version': ver,
                }:
                    log.info(
                        f'Established {ver} ws connection with id: {conid}'
                    )
                    continue

                case _:
                    print(f'UNHANDLED MSG: {msg}')
                    # yield msg


def normalize(ohlc: OHLC) -> dict:
    '''
    Norm an `OHLC` msg to piker's minimal (live-)quote schema.

    '''
    quote = ohlc.to_dict()
    quote['broker_ts'] = quote['time']
    quote['brokerd_ts'] = time.time()
    quote['symbol'] = quote['pair'] = quote['pair'].replace('/', '')
    quote['last'] = quote['close']
    quote['bar_wap'] = ohlc.vwap
    return quote


@acm
async def open_history_client(
    mkt: MktPair,

) -> AsyncGenerator[Callable, None]:

    symbol: str = mkt.bs_mktid

    # TODO implement history getter for the new storage layer.
    async with open_cached_client('kraken') as client:

        # lol, kraken won't send any more then the "last"
        # 720 1m bars.. so we have to just ignore further
        # requests of this type..
        queries: int = 0

        async def get_ohlc(
            timeframe: float,
            end_dt: Optional[datetime] = None,
            start_dt: Optional[datetime] = None,

        ) -> tuple[
            np.ndarray,
            datetime,  # start
            datetime,  # end
        ]:

            nonlocal queries
            if (
                queries > 0
                or timeframe != 60
            ):
                raise DataUnavailable(
                    'Only a single query for 1m bars supported')

            count = 0
            while count <= 3:
                try:
                    array = await client.bars(
                        symbol,
                        since=end_dt,
                    )
                    count += 1
                    queries += 1
                    break
                except DataThrottle:
                    log.warning(f'kraken OHLC throttle for {symbol}')
                    await trio.sleep(1)

            start_dt = pendulum.from_timestamp(array[0]['time'])
            end_dt = pendulum.from_timestamp(array[-1]['time'])
            return array, start_dt, end_dt

        yield get_ohlc, {'erlangs': 1, 'rate': 1}


async def stream_quotes(

    send_chan: trio.abc.SendChannel,
    symbols: list[str],
    feed_is_live: trio.Event,
    loglevel: str = None,

    # backend specific
    sub_type: str = 'ohlc',

    # startup sync
    task_status: TaskStatus[tuple[dict, dict]] = trio.TASK_STATUS_IGNORED,

) -> None:
    '''
    Subscribe for ohlc stream of quotes for ``pairs``.

    ``pairs`` must be formatted <crypto_symbol>/<fiat_symbol>.

    '''

    ws_pairs: list[str] = []
    init_msgs: list[FeedInit] = []

    async with (
        send_chan as send_chan,
    ):
        for sym_str in symbols:
            mkt, pair = await get_mkt_info(sym_str)
            init_msgs.append(
                FeedInit(mkt_info=mkt)
            )

            ws_pairs.append(pair.wsname)

        @acm
        async def subscribe(ws: NoBsWs):

            # XXX: setup subs
            # https://docs.kraken.com/websockets/#message-subscribe
            # specific logic for this in kraken's sync client:
            # https://github.com/krakenfx/kraken-wsclient-py/blob/master/kraken_wsclient_py/kraken_wsclient_py.py#L188
            ohlc_sub = {
                'event': 'subscribe',
                'pair': ws_pairs,
                'subscription': {
                    'name': 'ohlc',
                    'interval': 1,
                },
            }

            # TODO: we want to eventually allow unsubs which should
            # be completely fine to request from a separate task
            # since internally the ws methods appear to be FIFO
            # locked.
            await ws.send_msg(ohlc_sub)

            # trade data (aka L1)
            l1_sub = {
                'event': 'subscribe',
                'pair': ws_pairs,
                'subscription': {
                    'name': 'spread',
                    # 'depth': 10}
                },
            }

            # pull a first quote and deliver
            await ws.send_msg(l1_sub)

            yield

            # unsub from all pairs on teardown
            if ws.connected():
                await ws.send_msg({
                    'pair': ws_pairs,
                    'event': 'unsubscribe',
                    'subscription': ['ohlc', 'spread'],
                })

                # XXX: do we need to ack the unsub?
                # await ws.recv_msg()

        # see the tips on reconnection logic:
        # https://support.kraken.com/hc/en-us/articles/360044504011-WebSocket-API-unexpected-disconnections-from-market-data-feeds
        ws: NoBsWs
        async with (
            open_autorecon_ws(
                'wss://ws.kraken.com/',
                fixture=subscribe,
                reset_after=20,
            ) as ws,

            # avoid stream-gen closure from breaking trio..
            # NOTE: not sure this actually works XD particularly
            # if we call `ws._connect()` manally in the streaming
            # async gen..
            aclosing(process_data_feed_msgs(ws)) as msg_gen,
        ):
            # pull a first quote and deliver
            typ, ohlc_last = await anext(msg_gen)
            quote = normalize(ohlc_last)

            task_status.started((init_msgs,  quote))
            feed_is_live.set()

            # keep start of last interval for volume tracking
            last_interval_start: float = ohlc_last.etime

            # start streaming
            topic: str = mkt.bs_fqme
            async for typ, quote in msg_gen:
                match typ:

                    # TODO: can get rid of all this by using
                    # ``trades`` subscription..? Not sure why this
                    # wasn't used originally? (music queues) zoltannn..
                    # https://docs.kraken.com/websockets/#message-trade
                    case 'ohlc':
                        # generate tick values to match time & sales pane:
                        # https://trade.kraken.com/charts/KRAKEN:BTC-USD?period=1m
                        volume = quote.volume

                        # new OHLC sample interval
                        if quote.etime > last_interval_start:
                            last_interval_start: float = quote.etime
                            tick_volume: float = volume

                        else:
                            # this is the tick volume *within the interval*
                            tick_volume: float = volume - ohlc_last.volume

                        ohlc_last = quote
                        last = quote.close

                        quote = normalize(quote)
                        ticks = quote.setdefault(
                            'ticks',
                            [],
                        )
                        if tick_volume:
                            ticks.append({
                                'type': 'trade',
                                'price': last,
                                'size': tick_volume,
                            })

                    case 'l1':
                        # passthrough quote msg
                        pass

                    case _:
                        log.warning(f'Unknown WSS message: {typ}, {quote}')

                await send_chan.send({topic: quote})
