import logging
import os
import time
from hashlib import sha256
from io import StringIO
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
from jinjasql import JinjaSql
from validate_email import validate_email

from service.async_task_manager import io_task
from service.db import engine
from service.utils import appsync_communication as appsync
from service.utils.aws_connectors import pd_read_sql, run_query

logger = logging.getLogger(__name__)

VISIBLE_COLLECTIONS = ["source", "collection set", "segment"]
FAN_PROFILE_ID_SALT = "48cd56a56ea370469b870a893a922188ad3d4d8466d2b44cd30880ddbb05eacbce153913d8f86dc7b51674c6308c4a07269c394a68bed35affa705bc7787fe8b".encode()


def get_sql_from_template(query, bind_params):
    return query % bind_params if bind_params else query


def get_rendered_sql_template(params, template_sql):
    """
    Apply a JinjaSql template (string) substituting parameters (dict) and return
    the final SQL.
    """
    current_path = os.path.dirname(__file__)
    with open(f"{current_path}/../sql/{template_sql}", "r") as fp:
        template_file = fp.read()

    j = JinjaSql(param_style="pyformat")
    query, bind_params = j.prepare_query(template_file, params)
    return get_sql_from_template(query, bind_params)


def add_event_log(
    schema: str, collection_id: int, user_id: str, log_type: str, description: str
):

    run_query(
        f"""INSERT INTO {schema}.event_log (collection_id, user_id, type, description)
                  VALUES (%(collection_id)s, %(user_id)s, %(log_type)s,%(description)s);""",
        {
            "collection_id": collection_id,
            "user_id": user_id,
            "log_type": log_type,
            "description": description,
        },
        fetch=False,
    )


def generate_collection(
    collection_name: str,
    collection_source: str,
    schema_name: str,
    user_id: str,
    parent_id: int = None,
    status: str = "undefined",
    collection_type: str = "source",
    log_description: str = "not specified",
) -> int:
    """Loads collection into collection table and return collection id"""

    sql = f"""
            INSERT INTO {schema_name}.collection (name, source, parent_id, status, collection_type)
            VALUES (%(collection_name)s,
                    %(collection_source)s,
                    %(parent_id)s,
                    %(status)s,
                    %(collection_type)s )
            RETURNING id;"""
    collection_id = run_query(
        sql,
        {
            "collection_name": collection_name,
            "collection_source": collection_source,
            "parent_id": parent_id,
            "status": status,
            "collection_type": collection_type,
        },
    )[0][0]

    """ Update event log """
    add_event_log(
        schema=schema_name,
        collection_id=collection_id,
        user_id=user_id,
        log_type="collection_creation",  # TODO! this has to be from some agreed enum
        description=log_description,
    )

    logger.info(f"Generated collection {collection_id}:'{collection_name}' dataframe")
    change_collection_status(schema_name, user_id, collection_id, status)
    return collection_id


def get_source_from_collection(schema: str, collection_id: Union[int, List[int]]):
    """Gets file name that is associated with collection id"""
    if isinstance(collection_id, list):
        sql = f"""SELECT source FROM {schema}.collection
                    WHERE id IN ({','.join([str(int(x)) for x in collection_id])});"""
        file_names = run_query(sql)
        if file_names:
            return [fn[0] for fn in file_names]
        else:
            return []
    else:
        sql = f"""SELECT source FROM {schema}.collection
                    WHERE id = %(collection_id)s;"""
        file_name = run_query(sql, {"collection_id": collection_id})
        if file_name:
            return file_name[0][0]
        else:
            return ""


def get_collection_type(schema: str, collection_id: int) -> Optional[str]:
    """Gets collection type"""
    for collection_type in run_query(
        f"SELECT collection_type FROM {schema}.collection WHERE id = %(collection_id)s",
        {"collection_id": collection_id},
    ):
        return collection_type[0]
    return None


def get_collection_name(schema: str, collection_id: int) -> str:
    """Gets collection type"""
    for collection_name in run_query(
        f"SELECT name FROM {schema}.collection WHERE id = %(collection_id)s",
        {"collection_id": collection_id},
    ):
        return collection_name[0]
    return ""


def get_management_schema(schema: str, for_fan_table=False) -> str:
    """Returns management schema for alliances/workspaces. If it's company schema, just returns itself."""
    if for_fan_table and schema.startswith("a"):
        return schema

    if schema.startswith("c"):
        return schema

    for mgm in run_query(
        "SELECT management_company_id FROM commons.company_alliance WHERE alliance_id = %(schema)s",
        {"schema": schema},
    ):
        return mgm[0]

    raise RuntimeError("Unknown schema")


def get_collection_attribute_names(schema, collection_id=None):
    """Gets attribute names used by collection"""
    sql = f"""SELECT a.name FROM {schema}.attribute a """
    if collection_id is not None:
        sql += (
            f"INNER JOIN {schema}.collection_attribute ca ON a.id = ca.attribute_id "
            f"INNER JOIN {schema}.collection c ON ca.collection_id = c.id "
            f"WHERE c.id = %(collection_id)s OR c.parent_id = %(collection_id)s;"
        )
    attribute_names = run_query(sql, {"collection_id": collection_id})
    return [an[0] for an in attribute_names] if attribute_names else []


def get_attribute_id(schema, attribute_name: str):
    """Gets collection type"""
    sql = """SELECT a_id FROM commons.system_label
             WHERE a_system_name = %(attribute_name)s;"""
    attribute_id = run_query(sql, {"attribute_name": attribute_name})
    return attribute_id[0][0] if attribute_id else None


def rename_collection(
    schema: str, user_id: str, collection_id: int, collection_name: str
):
    """Gets file name that is associated with collection id"""
    return update_collection(schema, user_id, collection_id, name=collection_name)


def get_schema_collection_id_for_url_path(url_path: str) -> Tuple[str, int, str]:
    """Gets schema and collection_id that is associated with url_path.
    Also deletes the record from the table, as we never ever get the same event again, that would use the data.
    """
    sql = """DELETE FROM commons.upload_url WHERE url_path = %(url_path)s
             RETURNING schema_name, collection_id, user_id;"""
    result = run_query(sql, {"url_path": url_path})
    return result


def delete_garbage():
    """Delete rows which were created more than an hour ago"""
    sql = "DELETE FROM commons.upload_url WHERE upload_timestamp < NOW() - interval '1 hour';"
    run_query(query_sql=sql)


def delete_used_guesses(schema: str, collection_id: int):
    sql = f"DELETE FROM {schema}.guessed_file_upload_fields WHERE collection_id = %(collection_id)s"
    params = {"collection_id": collection_id}
    run_query(sql, params)


def read_saved_field_map(
    schema: str, collection_id: str, only_assigned=True, with_system_attribute_ids=True
):
    params = {"collection_id": collection_id}

    only_assigned_str = ""
    if only_assigned:
        """We need to check if new system fields are assigned to previously non system fields OR vise versa"""
        only_assigned_str = """AND ((g.system_field_name IS NOT NULL AND g.system_field_name != '')
                                OR (g.system_field_name = '' AND g.system_field_id < 10000))"""

    system_attribute_select = ""
    system_attribute_join = ""
    if with_system_attribute_ids:
        system_attribute_select = ",a.a_id system_field_id"
        system_attribute_join = (
            "LEFT JOIN commons.system_label a ON g.system_field_name = a.a_system_name"
        )
    sql = f"""
        SELECT g.system_field_name,
               g.file_field_name,
               g.system_field_id old_attribute_id
               {system_attribute_select}
        FROM {schema}.guessed_file_upload_fields g
        {system_attribute_join}
        WHERE g.collection_id = %(collection_id)s
        {only_assigned_str};
        """
    df = run_query(sql, params, return_type="df")
    return df


def save_field_map(schema: str, collection_id: int, field_map: list):
    collection_id = int(collection_id)
    sql = (
        f"UPDATE {schema}.guessed_file_upload_fields as g SET system_field_name = data.val "
        f"FROM (VALUES %s) as data(id, idx, val) "
        f"WHERE g.collection_id = data.id AND g.column_number = data.idx;"
    )
    values = [(collection_id, fm["index"], fm["value"]) for fm in field_map]
    run_query(sql, values, update_values=True)


def change_collection_status(
    schema: str,
    user_id: str,
    collection_id: int,
    status: str,
    collection_type: str = None,
    collection_backup=None,
):
    return update_collection(
        schema,
        user_id,
        collection_id,
        status=status,
        collection_type=collection_type,
        collection_backup=collection_backup,
    )


def update_collection(
    schema: str,
    user_id: str,
    collection_id: int,
    name: str = None,
    collection_type: str = None,
    status: str = None,
    collection_backup=None,
):
    """Updating the 'status' field of collection"""
    sql_params: Dict[str, Union[str, int]] = {}
    if name:
        sql_params["name"] = name
    if status:
        sql_params["status"] = status
    if collection_type:
        sql_params["collection_type"] = collection_type

    if sql_params:
        sql = (
            f"UPDATE {schema}.collection SET {','.join([f'{q} = %({q})s' for q in sql_params])} "
            f"WHERE id = %(collection_id)s RETURNING source, collection_type, "
            f"(SELECT status FROM {schema}.collection WHERE id = %(collection_id)s)"
        )
        sql_params["collection_id"] = int(collection_id)
        update_result = run_query(
            sql, sql_params
        )  # can be already deleted, so returns nothing
        try:
            template_sql = "get_collection_w_profiles.sql"
            params = {"schema_name": schema, "collection_id": collection_id}
            sql = get_rendered_sql_template(params, template_sql)

            for res in run_query(
                sql, return_type="dict"
            ):  # can be already deleted, so returns nothing
                collection = res
                break
            else:
                collection = {
                    "id": collection_id,
                    "name": (collection_backup or {}).get("name"),
                    "status": status,
                    "parentId": (collection_backup or {}).get("parentId"),
                    "dateCreated": (collection_backup or {}).get("dateCreated"),
                    "type": (collection_backup or {}).get("type"),
                    "totalProfiles": (collection_backup or {}).get("totalProfiles"),
                }

            if not status or not update_result or status != update_result[0][2]:
                # Only notify and log, when status actually changed
                ct = (
                    update_result[0][1]
                    if update_result
                    else collection_backup["type"]
                    if collection_backup and "type" in collection_backup
                    else None
                )
                parent_id = collection["parentId"] if "parentId" in collection else None
                if ct in VISIBLE_COLLECTIONS and (
                    parent_id is None or status in ("finished", "deleted")
                ):
                    # Only care about visible collection types or finished status for children.
                    try:
                        # Send collection fields defined in appsync-serverless to AppSync
                        appsync.send_collection(schema, collection)
                    except Exception as e:
                        logger.exception(e)
                # As a last step, store status log
                if status:
                    sql_params["collection_id"] = int(collection_id)
                    sql_params["user_id"] = user_id
                    sql = (
                        f"INSERT INTO {schema}.collection_status_log (collection_id, user_id, status) "
                        f"VALUES (%(collection_id)s, %(user_id)s, %(status)s ) "
                    )
                    run_query(sql, sql_params)
                    add_event_log(
                        schema,
                        collection_id,
                        user_id,
                        "collection_delete"
                        if status == "deleted"
                        else "collection_status_change",
                        status,
                    )
            if name:
                add_event_log(
                    schema, collection_id, user_id, "collection_name_change", name
                )
            return collection

        except Exception as e:
            logger.exception(e)
    return {}


def change_collection_source(schema: str, collection_id: int, source: str):
    """Updating the 'source' field of collection"""
    sql = f"UPDATE {schema}.collection SET source = %(source)s WHERE id = %(collection_id)s"
    run_query(sql, {"source": source, "collection_id": collection_id})


def string_to_db_async(df: pd.DataFrame, schema, table) -> None:
    """This method can be used to loading bigger dataframes in order to increase loading speed
    df has to have exact structure as the target DB table
    NOTE: can't use serial and other default db behaviour

    # TODO! Use dask to do group_by in parallel too
    Use io_task to commit those batches in async manner.
    """
    tasks = []
    # Create dataframes into roughly equal chunks for each worker
    io_workers = int(os.getenv("IO_WORKERS", 10)) - 1
    n = len(df) // io_workers
    for _, df in df.groupby(np.arange(len(df)) // n):
        df = df.astype(str)
        tasks.append(_async_generate_and_commit_string(df, schema, table))
    for task in tasks:
        task.wait_for_result()
    return


@io_task
def _async_generate_and_commit_string(df, schema, table, **kwargs):
    """Unfortunately, you cannot have two async cursors with one connection with psycopg2.
    https://github.com/aio-libs/aiopg/issues/535
    So if we have 10 async tasks, this means 10 connections are created.
    """
    connection = engine.raw_connection()
    f = StringIO()
    df.to_csv(f, index=False, header=False, sep="\t")  # remove header
    f.seek(0)  # move position to beginning of file before reading
    cursor = connection.cursor()
    # psycopg2-binary>2.9 breaks this: cursor.copy_from(f, f'{schema}.{table}', sep='\t')
    cursor.copy_expert(
        f"COPY {schema}.{table} FROM STDIN (FORMAT csv, DELIMITER '\t', HEADER False)",
        f,
    )
    connection.commit()
    connection.close()
    return


@io_task
def _async_load(schema, table, df, idx, if_exists="append", **kwargs):
    df.to_sql(
        table,
        engine,
        schema=schema,
        if_exists=if_exists,
        index=False,
        index_label=idx,  # this is the list of indices we want to create
        method="multi",
    )


@io_task
def _async_string_to_db(schema, df, table, **kwargs):
    """With smaller files (e.g. 20k rows), speed is similar. With larger (300k), the _async one is much faster."""
    string_to_db_async(df, schema, table)


def load_to_rds(schema, table_dict, if_exists="append") -> str:
    """
    Load multiple tables to RDS
    # TODO! add event_log function
    """
    loading = []
    for table, df in table_dict.items():
        if table == "fan_attribute":
            # We need to remove newline and tab character, else we build up messed up csv in later steps
            prepared_df = df[0]
            prepared_df["value"] = (
                prepared_df["value"]
                .replace(r"\n", " ", regex=True)
                .replace(r"\t", " ", regex=True)
                .replace(r"\r", " ", regex=True)
            )
            loading.append(
                _async_string_to_db(
                    schema=schema, df=prepared_df, table="fan_attribute"
                )
            )
        else:
            loading.append(
                _async_load(
                    schema=schema, table=table, df=df[0], idx=df[1], if_exists=if_exists
                )
            )
    for _, ll in enumerate(loading):
        try:
            ll.wait_for_result()
        except Exception as e:
            logger.error(f"Exception during multithreading waiting for results: {e}")
            ll.kwargs["df"].to_sql(
                ll.kwargs["table"],
                con=engine,
                schema=schema,
                if_exists=if_exists,
                index=False,
                method="multi",
            )

    # When everything is done, let's run "analyze" on all the tables, that got big inserts,
    # apparently that helps against things getting stuck
    sql = f"ANALYZE {','.join([f'{schema}.{table}' for table in table_dict])};"
    run_query(sql)

    nbr_tables = len(table_dict)
    return f"Unpacked and loaded data into {nbr_tables} tables in data model"


def attach_profile_id(df, fan_identifier_label) -> pd.DataFrame:
    """
    Obfuscated a field name, and return the dataframe with fan_id field that has this obufscated value.
    It also drops rows where fan_identifier_label is missing.
    """
    start = time.time()

    # Drop rows where fan_id is missing
    df = df[df[fan_identifier_label].notna()]

    # Filter by boolean indexing, to exclude non valid emails. We do this to avoid getting fan_id for TOTAL rows etc.
    df = df[df[fan_identifier_label].apply(validate_email)]

    # Hash our fan_identifier_label column with sha256 algorithm. That should be the same in absolutely every system.
    if hasattr(df, "parallel_apply"):
        df["profile_id"] = df[fan_identifier_label].parallel_apply(obfuscate_sha256)
    else:
        df["profile_id"] = df[fan_identifier_label].apply(obfuscate_sha256)
    end = time.time()
    logger.info(f"Hashing fan_id took {end - start} seconds")
    return df


def obfuscate_sha256(string: str):
    """Returns SHA 256 obfuscated string"""
    h = sha256(FAN_PROFILE_ID_SALT)
    h.update(str(string).lower().strip().encode())
    return h.hexdigest()


def get_related_collections(
    schema,
    collection_id: Union[int, List[int]],
    only_types=None,
    include_segment_sources=False,
    no_parents=False,
) -> list:
    """Returns a list of all collection_ids related with collection_id including the input collection_id
    include_segment_sources=True causes recursive collection of segment sources
    until real datasource collections are found
    """

    only_types = only_types or []
    internal_only_types = (
        only_types + ["segment"]
        if include_segment_sources and only_types and "segment" not in only_types
        else only_types
    )
    if isinstance(collection_id, (int, str)):
        collection_ids = [int(collection_id)]
    elif isinstance(collection_id, list):
        collection_ids = [int(cid) for cid in collection_id]
    else:
        raise RuntimeError("Check your inputs, please")
    if (
        len(collection_ids) == 0
        or (len(collection_ids) == 1 and collection_ids[0] == -1)
        or -1 in collection_ids
    ):
        parent_part = (
            f"parent as (SELECT id FROM {schema}.collection "
            f"WHERE collection_type = 'source')"
        )
    else:
        if no_parents:
            parent_part = (
                f"parent as (SELECT DISTINCT id "
                f"FROM {schema}.collection "
                f"WHERE id IN ({','.join(str(cid) for cid in collection_ids)}))"
            )
        elif include_segment_sources:
            parent_part = (
                f"RECURSIVE pre_parent as (SELECT DISTINCT id "
                f"FROM {schema}.collection "
                f"WHERE id IN ({','.join(str(cid) for cid in collection_ids)}) "
                f"UNION SELECT c.parent_id FROM {schema}.collection c "
                f"INNER JOIN pre_parent p ON p.id = c.id AND c.parent_id is not NULL), "
                f"parent as (SELECT DISTINCT id "
                f"FROM {schema}.collection "
                f"WHERE id IN (SELECT * from pre_parent) "
                f"UNION SELECT collection_id as id FROM {schema}.set_collection sc "
                f"INNER JOIN parent p ON p.id = sc.set_id)"
            )
        else:
            parent_part = (
                f"RECURSIVE parent as (SELECT DISTINCT id "
                f"FROM {schema}.collection "
                f"WHERE id IN ({','.join(str(cid) for cid in collection_ids)}) "
                f"UNION SELECT c.parent_id FROM {schema}.collection c "
                f"INNER JOIN parent p ON p.id = c.id AND c.parent_id is not NULL)"
            )

    data_types = (
        f""" AND collection_type IN ({','.join([f"'{tt}'" for tt in internal_only_types])})"""
        if internal_only_types
        else ""
    )
    query = f"""WITH {parent_part}
                SELECT id, collection_type, parent_id
                FROM {schema}.collection
                WHERE (id IN (SELECT * FROM parent)
                OR parent_id IN (SELECT * FROM parent)){data_types};"""
    result = pd_read_sql(query)
    try:
        segments = result.loc[
            (result["collection_type"] == "segment") & result["parent_id"].isnull()
        ]
        if len(segments):
            # remove all algos that are not directly under the earliest set,
            # if there are some later ones.
            if not result[
                (result["parent_id"].isin(segments["id"]))
                & (result["collection_type"] == "algo")
            ].empty:
                result.drop(
                    result[
                        (~result["parent_id"].isin(segments["id"]))
                        & (result["collection_type"] == "algo")
                    ].index,
                    inplace=True,
                )
    except Exception as e:
        logger.exception(e)
    try:
        if only_types:
            return list(result.loc[result["collection_type"].isin(only_types)]["id"])
        else:
            return list(result["id"])
    except Exception as e:
        logger.exception(f"Invalid collection error: {e}")
    return []


def get_parent(schema, collection_id):
    query = f"SELECT COALESCE(parent_id, {collection_id}) FROM {schema}.collection WHERE id = {collection_id}"
    res = run_query(query, {"collection_id": collection_id})
    if res:
        return res[0][0]
    else:
        return None


def get_table_structure(schema, table_name) -> list:
    """Returns a list of column names in specific table"""
    sql = f"""
        SELECT column_name
        FROM INFORMATION_SCHEMA.COLUMNS
        WHERE TABLE_NAME = '{table_name}'
            AND table_schema = '{schema}';
    """
    table_structure = pd_read_sql(sql)["column_name"]
    return table_structure


def refresh_materialized_queries(schema):
    template_sql = "refresh_materialized_views.sql"
    params = {"schema_name": schema}
    query = get_rendered_sql_template(params, template_sql)
    run_query(query)
    return True


def select_latest_alliance_analytics_segment(alliance_schema):
    for res in run_query(
        f"SELECT cc.id FROM {alliance_schema}.collection cc "
        f"WHERE cc.source = 'GLOBAL ANALYTICS' and cc.status = 'finished' "
        f"ORDER BY cc.id DESC LIMIT 1"
    ):
        return res[0]
    return None