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

'''
`nativedb`: a lulzy Apache-parquet file manager (that some might
            call a poor man's tsdb).

AKA a `piker`-native file-system native "time series database"
without needing an extra process and no standard TSDB features,
YET!

'''
# TODO: like there's soo much..
# - better name like "parkdb" or "nativedb" (lel)? bundle this lib with
#   others to make full system:
#   - tractor for failover and reliablity?
#   - borg for replication and sync?
#
# - use `fastparquet` for appends:
#   https://fastparquet.readthedocs.io/en/latest/api.html#fastparquet.write
#   (presuming it's actually faster then overwrites and
#   makes sense in terms of impl?)
#
# - use `polars` support for lazy scanning, processing and schema
#   validation?
#   - https://pola-rs.github.io/polars-book/user-guide/io/parquet/#scan
#   - https://pola-rs.github.io/polars-book/user-guide/concepts/lazy-vs-eager/
# - consider delta writes for appends?
#   - https://github.com/pola-rs/polars/blob/main/py-polars/polars/dataframe/frame.py#L3232
# - consider multi-file appends with appropriate time-range naming?
#   - https://pola-rs.github.io/polars-book/user-guide/io/multiple/
#
# - use `borg` for replication?
#   - https://borgbackup.readthedocs.io/en/stable/quickstart.html#remote-repositories
#   - https://github.com/borgbackup/borg
#   - https://borgbackup.readthedocs.io/en/stable/faq.html#usage-limitations
#   - https://github.com/borgbackup/community
#   - https://github.com/spslater/borgapi
#   - https://nixos.wiki/wiki/ZFS
from __future__ import annotations
from collections.abc import Iterator
from contextlib import asynccontextmanager as acm
from dataclasses import dataclass
from datetime import datetime
import os
from pathlib import Path
from tempfile import NamedTemporaryFile
import time

from bidict import bidict
import numpy as np
import polars as pl
from pendulum import (
    from_timestamp,
)

from piker import config
from piker import tsp
from tractor.ipc._shm import ShmArray
from piker.data import def_iohlcv_fields
from piker.log import get_logger
from . import TimeseriesNotFound


log = get_logger('storage.nativedb')
_price_fields: tuple[str, ...] = (
    'open',
    'high',
    'low',
    'close',
)


@dataclass(frozen=True)
class NativeSeriesRef:
    '''
    Identify one canonical NativeDB OHLCV file.

    '''
    fqme: str
    period_s: int
    path: Path


def detect_period(shm: ShmArray) -> float:
    '''
    Attempt to detect the series time step sampling period
    in seconds.

    '''
    # TODO: detect sample rate helper?
    # calc ohlc sample period for naming
    ohlcv: np.ndarray = shm.array
    times: np.ndarray = ohlcv['time']
    period: float = times[-1] - times[-2]
    if period == 0:
        # maybe just last sample is borked?
        period: float = times[-2] - times[-3]

    return period


def mk_ohlcv_shm_keyed_filepath(
    fqme: str,
    period: float | int,  # ow known as the "timeframe"
    datadir: Path,

) -> Path:

    if (
        period < 1.
        or
        int(period) != period
    ):
        raise ValueError(
            'Sample period must be positive whole seconds'
        )

    path: Path = (
        datadir
        /
        f'{fqme}.ohlcv{int(period)}s.parquet'
    )
    return path


def parse_ohlcv_parquet_path(
    path: Path,

) -> NativeSeriesRef|None:
    '''
    Parse one canonical NativeDB OHLCV filename.

    '''
    name: str = path.name
    if not name.endswith('.parquet'):
        return None

    stem: str = name.removesuffix('.parquet')
    fqme, marker, period_text = stem.rpartition('.ohlcv')
    if (
        marker != '.ohlcv'
        or
        not fqme
        or
        not period_text.endswith('s')
    ):
        return None

    digits: str = period_text.removesuffix('s')
    if (
        not digits.isascii()
        or
        not digits.isdecimal()
        or
        digits.startswith('0')
        or
        '..' in path.parts
    ):
        return None

    period_s: int = int(digits)
    if (
        fqme in {'.', '..'}
        or
        Path(fqme).name != fqme
        or
        '/' in fqme
        or
        '\\' in fqme
    ):
        return None

    expected: Path = mk_ohlcv_shm_keyed_filepath(
        fqme,
        period_s,
        path.parent,
    )
    if expected.name != name:
        return None

    return NativeSeriesRef(
        fqme=fqme,
        period_s=period_s,
        path=path,
    )


def iter_native_series(
    datadir: Path,
) -> Iterator[NativeSeriesRef]:
    '''
    Yield canonical NativeDB series in deterministic order.

    '''
    if not datadir.is_dir():
        return

    for path in sorted(
        datadir.iterdir(),
        key=lambda candidate: candidate.name,
    ):
        if (
            path.is_symlink()
            or
            not path.is_file()
        ):
            continue
        ref: NativeSeriesRef|None = parse_ohlcv_parquet_path(path)
        if ref is not None:
            yield ref


def unpack_fqme_from_parquet_filepath(path: Path) -> str:
    ref: NativeSeriesRef|None = parse_ohlcv_parquet_path(path)
    if ref is None:
        raise ValueError(f'Invalid NativeDB OHLCV path: {path}')
    return ref.fqme


# Only hydrate fields shared by canonical storage and provider SHM.
# Provider-only fields remain at their buffer defaults on restart.
ohlc_key_map = bidict({
    name: name
    for name, _ in def_iohlcv_fields
    if name != 'index'
})


class NativeStorageClient:
    '''
    High level storage api for OHLCV time series stored in
    a (modern) filesystem as apache parquet files B)

    Part of a grander scheme to use arrow and parquet as our main
    lowlevel data framework: https://arrow.apache.org/faq/.

    '''
    name: str = 'nativedb'

    def __init__(
        self,
        datadir: Path,

    ) -> None:
        self._datadir = datadir
        self._index: dict[tuple[str, int], NativeSeriesRef] = {}

        # series' cache from tsdb reads
        self._dfs: dict[str, dict[str, pl.DataFrame]] = {}

    @property
    def address(self) -> str:
        return self._datadir.as_uri()

    @property
    def cardinality(self) -> int:
        return len({fqme for fqme, _ in self._index})

    @property
    def series_cardinality(self) -> int:
        return len(self._index)

    # @property
    # def compression(self) -> str:
    #     ...

    async def list_keys(self) -> list[str]:
        return sorted({fqme for fqme, _ in self._index})

    async def list_series(self) -> list[NativeSeriesRef]:
        return list(self._index.values())

    def index_files(self):
        self._index.clear()
        for ref in iter_native_series(self._datadir):
            self._index[(ref.fqme, ref.period_s)] = ref

        return self._index


    # async def search_keys(self, pattern: str) -> list[str]:
    #     '''
    #     Search for time series key in the storage backend.

    #     '''
    #     ...

    # async def write_ticks(self, ticks: list) -> None:
    #     ...

    async def load(
        self,
        fqme: str,
        timeframe: int,

    ) -> tuple[
        np.ndarray,  # timeframe sampled array-series
        datetime | None,  # first dt
        datetime | None,  # last dt
    ] | None:
        try:
            array: np.ndarray = await self.read_ohlcv(
                fqme,
                timeframe,
            )
        except FileNotFoundError as fnfe:

            bs_fqme, _, *_ = fqme.rpartition('.')

            possible_matches: list[str] = []
            for tsfqme, _period in self._index:
                if bs_fqme in tsfqme:
                    possible_matches.append(tsfqme)

            match_str: str = '\n'.join(sorted(set(possible_matches)))
            raise TimeseriesNotFound(
                f'No entry for `{fqme}`?\n'
                f'Maybe you need a more specific fqme-key like:\n\n'
                f'{match_str}'
            ) from fnfe

        times = array['time']
        if array.size == 0:
            return None
        return (
            array,
            from_timestamp(times[0]),
            from_timestamp(times[-1]),
        )

    def mk_path(
        self,
        fqme: str,
        period: float,

    ) -> Path:
        return mk_ohlcv_shm_keyed_filepath(
            fqme=fqme,
            period=period,
            datadir=self._datadir,
        )

    def _cache_df(
        self,
        fqme: str,
        df: pl.DataFrame,
        timeframe: float,

    ) -> None:
        # cache df for later usage since we (currently) need to
        # convert to np.ndarrays to push to our `ShmArray` rt
        # buffers subsys but later we may operate entirely on
        # pyarrow arrays/buffers so keeping the dfs around for
        # a variety of purposes is handy.
        self._dfs.setdefault(
            timeframe,
            {},
        )[fqme] = df

    def _canonicalize_ohlcv(
        self,
        df: pl.DataFrame,

    ) -> pl.DataFrame:
        '''
        Normalize a provider frame to the durable OHLCV schema.

        Provider frames may omit the derived ``index`` field and add
        provider-only fields. Durable frames contain only canonical
        fields in their declared order.

        '''
        required: tuple[str, ...] = tuple(
            name
            for name, _ in def_iohlcv_fields
        )
        schema = {
            name: pl.Int64 if field_type is int else pl.Float64
            for name, field_type in def_iohlcv_fields
        }
        missing: set[str] = set(required).difference(df.columns)
        missing.discard('index')
        if missing:
            raise ValueError(
                f'OHLCV frame is missing columns: {sorted(missing)}'
            )

        times: np.ndarray = df['time'].to_numpy()
        if (
            not np.issubdtype(times.dtype, np.number)
            or
            not np.all(np.isfinite(times))
        ):
            raise ValueError(
                "OHLCV column 'time' must contain finite numeric data"
            )
        if np.any(times != np.floor(times)):
            raise ValueError(
                'OHLCV timestamps must use whole-second values'
            )

        return (
            df
            .with_columns(
                pl.Series('index', np.arange(df.height))
            )
            .select(required)
            .cast(schema)
        )

    def _repair_stored_ohlcv(
        self,
        df: pl.DataFrame,
        timeframe: int,

    ) -> pl.DataFrame:
        '''
        Drop known legacy sentinels before hydration or merge.

        '''
        stored: pl.DataFrame = self._canonicalize_ohlcv(df)
        zero_prices: pl.Expr = pl.all_horizontal([
            pl.col(field) == 0
            for field in _price_fields
        ]).fill_null(False)
        invalid: pl.Expr = (
            (pl.col('time') <= 0)
            |
            (
                (pl.col('time') <= timeframe)
                &
                zero_prices
            )
        )
        stored_len: int = stored.height
        stored = stored.filter(~invalid)
        dropped: int = stored_len - stored.height
        if dropped:
            log.warning(
                f'Dropping {dropped} invalid persisted OHLCV '
                f'row(s) during read repair'
            )
            stored = self._canonicalize_ohlcv(stored)
        return stored

    async def read_ohlcv(
        self,
        fqme: str,
        timeframe: int | str,
        end: float | None = None,  # epoch or none
        # limit: int = int(200e3),

    ) -> np.ndarray:
        path: Path = self.mk_path(
            fqme,
            period=int(timeframe),
        )
        df: pl.DataFrame = self._repair_stored_ohlcv(
            pl.read_parquet(path),
            timeframe=int(timeframe),
        )

        self._cache_df(
            fqme=fqme,
            df=df,
            timeframe=timeframe,
        )
        # TODO: filter by end and limit inputs
        # times: pl.Series = df['time']
        array: np.ndarray = tsp.pl2np(
            df,
            dtype=np.dtype(def_iohlcv_fields),
        )
        return array

    async def as_df(
        self,
        fqme: str,
        period: int = 60,
        load_from_offline: bool = True,

    ) -> pl.DataFrame:
        try:
            return self._dfs[period][fqme]
        except KeyError:
            if not load_from_offline:
                raise

            await self.read_ohlcv(fqme, period)
            return self._dfs[period][fqme]

    def _write_ohlcv(
        self,
        fqme: str,
        ohlcv: np.ndarray | pl.DataFrame,
        timeframe: int,

    ) -> Path:
        '''
        Sync version of the public interface meth, since we don't
        currently actually need or support an async impl.

        '''
        path: Path = mk_ohlcv_shm_keyed_filepath(
            fqme=fqme,
            period=timeframe,
            datadir=self._datadir,
        )
        if isinstance(ohlcv, np.ndarray):
            df: pl.DataFrame = tsp.np2pl(ohlcv)
        else:
            df = ohlcv
        df = self._canonicalize_ohlcv(df)
        self._validate_ohlcv(df, timeframe)

        # TODO: in terms of managing the ultra long term data
        # -[ ] use a proper profiler to measure all this IO and
        #   roundtripping!
        # -[ ] implement parquet append!? see issue:
        #   https://github.com/pikers/piker/issues/536
        #   -[ ] try out ``fastparquet``'s append writing:
        #     https://fastparquet.readthedocs.io/en/latest/api.html#fastparquet.write
        start = time.time()
        with NamedTemporaryFile(
            dir=path.parent,
            prefix=f'.{path.name}.',
            suffix='.tmp',
            delete=False,
        ) as tmp_file:
            tmp_path = Path(tmp_file.name)

        try:
            df.write_parquet(tmp_path)
            committed: pl.DataFrame = pl.read_parquet(tmp_path)
            self._validate_ohlcv(committed, timeframe)
            if not committed.equals(df):
                raise IOError(
                    'Temporary parquet differs from input frame'
                )

            with tmp_path.open('rb') as tmp_file:
                os.fsync(tmp_file.fileno())
            tmp_path.replace(path)

            self._cache_df(
                fqme=fqme,
                df=committed,
                timeframe=timeframe,
            )

            dir_fd: int = os.open(path.parent, os.O_RDONLY)
            try:
                os.fsync(dir_fd)
            finally:
                os.close(dir_fd)
        except BaseException:
            tmp_path.unlink(missing_ok=True)
            raise
        delay: float = round(
            time.time() - start,
            ndigits=6,
        )
        log.info(
            f'parquet write took {delay} secs\n'
            f'file path: {path}'
        )
        return path

    def _validate_ohlcv(
        self,
        df: pl.DataFrame,
        timeframe: int,

    ) -> None:
        '''
        Validate timestamp invariants before publishing a series.

        '''
        if df.is_empty():
            raise ValueError('Can not write an empty OHLCV series')

        required: tuple[str, ...] = tuple(
            name
            for name, _ in def_iohlcv_fields
        )
        missing: set[str] = set(required).difference(df.columns)
        if missing:
            raise ValueError(
                f'OHLCV frame is missing columns: {sorted(missing)}'
            )

        for field in required:
            values: np.ndarray = df[field].to_numpy()
            if (
                not np.issubdtype(values.dtype, np.number)
                or
                not np.all(np.isfinite(values))
            ):
                raise ValueError(
                    f'OHLCV column {field!r} must contain '
                    f'finite numeric data'
                )

        times: np.ndarray = df['time'].to_numpy()
        if (
            not np.all(np.isfinite(times))
            or
            np.any(times <= 0)
        ):
            raise ValueError(
                'OHLCV timestamps must be finite and positive'
            )

        prices: np.ndarray = df.select(_price_fields).to_numpy()
        epoch_sentinel: np.ndarray = (
            (times <= timeframe)
            &
            np.all(prices == 0, axis=1)
        )
        if np.any(epoch_sentinel):
            raise ValueError(
                'OHLCV frame contains a zero-initialized epoch sentinel'
            )

        if np.any(np.diff(times) <= 0):
            raise ValueError(
                'OHLCV timestamps must be strictly increasing'
            )

    async def update_ohlcv(
        self,
        fqme: str,
        ohlcv: np.ndarray | pl.DataFrame,
        timeframe: int,

    ) -> Path:
        '''
        Merge an incremental frame into durable OHLCV history.

        Incoming rows replace stored rows at matching timestamps.
        Writer ownership belongs to the persistent ``datad`` feed task.

        '''
        if isinstance(ohlcv, np.ndarray):
            incoming: pl.DataFrame = tsp.np2pl(ohlcv)
        else:
            incoming = ohlcv
        incoming = self._canonicalize_ohlcv(incoming)
        self._validate_ohlcv(incoming, timeframe)

        path: Path = self.mk_path(fqme, timeframe)
        if path.exists():
            stored: pl.DataFrame = self._repair_stored_ohlcv(
                pl.read_parquet(path),
                timeframe=timeframe,
            )
            merged: pl.DataFrame = pl.concat(
                [stored, incoming],
                how='diagonal_relaxed',
            )
            merged = (
                merged
                .unique(
                    subset='time',
                    keep='last',
                )
                .sort('time')
            )
        else:
            merged = incoming

        self._validate_ohlcv(merged, timeframe)
        return self._write_ohlcv(
            fqme,
            merged,
            timeframe,
        )

    async def write_ohlcv(
        self,
        fqme: str,
        ohlcv: np.ndarray | pl.DataFrame,
        timeframe: int,

    ) -> Path:
        '''
        Write input ohlcv time series for fqme and sampling period
        to (local) disk.

        '''
        return self._write_ohlcv(
            fqme,
            ohlcv,
            timeframe,
        )

    async def delete_ts(
        self,
        key: str,
        timeframe: int | None = None,

    ) -> bool:
        path: Path = mk_ohlcv_shm_keyed_filepath(
            fqme=key,
            period=timeframe,
            datadir=self._datadir,
        )
        if path.is_file():
            path.unlink()
            log.warning(f'Deleting parquet entry:\n{path}')
        else:
            log.error(f'No path exists:\n{path}')

        return path

    # TODO: allow wiping and refetching a segment of the OHLCV timeseries
    # data.
    # def clear_range(
    #     self,
    #     key: str,
    #     start_dt: datetime,
    #     end_dt: datetime,
    #     timeframe: int | None = None,
    # ) -> pl.DataFrame:
    #     '''
    #     Clear and re-fetch a range of datums for the OHLCV time series.

    #     Useful for series editing from a chart B)

    #     '''
    #     ...


# TODO: does this need to be async on average?
# I guess for any IPC connected backend yes?
@acm
async def get_client(

    # TODO: eventually support something something apache arrow
    # transport over ssh something..?
    # host: str | None = None,

    **kwargs,

) -> NativeStorageClient:
    '''
    Load a ``anyio_marketstore`` grpc client connected
    to an existing ``marketstore`` server.

    '''
    datadir: Path = config.get_conf_dir() / 'nativedb'
    if not datadir.is_dir():
        log.info(f'Creating `nativedb` dir: {datadir}')
        datadir.mkdir()

    client = NativeStorageClient(datadir)
    client.index_files()
    yield client
