from datetime import date, time, datetime, timedelta
import os
from inspect import signature, Parameter
from pathlib import Path
import pprint
from string import Formatter
import sys
import traceback

import cx_Oracle
import pandas as pd
from sa_gwdata import *

from .utils import (
    SQL,
    apply_well_id,
    cleanup_columns,
    chunk,
    get_logger,
    parse_query_metadata,
)
from .config import *

from sageodata_db.utils import get_logger


logger = get_logger()

for test_path in (r"C:\Program Files\Oracle\Instant Client v19.9",):
    if Path(test_path).is_dir():
        try:
            cx_Oracle.init_oracle_client(test_path)
        except cx_Oracle.ProgrammingError:
            logger.info(
                "Ignoring cx_Oracle.ProgrammingError - already initialised likely"
            )


def load_predefined_query(filename, stub="queries"):
    if not Path(filename).is_file():
        filename = Path(__file__).parent / stub / filename
    if not Path(filename).is_file():
        return load_predefined_query(str(filename) + ".sql")
    with open(str(filename)) as f:
        query = f.read()
    return query


def get_predefined_query_filenames():
    return [p.name for p in (Path(__file__).parent / "queries").glob("*.sql")]


def connect(user="gwquery", password="gwquery", **kwargs):
    """Connect to SA Geodata.

    Args:
        user (str): oracle user
        password (str): password
        service_name (str): version of SA Geodata you would like to connect
            to - options are "DMED.WORLD" or "dev"; "DMET.WORLD" or "test" or
            "QA"; or "DMEP.WORLD" or "prod" - see
            :func:`sageodata_db.config.normalize_service_name` for details.

    Other keyword arguments are passed to
    :func:`sageodata_db.config.make_connection_string`.

    Returns:
        a :class:`sageodata_db.SAGeodataConnection` object.

    Example:

        >>> from sageodata_db import connect
        >>> db = connect()
        >>> db
        <sageodata_db.connection.SAGeodataConnection to gwquery@pirsapd07.pirsa.sa.gov.au:1521/DMEP.World>

    """
    connect_string = make_connection_string(**kwargs)
    return SAGeodataConnection(user=user, password=password, dsn=connect_string)


class SAGeodataConnection(cx_Oracle.Connection):
    """:class:`cx_Oracle.Connection` with methods to read groundwater data.

    This should not be instantiated directly. Instead use::

        >>> from sageodata_db import connect
        >>> db = connect()
        >>> type(db)
        <class 'sageodata_db.connection.SAGeodataConnection'>

    See :ref:`predefined-queries-label` for more information on accessing
    data using predefined SQL queries.

    """

    SQL = SQL

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        from sageodata_db import __version__

        self.__version__ = __version__
        self._all_replacement_drillholes = self.all_replacement_drillholes()

    def _make_predefined_query_lambda(self, query_name):
        return lambda *args, **kwargs: self._predefined_query(
            query_name, *args, **kwargs
        )

    def __repr__(self):
        return super().__repr__()

    def test_alive(self):
        """Test whether connection is functioning.

        Returns:
            `True` or `False`

        """
        df = self.query("select distinct owner from all_tables")
        try:
            owners = df["owner"]
        except cx_Oracle.DatabaseError:
            return False
        if "SYS" in owners.values:
            return True
        return False

    def _predefined_query(
        self,
        filename,
        *args,
        stub="queries",
        query_method="query",
        include_replacements=False,
        **kwargs,
    ):
        """Run a predefined query on the database.

        Args:
            query_name (str): see below for possible predefined queries.

        Other args and kwargs are passed to `_predefined_query()` along with the fixed
        kwargs `stub="queries", query_method="query"`.

        Returns:
            pandas DataFrame.

        """
        sql = load_predefined_query(filename, stub=stub)
        query_method = getattr(self, query_method)
        df = query_method(SQL(sql, *args, **kwargs))
        cols = [x for x in df.columns]
        if (
            "obs_no" in df.columns
            and "unit_hyphen" in df.columns
            and "dh_no" in df.columns
        ):
            if len(df):
                df["well_id"] = df.apply(apply_well_id, axis="columns")
            else:
                df = df.assign(well_id=None)
            cols = ["well_id"] + cols
        df = df[cols]
        logger.debug(f"Include replacements? {include_replacements}")
        if include_replacements:
            dfs = []
            for dh_no in df.dh_no.unique():
                all_dhs = self.find_replacements([dh_no])
                if len(all_dhs) > 0:
                    missing_dhs = set(all_dhs.old_dh_no).difference(set(df.dh_no))
                    logger.debug(f"missing_dhs {missing_dhs}")
                    all_df = self._predefined_query(
                        filename,
                        list(missing_dhs),
                        include_replacements=False,
                        **kwargs,
                    )
                    df2 = pd.concat([df, all_df])
                    if "obs_date" in df2.columns:
                        dt_col = "obs_date"
                    elif "collected_date" in df2.columns:
                        dt_col = "collected_date"

                    df3 = df2[df2.dh_no == all_dhs.iloc[0].old_dh_no]

                    for idx, row in all_dhs.sort_values("replaced_from").iterrows():
                        logger.debug(
                            f"Adding in data from replacement: {row.to_dict()}"
                        )
                        if ~pd.isnull(row.new_dh_no) and ~pd.isnull(row.replaced_from):
                            # Remove measurements from the old_dh_no after or on the replacement date.
                            remove_old = (df3.dh_no == row.old_dh_no) & (
                                df3[dt_col] >= row.replaced_from
                            )
                            df3 = df3[~remove_old]
                            # Add measurements from the new_dh_no on and after the replacement date.
                            keep_new = (df2.dh_no == row.new_dh_no) & (
                                df2[dt_col] >= row.replaced_from
                            )
                            df3 = pd.concat([df3, df2[keep_new]])
                    df3 = df3.sort_values([dt_col])

                    # Improve columns
                    current = (
                        self.drillhole_details(all_dhs.index.unique()).iloc[0].to_dict()
                    )
                    logger.debug(f"Current well is: {current}")
                    well_cols = [
                        "well_id",
                        "dh_no",
                        "unit_long",
                        "unit_hyphen",
                        "obs_no",
                        "dh_name",
                        "easting",
                        "northing",
                        "zone",
                        "latitude",
                        "longitude",
                        "aquifer",
                    ]
                    for col in ["dh_no", "unit_long", "unit_hyphen", "obs_no"]:
                        df3[f"from_{col}"] = df3[col]
                    for col in well_cols[::-1]:
                        df3[col] = current[col]
                else:
                    df3 = df[df.dh_no == dh_no]
                dfs.append(df3)
            df = pd.concat(dfs)
        return df

    def query(self, query, lowercase_columns=True, *args, conn=None, **kwargs):
        """Run an SQL query on SA Geodata and return a :class:`pandas.DataFrame`.

        Args:
            query (str, SQL): SQL query as string or
                :class:`sageodata_db.connection.SQL`

        Returns:
            a :class:`pandas.DataFrame` object.

        """
        if conn is None:
            conn = self
        if isinstance(query, str):
            query = SQL(query, *args, **kwargs)
        dfs = []
        for sql in query:
            logger.debug("running query:\n{}".format(sql))
            dfs.append(pd.read_sql_query(sql, conn))
        df = pd.concat(dfs)
        if lowercase_columns:
            df = df.rename(columns=str.lower)
        return df

    def _create_well_instances(self, dh_nos):
        """Private method to create Well instances from dh_nos.

        Args:
            dh_nos (array-like): drillhole numbers as integers.

        Returns:
            :class:`sa_gwdata.Wells` object (sequence of :class:`sa_gwdata.Well` objects).

        """
        df = self.drillhole_details(dh_nos)
        df.loc[:, "name"] = df["dh_name"]
        df.loc[:, "unit_no.hyphen"] = df["unit_hyphen"]
        df.loc[:, "unit_no"] = df["unit_hyphen"]
        df.loc[:, "dh_no"] = df["dh_no"]
        return Wells([Well(**vals.to_dict()) for _, vals in df.iterrows()])

    def find_replacement_history(self, dh_no):
        """Find the chain of drillholes that replaced one another, for *dh_no*.

        Args:
            dh_no (int): drillhole number

        Returns:
            a pandas dataframe with columns "old_dh_no", "new_dh_no", and
            "replaced_from".

        """

        df = self._all_replacement_drillholes

        def _find_replacement_history(
            dh_no_chain, check_for_newer_wells=True, check_for_older_wells=True
        ):
            earliest_dh_no = dh_no_chain[0]
            latest_dh_no = dh_no_chain[-1]
            if check_for_newer_wells:
                replacements = df[df["dh_no"] == latest_dh_no]
                if len(replacements) == 1:
                    row = replacements.iloc[0]
                    dh_no_chain = list(dh_no_chain) + [row.new_dh_no]
                    return _find_replacement_history(
                        dh_no_chain, check_for_newer_wells=True
                    )
            if check_for_older_wells:
                replaced = df[df["new_dh_no"] == earliest_dh_no]
                if len(replaced) == 0:
                    return dh_no_chain
                elif len(replaced) == 1:
                    row = replaced.iloc[0]
                    previous_dh_no = row.dh_no
                    dh_no_chain = [previous_dh_no] + list(dh_no_chain)
                    return _find_replacement_history(
                        dh_no_chain, check_for_newer_wells=False
                    )

        dh_no_chain = _find_replacement_history([dh_no])
        hdf = df[df.dh_no.isin(dh_no_chain)]
        hdf = hdf.set_index("dh_no")
        hdf2 = hdf.reindex(dh_no_chain)
        hdf3 = hdf2.reset_index().rename(columns={"dh_no": "old_dh_no"})
        return hdf3

    def find_replacements(self, wells):
        """Find replacements for a set of wells.

        Args:
            wells (Wells, pandas dataframe, or list of ints): either must
                have an attribute "dh_no", or a list of drillhole numbers.

        Returns:
            pandas DataFrame with index "current_dh_no" and
            columns "queried_dh_no", "old_dh_no", "new_dh_no", "replaced_from"

        """
        if hasattr(wells, "dh_no"):
            dh_nos = [getattr(o, "dh_no") for o in wells]
        else:
            dh_nos = wells
        dfs = []
        for dh_no in dh_nos:
            df = self.find_replacement_history(dh_no)
            if len(df) > 1:
                df.insert(0, "current_dh_no", _find_current_for_replacement_chain(df))
                df.insert(0, "queried_dh_no", dh_no)
            else:
                df = df.assign(current_dh_no=dh_no, queried_dh_no=dh_no)
            dfs.append(df)
        return pd.concat(dfs).set_index("current_dh_no")

    def find_wells(self, input_text, **kwargs):
        """Find wells and retrieve some summary information.

        Args:
            input_text (str): any well identifiers to parse. See
                :func:`sa_gwdata.parse_well_ids_plaintext` for details of
                other keyword arguments you can pass here.

        Returns:
            a list of :class:`sa_gwdata.Well` objects.

        For example::

            >>> from sageodata_db import connect
            >>> db = connect()
            >>> wells = db.find_wells(": G662801265, 6628-14328, YAT30, and ULE 205.")
            >>> wells
            ['6628-1265', 'YAT030', 'MOR201', 'ULE205']

        See :func:`sa_gwdata.parse_well_ids` for more details.

        """
        ids = parse_well_ids(input_text, **kwargs)
        dh_nos = [x for id_type, x in ids if id_type == "dh_no"]
        unit_nos = [x for id_type, x in ids if id_type == "unit_no"]
        obs_nos = [x for id_type, x in ids if id_type == "obs_no"]
        unit_long_ints = [UnitNumber(x).long_int for x in unit_nos]
        logger.debug("unit_no.long_int: {}".format(unit_long_ints))
        r_obs_nos = []

        if obs_nos:
            r1 = self._predefined_query("drillhole_no_by_obs_no.sql", obs_nos)
            r_obs_nos = r1.dh_no.tolist()
        r_unit_nos = []
        if unit_long_ints:
            r2 = self._predefined_query("drillhole_no_by_unit_long.sql", unit_long_ints)
            r_unit_nos = r2.dh_no.tolist()
        logger.debug("obs_nos -> {}".format(r_obs_nos))
        logger.debug("unit_nos -> {}".format(r_unit_nos))
        all_dh_nos = list(set(dh_nos + r_obs_nos + r_unit_nos))
        return self._create_well_instances(all_dh_nos)

    def find_wells_from_df(
        self,
        df,
        return_id_cols=("unit_hyphen",),
        remap_id_col_names=True,
        **kwargs,
    ):
        """Find wells identifier in a dataframe and merge them together.

        Args:
            df (pd.DataFrame): dataframe containing well IDs
            return_id_cols (list of str): columns containing identified well
                IDs which are returned. If *remap_id_col_names* is True, you can
                pick from: dh_no, well_id, dh_name, unit_hyphen, unit_long, unit_wilma
                unit_hydstra, obs_no, and obs_egis. If *remap_id_col_names* is False,
                you can choose from dh_no, id, name, unit_no.map, unit_no.seq,
                unit_no.hyphen, unit_no.long_int, unit_no.long,
                unit_no.wilma, unit_no.hydstra, obs_no.plan, obs_no.seq, obs_no.id,
                and obs_no.egis.
            remap_id_col_names (bool): remap from Wells.df to the typical columns
                found in other data queries.

        Other kwargs are passed to parse_well_ids().

        Returns:
            a copy of *df* with *return_id_cols* added at the start, with all
            well identifiers matched.

        """
        if isinstance(df, pd.Series):
            df = df.to_frame()

        # Parse any well identifier for each row of the
        # dataframe.
        wids = df.apply(
            lambda series: parse_well_ids("|".join(map(str, series.values)), **kwargs),
            axis="columns",
        )

        # Use only the first of each well identifier in each row
        wids = [wid[0] if len(wid) else "" for wid in wids]

        search_lines = list(set([str(wid[1]) if len(wid) == 2 else "" for wid in wids]))
        vwids_dfs = []
        for search in chunk(search_lines, 1000):
            logger.info(f"Searching for {len(search)} well IDs")
            vwids_dfs.append(self.find_wells("/".join(search), **kwargs).df())
        vwids = pd.concat(vwids_dfs).drop_duplicates()
        values = []
        for wid in wids:
            row = False
            if wid == "":
                pass
            elif wid[0] == "unit_no":
                row_sel = vwids[vwids["unit_no.hyphen"] == wid[1]]
                if len(row_sel) == 1:
                    row = row_sel.iloc[0]
            elif wid[0] == "obs_no":
                row_sel = vwids[vwids["obs_no.id"] == wid[1]]
                if len(row_sel) == 1:
                    row = row_sel.iloc[0]
            elif wid[0] == "dh_no":
                row_sel = vwids[vwids.dh_no == wid[1]]
                if len(row_sel) == 1:
                    row = row_sel.iloc[0]
            if row is False:
                row = pd.Series([None for x in vwids.columns], index=vwids.columns)
            row.loc["parsed_id"] = wid
            values.append(row)
        id_cols = pd.DataFrame(values, index=df.index.values)
        if remap_id_col_names:
            id_cols = id_cols.rename(
                columns={
                    "id": "well_id",
                    "name": "dh_name",
                    "unit_no.hyphen": "unit_hyphen",
                    "unit_no.long_int": "unit_long",
                    "unit_no.wilma": "unit_wilma",
                    "unit_no.hydstra": "unit_hydstra",
                    "obs_no.id": "obs_no",
                    "obs_no.egis": "obs_egis",
                }
            )
        if return_id_cols is None:
            return_id_cols = id_cols.columns
        else:
            return_id_cols = ["parsed_id"] + list(return_id_cols)
        return id_cols[return_id_cols].T.drop_duplicates().T

    def find_edits_by(self, *modified_by):
        """Return a table of edits made to certain tables by a set of users.

        The tables which are current queried are water levels, salinity samples,
        and elevation surveys.

        Args:
            *modified_by (str): SA Geodata usernames

        Returns:
            pandas DataFrame

        """
        wls = self.water_level_edits(modified_by)
        sals = self.salinity_edits(modified_by)
        elevs = self.elevation_edits(modified_by)
        wls = _remove_elev_mods_from_wl(wls, elevs)
        wls = wls.groupby(["well_id", "modified_date"]).filter(lambda x: len(x) <= 1)
        df = pd.concat(
            [
                wls.assign(table="WL").rename(columns={"wl": "data_values"}),
                sals.assign(table="Salinity").rename(
                    columns={"salinity": "data_values"}
                ),
                elevs.assign(table="Elevation").rename(columns={"elev": "data_values"}),
            ]
        ).sort_values("modified_date", ascending=False)
        cols = [x for x in df.columns if not x == "table"]
        cols.insert(cols.index("well_id") + 1, "table")
        return df[cols]

    def find_additions_by(self, *created_by):
        """Return a table of additions made to certain tables by a set of users.

        The tables which are current queried are water levels, salinity samples,
        and elevation surveys.

        Args:
            *created_by (str): SA Geodata usernames

        Returns:
            pandas DataFrame.

        """
        wls = self.water_level_additions(created_by)
        sals = self.salinity_additions(created_by)
        elevs = self.elevation_additions(created_by)
        wls = wls.groupby(["well_id", "creation_date"]).filter(lambda x: len(x) <= 1)
        df = pd.concat(
            [
                wls.assign(table="WL").rename(columns={"wl": "data_values"}),
                sals.assign(table="Salinity").rename(
                    columns={"salinity": "data_values"}
                ),
                elevs.assign(table="Elevation").rename(columns={"elev": "data_values"}),
            ]
        ).sort_values("creation_date", ascending=False)
        cols = [x for x in df.columns if not x == "table"]
        cols.insert(cols.index("well_id") + 1, "table")
        return df[cols]


def _remove_elev_mods_from_wl(wl, elev):
    """Remove any edit records from WL data when there is an edit made to
    the elevation table within 5 seconds.

    """
    remove_indices = []
    wl_copy = pd.DataFrame(wl)
    for (dh_no, modified_date), gdf in elev.groupby(["dh_no", "modified_date"]):
        n1 = len(wl[wl.dh_no == dh_no])
        date1 = modified_date - timedelta(seconds=2)
        date2 = modified_date + timedelta(seconds=2)
        remove_idx = (
            (wl.dh_no == dh_no)
            & (wl.modified_date >= date1)
            & (wl.modified_date <= date2)
        )
        n2 = len(wl[remove_idx])
        logger.debug(f"{dh_no}, {modified_date}: wl {n1} -> {n2}")
        remove_indices.append(remove_idx)
        wl_copy = wl_copy[~remove_idx]
    return wl_copy


def _find_current_for_replacement_chain(df):
    """Identify which drillhole is current in a chain of
    replacement drillholes.

    Args:
        df (pandas DataFrame): returned by
            SAGeodataConnection.find_replacement_history()

    Returns:
        int, the drillhole number of the current well.

    """
    candidate = df[df.new_dh_no.isnull()].iloc[0]
    candidate_start = df.loc[df.new_dh_no == candidate.old_dh_no, "replaced_from"].iloc[
        0
    ]
    other_candidates = df.loc[df.replaced_from > candidate_start, :]
    if len(other_candidates) > 0:
        return int(other_candidates.sort_values("replaced_from").iloc[-1].new_dh_no)
    else:
        return int(candidate.old_dh_no)


def __monkey_patch_predefined_queries(cls, method_name):
    """Monkey-patch predefined query methods on to a class.

    Args:
        cls (class): class to monkey-patch methods on to
        method_name (str): name of method in cls which takes
            the query_name as its first argument. Other args
            and kwargs are passed through to this method

    Returns:
        nothing.

    """

    def __make_lambda(query_name, method_name):
        return lambda self, *args, **kwargs: getattr(self, method_name)(
            query_name, *args, **kwargs
        )

    docstring_lines = []
    query_names = get_predefined_query_filenames()
    for query_name in query_names:
        query = load_predefined_query(query_name)
        columns, param_defns, docstring_parts = parse_query_metadata(query)
        query_name = query_name.replace(".sql", "")
        func = __make_lambda(query_name, method_name=method_name)
        func.__name__ = query_name
        field_names = [fname for _, fname, _, _ in Formatter().parse(query) if fname]
        if field_names:
            params = [Parameter("self", Parameter.POSITIONAL_OR_KEYWORD)]
            for field_name in field_names:
                param_meta = param_defns.get(field_name, {"default": None, "descr": ""})
                if not param_meta["default"]:
                    param_meta["default"] = Parameter.empty
                params.append(
                    Parameter(
                        field_name,
                        Parameter.POSITIONAL_OR_KEYWORD,
                        default=param_meta["default"],
                    )
                )
            if "dh_no" in [n.lower() for n in field_names]:
                params.append(
                    Parameter(
                        "include_replacements",
                        Parameter.POSITIONAL_OR_KEYWORD,
                        default=False,
                    )
                )
                param_defns["include_replacements"] = {
                    "descr": (
                        "include_replacements (bool): default ``False``. "
                        "Fold data from replaced drillholes the newest "
                        "drillhole (e.g. data from dh_no 12345 will be included"
                        " as part of drillhole 23500 if 23500 is marked as a "
                        "replacement for 12345)"
                    ),
                    "default": False,
                }

            unique_params = []
            for p in params:
                if not p in unique_params:
                    unique_params.append(p)
            func.__signature__ = signature(func).replace(parameters=unique_params)

        doclines = [f"Run predefined query '{query_name}.sql'.", ""]
        if docstring_parts[0]:
            doclines += docstring_parts[0]
            doclines += [""]
        if param_defns:
            doclines += ["Args:"]
            doclines += [
                f"    {meta['descr'].strip()}" for name, meta in param_defns.items()
            ]
            doclines += [""]
        if len(columns):
            doclines += ["Returns:", "    :class:`pandas.DataFrame`", ""]
            doclines += ["The returned DataFrame has the following columns:", ""]
            doclines += [
                f" - ``'{name}'`` {meta['descr']}" for name, meta in columns.items()
            ]
            doclines += [""]
        if docstring_parts[1]:
            doclines += docstring_parts[1]
            doclines += [""]
        # doclines += [
        #     "The SQL template (used by :class:`sageodata_db.SQL`) for this query:",
        #     "",
        #     ".. code-block:: plpgsql",
        #     "",
        # ]
        # doclines += [
        #     f"    {line}"
        #     for line in query.splitlines()
        #     if not line.strip().startswith("-- ")
        # ]
        # doclines += [""]

        func.__doc__ = "\n".join([f"    {line}" for line in doclines])
        setattr(cls, query_name, func)


__monkey_patch_predefined_queries(SAGeodataConnection, method_name="_predefined_query")
