##!/usr/bin/env python
# crate_anon/nlp_manager/input_field_config.py

import logging
from typing import Any, Dict, Iterator, List, Optional, Tuple

from cardinal_pythonlib.rnc_db import (
    ensure_valid_field_name,
    ensure_valid_table_name,
)
from sqlalchemy import BigInteger, Column, Index, String, Table
from sqlalchemy.sql import column, func, select, table
# from sqlalchemy.sql.elements import quoted_name

from crate_anon.nlp_manager.constants import MAX_STRING_PK_LENGTH
from crate_anon.common.hash import hash64
from crate_anon.common.sqla import (
    is_sqlatype_integer,
    get_column_type,
    table_exists,
)
from crate_anon.nlp_manager.constants import SqlTypeDbIdentifier
from crate_anon.nlp_manager.models import NlpRecord
from crate_anon.nlp_manager.nlp_definition import NlpDefinition

log = logging.getLogger(__name__)

FN_SRCDB = '_srcdb'
FN_SRCTABLE = '_srctable'
FN_SRCPKFIELD = '_srcpkfield'
FN_SRCPKVAL = '_srcpkval'
FN_SRCPKSTR = '_srcpkstr'
FN_SRCFIELD = '_srcfield'


# =============================================================================
# Input field definition
# =============================================================================

class InputFieldConfig(object):
    """
    Class defining configuration for an input field (containing text).
    """

    def __init__(self, nlpdef: NlpDefinition, section: str) -> None:
        """
        Read config from a configparser section.
        """
        def opt_str(option: str) -> str:
            return nlpdef.opt_str(section, option, required=True)

        def opt_strlist(option: str,
                        required: bool = False,
                        lower: bool = True) -> List[str]:
            return nlpdef.opt_strlist(section, option, as_words=False,
                                      lower=lower, required=required)

        self._nlpdef = nlpdef

        self._srcdb = opt_str('srcdb')
        self._srctable = opt_str('srctable')
        self._srcpkfield = opt_str('srcpkfield')
        self._srcfield = opt_str('srcfield')
        # Make these case-sensitive to avoid our failure in renaming SQLA
        # Column objects to be lower-case:
        self._copyfields = opt_strlist('copyfields', lower=False)  # fieldnames
        self._indexed_copyfields = opt_strlist('indexed_copyfields',
                                               lower=False)

        ensure_valid_table_name(self._srctable)
        ensure_valid_field_name(self._srcpkfield)
        ensure_valid_field_name(self._srcfield)

        if len(set(self._indexed_copyfields)) != len(self._indexed_copyfields):
            raise ValueError("Redundant indexed_copyfields: {}".format(
                self._indexed_copyfields))

        if len(set(self._copyfields)) != len(self._copyfields):
            raise ValueError("Redundant copyfields: {}".format(
                self._copyfields))

        indexed_not_copied = set(self._indexed_copyfields) - set(
            self._copyfields)
        if indexed_not_copied:
            raise ValueError("Fields in index_copyfields but not in "
                             "copyfields: {}".format(indexed_not_copied))

        # allfields = [self._srcpkfield, self._srcfield] + self._copyfields
        # if len(allfields) != len(set(allfields)):
        #     raise ValueError(
        #         "Field overlap in InputFieldConfig: {}".format(section))
        # RE-THOUGHT: OK to copy source text fields etc. if desired.
        # It's fine in SQL to say SELECT a, a FROM mytable;

        self._db = nlpdef.get_database(self._srcdb)

    def get_srcdb(self) -> str:
        return self._srcdb

    def get_srctable(self) -> str:
        return self._srctable

    def get_srcpkfield(self) -> str:
        return self._srcpkfield

    def get_srcfield(self) -> str:
        return self._srcfield

    def _get_source_session(self):
        return self._db.session

    def _get_source_metadata(self):
        return self._db.metadata

    def _get_source_engine(self):
        return self._db.engine

    def _get_progress_session(self):
        return self._nlpdef.get_progdb_session()

    @staticmethod
    def get_srcref_columns_for_dest() -> List[Column]:
        """Columns referring to the source."""
        return [
            Column(FN_SRCDB, SqlTypeDbIdentifier,
                   doc="Source database name (from CRATE NLP config)"),
            Column(FN_SRCTABLE, SqlTypeDbIdentifier,
                   doc="Source table name"),
            Column(FN_SRCPKFIELD, SqlTypeDbIdentifier,
                   doc="PK field (column) in source table"),
            Column(FN_SRCPKVAL, BigInteger,
                   doc="PK of source record (or hash of PK if the PK is a "
                       "string"),
            Column(FN_SRCPKSTR, String(MAX_STRING_PK_LENGTH),
                   doc="NULL if the table has an integer PK, but the PK if "
                       "the PK was a string, to deal with hash collisions. "
                       "Max length: {}".format(MAX_STRING_PK_LENGTH)),
            Column(FN_SRCFIELD, SqlTypeDbIdentifier,
                   doc="Field (column) name of source text"),
        ]

    @staticmethod
    def get_srcref_indexes_for_dest() -> List[Index]:
        """Indexes for columns referring to the source."""
        # http://stackoverflow.com/questions/179085/multiple-indexes-vs-multi-column-indexes  # noqa
        return [
            Index('_idx_srcref',
                  FN_SRCDB, FN_SRCTABLE, FN_SRCPKFIELD, FN_SRCPKVAL,
                  FN_SRCPKSTR),
        ]

    def _require_table_exists(self) -> None:
        if not table_exists(self._get_source_engine(), self._srctable):
            msg = "Missing source table: {}.{}".format(self._srcdb,
                                                       self._srctable)
            log.critical(msg)
            raise ValueError(msg)

    def get_copy_columns(self) -> List[Column]:
        # We read the column type from the source database.
        self._require_table_exists()
        meta = self._get_source_metadata()
        t = Table(self._srctable, meta, autoload=True)
        copy_columns = []
        processed_copy_column_names = []
        for c in t.columns:
            # if c.name.lower() in self._copyfields:
            if c.name in self._copyfields:
                copied = c.copy()
                # Force lower case:
                # copied.name = copied.name.lower()
                # copied.name = quoted_name(copied.name.lower(), None)
                # ... this is not working properly. Keep getting an
                # "Unconsumed column names" error with e.g. a source field of
                # "Text".
                # Try making copyfields case-sensitive instead.
                copy_columns.append(copied)
                processed_copy_column_names.append(c.name)
        # Check all requested fields are present:
        missing = set(self._copyfields) - set(processed_copy_column_names)
        if missing:
            raise ValueError(
                "The following fields were requested to be copied but are "
                "absent from the source (NB case-sensitive): {}".format(
                    missing))
        # log.critical(copy_columns)
        return copy_columns

    def get_copy_indexes(self) -> List[Index]:
        self._require_table_exists()
        meta = self._get_source_metadata()
        t = Table(self._srctable, meta, autoload=True)
        copy_indexes = []
        processed_copy_index_col_names = []
        for c in t.columns:
            # if c.name.lower() in self._indexed_copyfields:
            if c.name in self._indexed_copyfields:
                copied = c.copy()
                # See above re case.
                idx_name = "idx_{}".format(c.name)
                copy_indexes.append(Index(idx_name, copied))
                processed_copy_index_col_names.append(c.name)
        missing = set(self._indexed_copyfields) - set(
            processed_copy_index_col_names)
        if missing:
            raise ValueError(
                "The following fields were requested to be copied/indexed but "
                "are absent from the source (NB case-sensitive): {}".format(
                    missing))
        return copy_indexes

    def gen_src_pks(self) -> Iterator(int):
        """
        Generate integer PKs from the source table specified for the
        InputFieldConfig.
        """
        session = self._get_source_session()
        query = (
            select([column(self._srcpkfield)]).
            select_from(table(self._srctable))
        )
        result = session.execute(query)
        for row in result:
            yield row[0]

    def gen_text(self,
                 tasknum: int = 0,
                 ntasks: int = 1) -> Iterator(Tuple[str, Dict[str, Any], bool]):
        """
        Generate text strings from the input database.
        Yields tuple of (text, dict), where the dict is a column-to-value
        mapping for all other fields (source reference fields, copy fields).
        """
        if 1 < ntasks <= tasknum:
            raise Exception("Invalid tasknum {}; must be <{}".format(
                tasknum, ntasks))
        base_dict = {
            FN_SRCDB: self._srcdb,
            FN_SRCTABLE: self._srctable,
            FN_SRCPKFIELD: self._srcpkfield,
            FN_SRCFIELD: self._srcfield,
        }
        session = self._get_source_session()
        pkcol = column(self._srcpkfield)
        # ... don't use is_sqlatype_integer with this; it's a column clause,
        # not a full column definition.
        pkcoltype = get_column_type(self._get_source_engine(), self._srctable,
                                    self._srcpkfield)
        if not pkcoltype:
            raise ValueError("Unable to get column type for column "
                             "{}.{}".format(self._srctable, self._srcpkfield))
        pk_is_integer = is_sqlatype_integer(pkcoltype)
        # log.debug("pk_is_integer: {} -> {}".format(repr(pkcoltype),
        #                                            pk_is_integer))

        selectcols = [pkcol, column(self._srcfield)]
        for extracol in self._copyfields:
            selectcols.append(column(extracol))
        query = (
            select(selectcols).
            select_from(table(self._srctable)).
            order_by(pkcol)
        )
        distribute_by_hash = False
        if ntasks > 1:
            if pk_is_integer:
                # Integer PK, so we can be efficient and bake the parallel
                # processing work division into the SQL:
                query = query.where(pkcol % ntasks == tasknum)
            else:
                distribute_by_hash = True
        hashed_pk = None
        for row in session.execute(query):  # ... a generator itself
            pkval = row[0]
            text = row[1]
            other_values = dict(zip(self._copyfields, row[2:]))
            if pk_is_integer:
                other_values[FN_SRCPKVAL] = pkval
                other_values[FN_SRCPKSTR] = None
            else:
                hashed_pk = hash64(pkval)
                other_values[FN_SRCPKVAL] = hashed_pk
                other_values[FN_SRCPKSTR] = pkval
            other_values.update(base_dict)
            if distribute_by_hash and hashed_pk % ntasks != tasknum:
                # We convert some non-integer thing into a deterministic but
                # roughly randomly distributed integer using hash64.
                # That produces a signed integer, but that doesn't
                # matter, because % works nonetheless.
                # This is obviously less efficient than dividing the work up
                # via SQL, because we have to fetch and hash something.
                continue
            yield text, other_values

    def get_count_max(self) -> Tuple[int, Optional[int]]:
        """
        Counts records in the input table for the given InputFieldConfig.
        Used for progress monitoring.
        """
        session = self._get_source_session()
        pkcol = column(self._srcpkfield)
        query = (
            select([func.count(), func.max(pkcol)]).
            select_from(table(self._srctable))
        )
        result = session.execute(query)
        return result.fetchone()  # count, maximum

    def get_progress_record(self,
                            srcpkval: int,
                            srchash: str = None,
                            srcpkstr: str = None) -> Optional[NlpRecord]:
        """
        Fetch a progress record (NlpRecord) for the given source record, if one
        exists.
        """
        session = self._get_progress_session()
        query = (
            session.query(NlpRecord).
            filter(NlpRecord.srcdb == self._srcdb).
            filter(NlpRecord.srctable == self._srctable).
            filter(NlpRecord.srcpkfield == self._srcpkfield).
            filter(NlpRecord.srcpkval == srcpkval).
            filter(NlpRecord.srcfield == self._srcfield)
        )
        if srchash is not None:
            query = query.filter(NlpRecord.srchash == srchash)
        if srcpkstr is not None:
            query = query.filter(NlpRecord.srcpkstr == srcpkstr)
        return query.one_or_none()

    def delete_progress_records_where_srcpk_not(self,
                                                src_pks: List[int]) -> None:
        progsession = self._get_progress_session()
        log.debug("delete_progress_records_where_srcpk_not... {}.{} -> "
                  "progressdb".format(self._srcdb, self._srctable))
        prog_deletion_query = (
            progsession.query(NlpRecord).
            filter(NlpRecord.srcdb == self._srcdb).
            filter(NlpRecord.srctable == self._srctable).
            filter(NlpRecord.srcpkfield == self._srcpkfield).
            filter(NlpRecord.nlpdef == self._nlpdef.get_name())
        )
        if src_pks:
            log.debug("... deleting selectively")
            prog_deletion_query = prog_deletion_query.filter(
                ~NlpRecord.srcpkval.in_(src_pks)
            )
        else:
            log.debug("... deleting all")
        progsession.execute(prog_deletion_query)
        progsession.commit()

