import logging
from datetime import datetime, timedelta, date
from dateutil import tz

from influxdb_client.client.flux_table import TableList


from algo.data.connection import DataConnection
from algo.data.common import PricePoint, PriceStream
from algo.util import d, dt_to_rfc3339, get_week_start_day, is_dt_naive
from algo.constant import Symbol
from algo.config import config


log = logging.getLogger(__name__)


class DataReader(DataConnection):
    def __init__(self, url, token=None, org=None):
        super().__init__(url, token, org)
        self._query_api = self._client.query_api()

    # TBD: consider removing
    def hour_data(self, hour_start: datetime, symbol: Symbol) -> PriceStream:
        """
        treats timezone naive data in hour_data as UTC
        """
        start = hour_start.replace(tzinfo=tz.tzutc()) if is_dt_naive(hour_start) else hour_start
        end = start + timedelta(hours=1)
        yield from self.get_data_stream(start, end, symbol)

    def day_data(self, day: date, symbol: Symbol) -> PriceStream:
        start = datetime.combine(
            day, datetime.min.time(), tz.tzutc()
        )
        end = start + timedelta(days=1)
        yield from self.get_data_stream(start, end, symbol)

    def week_data(self, week_day: date, symbol: Symbol) -> PriceStream:
        week_start_date = get_week_start_day(week_day)
        for days_offset in range(7):
            current_day = week_start_date + timedelta(days=days_offset)
            log.info('getting week data day by the day', extra=d(week_start=week_day, current_day=current_day))
            yield from self.day_data(current_day, symbol)

    def get_result_set(self, symbol: Symbol, start: datetime, end: datetime) -> TableList:
        query_str = f"""
            from(bucket: "tsdata")
              |> range(start: {dt_to_rfc3339(start)}, stop: {dt_to_rfc3339(end)})
              |> filter(fn: (r) => r["instrument"] == "{symbol.value}")
              |> filter(fn: (r) => r["_measurement"] == "price")
              |> pivot(rowKey: ["_time"], columnKey: ["_field"], valueColumn: "_value")
        """
        log.debug('q: %s', query_str)
        return self._query_api.query(query_str)

    def get_data_stream(
            self, start: datetime, end: datetime, symbol: Symbol
    ) -> PriceStream:
        log.debug('getting chunk starting', extra=d(start))
        results = self.get_result_set(symbol, start, end)
        log.info('query-result', extra={'r': results})

        for _time, ask, bid in results.to_values(columns=('_time', 'ask', 'bid')):
            yield PricePoint(_time, ask, bid)


def get_data_reader() -> DataReader:
    return DataReader(
        config.influx_url,
        config.influx_token,
        config.influx_org
    )
