"""
Utils for banks
"""

from __future__ import annotations

import contextlib
import os
import re
import sqlite3
import time
import warnings
from collections.abc import Sequence

import numpy as np
import obspy
import pandas as pd
from tables.exceptions import ClosedNodeError

from obsplus.constants import (
    LARGEDT64,
    NSLC,
    READ_HDF5_KWARGS,
    SMALLDT64,
    WAVEFORM_NAME_STRUCTURE,
    WAVEFORM_STRUCTURE,
)
from obsplus.exceptions import UnsupportedKeywordError
from obsplus.utils.geodetics import map_longitudes
from obsplus.utils.misc import READ_DICT, _get_path
from obsplus.utils.mseed import summarize_mseed
from obsplus.utils.time import _dict_times_to_ns, to_datetime64

# functions for summarizing the various formats
summarizing_functions = dict(mseed=summarize_mseed)

# extensions
WAVEFORM_EXT = ".mseed"
EVENT_EXT = ".xml"
STATION_EXT = ".xml"


# name structures


def _get_time_values(time1, time2=None):
    """Get the time values from a UTCDateTime object or two"""
    tvals = "year month day hour minute second microsecond".split()
    utc1 = time1
    split = re.split("-|:|T|[.]", str(utc1).replace("Z", ""))
    assert len(tvals) == len(split)
    out = {key: val for key, val in zip(tvals, split)}
    out["julday"] = "%03d" % utc1.julday
    out["starttime"] = utc1.timestamp
    if time2:
        out["endtime"] = time2.timestamp
    out["time"] = str(utc1).replace(":", "-").split(".")[0]
    return out


def summarize_generic_stream(path, format=None) -> list[dict]:
    """
    Return summary information for a stream stored on disk.

    Parameters
    ----------
    path
        The path to the stream.
    format
        The format code, as used by obspy.
    """
    st = _try_read_stream(path, format=format) or _try_read_stream(path) or []
    out = []
    for tr in st:
        summary = {
            "starttime": tr.stats.starttime._ns,
            "endtime": tr.stats.endtime._ns,
            "sampling_period": int(tr.stats.delta * 1_000_000_000),
            "path": path,
        }
        summary.update(dict((x, c) for x, c in zip(NSLC, tr.id.split("."))))
        out.append(summary)
    return out


def _summarize_wave_file(path, format, summarizer=None):
    """
    Summarize waveform files for indexing.

    Note: this function is a bit ugly, but it gets called *a lot* so be sure
    to profile before refactoring.
    """
    if summarizer is not None:
        try:
            return summarizer(path)
        except Exception:
            pass
    return summarize_generic_stream(path, format)


def _summarize_trace(
    trace: obspy.Trace,
    path: str | None = None,
    name: str | None = None,
    path_struct: str | None = None,
    name_struct: str | None = None,
) -> dict:
    """
    Function to extract info from traces for indexing.

    Parameters
    ----------
    trace
        The trace object
    path
        Other Parameters to the file
    name
        Name of the file
    path_struct
        directory structure to create
    name_struct
        structure of the file name if name is not used.
    """
    assert hasattr(trace, "stats"), "only a trace object is accepted"
    out = {"seedid": trace.id, "ext": WAVEFORM_EXT}
    t1, t2 = trace.stats.starttime, trace.stats.endtime
    out.update(_get_time_values(t1, t2))
    out.update(dict((x, c) for x, c in zip(NSLC, trace.id.split("."))))

    path_struct = path_struct if path_struct is not None else WAVEFORM_STRUCTURE
    name_struct = name_struct or WAVEFORM_NAME_STRUCTURE

    out.update(_get_path(out, path, name, path_struct, name_struct))
    return out


def _remove_base_path(series: pd.Series, base="") -> pd.Series:
    """
    Ensure paths stored in column name use unix style paths and have base
    path removed.
    """
    if series.empty:
        return series
    unix_paths = series.str.replace(os.sep, "/", regex=False)
    unix_base_path = str(base).replace(os.sep, "/")
    # NOTE: this uses "/" not os.sep because we've already swapped out os.sep
    return unix_paths.str.replace(unix_base_path + "/", "", regex=False)


def _natify_paths(series: pd.Series) -> pd.Series:
    """
    Natify paths in a series. IE, on windows replace / with \
    """
    if series.empty:
        return series
    return series.str.replace("/", os.sep)


class _IndexCache:
    """A simple class for caching indexes"""

    def __init__(self, bank, cache_size=5):
        self.max_size = cache_size
        self.bank = bank
        self.cache = pd.DataFrame(
            index=range(cache_size), columns="t1 t2 kwargs cindex".split()
        )
        self._current_index = 0
        # self.next_index = itertools.cycle(self.cache.index)

    def __call__(self, starttime, endtime, buffer, **kwargs):
        """Get start and end times, perform in kernel lookup"""
        starttime, endtime = self._get_times(starttime, endtime)
        self._validate_kwargs(kwargs)
        # find out if the query falls within one cached times
        con1 = self.cache.t1 <= starttime
        con2 = self.cache.t2 >= endtime
        con3 = self.cache.kwargs == self._kwargs_to_str(kwargs)
        cached_index = self.cache[con1 & con2 & con3]
        if not len(cached_index):  # query is not cached get it from hdf5 file
            where = _get_kernel_query(
                starttime.astype(np.int64), endtime.astype(np.int64), int(buffer)
            )
            raw_index = self._get_index(where, **kwargs)
            # replace "None" with None
            ic = self.bank.index_str
            raw_index.loc[:, ic] = raw_index.loc[:, ic].replace(["None"], [None])
            # convert data types used by bank back to those seen by user
            index = raw_index.astype(dict(self.bank._dtypes_output))
            self._set_cache(index, starttime, endtime, kwargs)
        else:
            index = cached_index.iloc[0]["cindex"]
        # trim down index
        con1 = index["starttime"] >= (endtime + buffer)
        con2 = index["endtime"] <= (starttime - buffer)
        return index[~(con1 | con2)]

    @staticmethod
    def _get_times(starttime, endtime):
        """Return starttimes and endtimes."""
        # get defaults if starttime or endtime is none
        starttime = None if pd.isnull(starttime) else starttime
        endtime = None if pd.isnull(endtime) else endtime
        starttime = to_datetime64(starttime or SMALLDT64)
        endtime = to_datetime64(endtime or LARGEDT64)
        if starttime is not None and endtime is not None:
            if starttime > endtime:
                msg = "starttime cannot be greater than endtime."
                raise ValueError(msg)
        return starttime, endtime

    def _validate_kwargs(self, kwargs):
        """Ensure kwargs are supported."""
        kwarg_set = set(kwargs)
        if not kwarg_set.issubset(READ_HDF5_KWARGS):
            bad_kwargs = kwarg_set - set(READ_HDF5_KWARGS)
            msg = f"The following kwargs are not supported: {bad_kwargs}. "
            raise UnsupportedKeywordError(msg)

    def _set_cache(self, index, starttime, endtime, kwargs):
        """Cache the current index"""
        ser = pd.Series(
            {
                "t1": starttime,
                "t2": endtime,
                "cindex": index,
                "kwargs": self._kwargs_to_str(kwargs),
            }
        )
        self.cache.loc[self._get_next_index()] = ser

    def _kwargs_to_str(self, kwargs):
        """Convert kwargs to a string"""
        keys = sorted(list(kwargs.keys()))
        ou = str([(item, kwargs[item]) for item in keys])
        return ou

    def _get_index(self, where, fail_counts=0, **kwargs):
        """Read the hdf5 file"""
        try:
            return pd.read_hdf(
                self.bank.index_path, self.bank._index_node, where=where, **kwargs
            )

        except (ClosedNodeError, Exception) as e:
            # Sometimes in concurrent updates the nodes need time to open/close
            if fail_counts > 10:
                raise e
            # Wait a bit and try again (up to 10 times)
            time.sleep(0.1)
            return self._get_index(where, fail_counts=fail_counts + 1, **kwargs)

    def clear_cache(self):
        """Removes all cached dataframes."""
        self.cache = pd.DataFrame(
            index=range(self.max_size), columns="t1 t2 kwargs cindex".split()
        )

    def _get_next_index(self):
        """
        Get the next index value on cache.

        Note we can't use itertools.cycle here because it cant be pickled.
        """
        if self._current_index == len(self.cache.index) - 1:
            self._current_index = 0
        else:
            self._current_index += 1
        return self.cache.index[self._current_index]


@contextlib.contextmanager
def sql_connection(path, **kwargs):
    """
    A simple context manager for a sqlite connection.

    Parameters
    ----------
    path
        A path to the sqlite database.
    """
    con = sqlite3.connect(str(path), **kwargs)
    with con:
        yield con
    con.close()  # this is needed on windows but not linux, weird...


def _get_kernel_query(starttime: int, endtime: int, buffer: int):
    """
    Create a HDF5 kernel query based on start and end times.

    This is necessary because hdf5 doesnt accept inverted conditions.
    A slight buffer is applied to the ranges to make sure no edge files
    are excluded.
    """
    t1 = starttime - buffer
    t2 = endtime + buffer
    con = (
        f"(starttime>{t1:d} & starttime<{t2:d}) | "
        f"((endtime>{t1:d} & endtime<{t2:d}) | "
        f"(starttime<{t1:d} & endtime>{t2:d}))"
    )
    return con


# --- SQL stuff


def _str_of_params(value):
    """
    Make sure a list of params is returned.

    This allows user to specify a single parameter, a list, set, nparray, etc.
    to match on.
    """
    if isinstance(value, str):
        return value
    else:
        # try to coerce in a list of str
        try:
            return [str(x) for x in value]
        except TypeError:  # else fallback to str repr
            return str(value)


def _make_wheres(queries):
    """Create the where queries, join with AND clauses"""

    def _rename_keys(kwargs):
        """re-word some keys to make automatic sql generation easier"""
        if "eventid" in kwargs:
            kwargs["event_id"] = kwargs.pop("eventid")
        if "event_id" in kwargs:
            kwargs["event_id"] = _str_of_params(kwargs["event_id"])
        if "event_description" in kwargs:
            kwargs["event_description"] = _str_of_params(kwargs["event_description"])
        if "endtime" in kwargs:
            kwargs["maxtime"] = kwargs.pop("endtime")
        if "starttime" in kwargs:
            kwargs["mintime"] = kwargs.pop("starttime")
        return kwargs

    def _handle_nat(kwargs):
        """Add a mintime that will exclude NaT values if endtime is used"""
        if "maxtime" in kwargs and "mintime" not in kwargs:
            kwargs["mintime"] = SMALLDT64.astype(np.int64) + 1
        return kwargs

    def _handle_dateline_transversal(kwargs, out):
        """Check if dateline should be transversed by query."""
        # if longitudes aren't being used bail out
        if not {"minlongitude", "maxlongitude"}.issubset(set(kwargs)):
            return kwargs, out
        # if dateline is not to be transversed by query bail out
        long_array = np.array([kwargs["minlongitude"], kwargs["maxlongitude"]])
        minlong, maxlong = map_longitudes(long_array)
        if not minlong > maxlong:
            return kwargs, out
        # remove min/max long from query dict and reform to two queries.
        kwargs.pop("minlongitude"), kwargs.pop("maxlongitude")
        cond = f"(( longitude > {minlong}) OR ( longitude < {maxlong})) "
        out.append(cond)
        return kwargs, out

    def _build_query(kwargs):
        """Iterate each key/value and build query"""
        out = []
        kwargs, out = _handle_dateline_transversal(kwargs, out)
        for key, val in kwargs.items():
            # deal with simple min/max
            if key.startswith("min"):
                out.append(f"{key.replace('min', '')} > {val}")
            elif key.startswith("max"):
                out.append(f"{key.replace('max', '')} < {val}")
            # deal with equals or ins
            elif isinstance(val, Sequence):
                if isinstance(val, str):
                    val = [val]
                tup = str(tuple(val)).replace(",)", ")")  # no trailing comma
                out.append(f"{key} IN {tup}")
            else:
                out.append(f"{key} = {val}")
        return " AND ".join(out).replace("'", '"')

    kwargs = _dict_times_to_ns(queries)
    kwargs = _rename_keys(kwargs)
    kwargs = _handle_nat(kwargs)
    return _build_query(kwargs)


def _make_sql_command(cmd, table_name, columns=None, **kwargs) -> str:
    """Build a sql command"""
    # get columns
    if columns:
        col = [columns] if isinstance(columns, str) else columns
        col += ["event_id"]  # event_id is used as index
        columns = ", ".join(col)
    elif cmd.upper() == "DELETE":
        columns = ""
    else:
        columns = "*"
    limit = kwargs.pop("limit", None)
    wheres = _make_wheres(kwargs)
    sql = f'{cmd.upper()} {columns} FROM "{table_name}"'
    if wheres:
        sql += f" WHERE {wheres}"
    if limit:
        sql += f" LIMIT {limit}"
    return sql + ";"


def _read_table(table_name, con, columns=None, **kwargs) -> pd.DataFrame:
    """
    Read a SQLite table.

    Parameters
    ----------
    table_name
    con
    columns
    """
    # first ensure all times are ns (as ints)
    sql = _make_sql_command("select", table_name, columns=columns, **kwargs)
    # replace "None" with None
    return pd.read_sql(sql, con)


def _get_tables(con):
    """Return a list of table in sqlite database"""
    out = con.execute("SELECT name FROM sqlite_master WHERE type='table';")
    return set(out)


def _drop_rows(table_name, con, columns=None, **kwargs):
    """Drop indicies in table"""
    sql = _make_sql_command("delete", table_name, columns=columns, **kwargs)
    con.execute(sql)


def _try_read_stream(stream_path, format=None, **kwargs):
    """Try to read a waveforms from file, if raises return None"""
    read = READ_DICT.get(format, obspy.read)
    stt = None
    try:
        stt = read(stream_path, **kwargs)
    except Exception:
        try:
            stt = obspy.read(stream_path, **kwargs)
        except Exception:
            warnings.warn(f"obspy failed to read {stream_path}", UserWarning)
        else:
            msg = f"{stream_path} was read but is not of format {format}"
            warnings.warn(msg, UserWarning)
    finally:
        return stt if stt else None
