import logging

from service.tasks.airflow.enrichment import AirflowEnrichmentRunner
from service.tasks.filtering import Set
from service.utils.aws_connectors import run_query
from service.utils.data_model_utils import (
    change_collection_status,
    generate_collection,
    get_rendered_sql_template,
)

logger = logging.getLogger(__name__)


def merge_collections_into_alliance(
    source_schema,
    management_schema,
    source_collection_ids,
    alliance_schema,
    user_id,
    task_handle,
):
    """Assumes, that source_collection_ids contains only sources and enrichments
    and all the parents for every child present"""

    # separate source collections into parents and children
    try:
        parents = run_query(
            f"SELECT id, name, collection_type FROM {source_schema}.collection "
            f"WHERE parent_id is null and id in ({','.join(str(cid) for cid in source_collection_ids)})"
        )

        if not parents:
            raise RuntimeError("Bad collection IDs")

        # create collections for parents
        coll_id_map = {
            cid: generate_collection(
                collection_name=f"{name}",
                collection_source=f"{source_schema}/{cid}/{user_id}",
                collection_type=coll_type,
                schema_name=alliance_schema,
                user_id=user_id,
                status="importing data",
                log_description="data import from workspace",
            )
            for cid, name, coll_type in parents
        }
    except Exception as e:
        logger.exception("BAD collection IDs", exc_info=e)
        task_handle.send_immediate_error(RuntimeError("Bad collection IDs"))
    # now let the sync part continue, as all the parent collections have already appeared.
    task_handle.send_immediate_response("merging")
    children = run_query(
        f"SELECT id, name, collection_type, parent_id FROM {source_schema}.collection "
        f"WHERE parent_id is not null "
        f"and id in ({','.join(str(cid) for cid in source_collection_ids if cid not in coll_id_map)})"
    )
    # create collections for children
    coll_id_map.update(
        {
            cid: generate_collection(
                collection_name=f"{name}",
                collection_source=f"{source_schema}/{cid}/{user_id}",
                parent_id=coll_id_map[parent_id],
                collection_type=coll_type,
                schema_name=alliance_schema,
                user_id=user_id,
                status="importing data",
                log_description="data import from workspace",
            )
            for cid, name, coll_type, parent_id in children
        }
    )

    params = {
        "id_map_table": f"merge_alliance_{alliance_schema[:20]}_id_{list(coll_id_map.values())[0]}",
        "id_map_data": f"{','.join(f'({i1},{i2})' for i1, i2 in coll_id_map.items())}",
        "temp_table": f"merge_alliance_{alliance_schema[:20]}_{list(coll_id_map.values())[0]}",
        "source_collection_ids": ",".join(str(cid) for cid in source_collection_ids),
        "source_schema": source_schema,
        "alliance_schema": alliance_schema,
        "management_schema": management_schema,
    }
    template_sql = "merge_collections_into_alliance.sql"
    sql = get_rendered_sql_template(params=params, template_sql=template_sql)
    run_query(query_sql=sql)

    for cid in coll_id_map.values():
        change_collection_status(alliance_schema, user_id, cid, "finished")


def generate_alliance_segments(schema, user_id):

    new_global_set = Set(schema, user_id).new_segment(
        [{"id": "0", "values": []}],
        name="All fans in alliance",
        source="GLOBAL ANALYTICS",
        finished_status="processing",
    )
    collection_id = new_global_set.set_id

    if new_global_set.json_for_appsync()["totalProfiles"] == 0:
        change_collection_status(schema, user_id, collection_id, "finished")
        logger.info(
            f"{schema} - alliance collection {collection_id} is empty - nothing to enrich."
        )
        return None

    airflow_runner = AirflowEnrichmentRunner(
        user_id=user_id, alliance_schema=schema, collection_id=collection_id
    )
    try:
        airflow_runner.execute_alliance_enrichment()
        change_collection_status(schema, user_id, collection_id, "enriching")
    except Exception as e:
        logger.error(e, exc_info=e)
        change_collection_status(schema, user_id, collection_id, "failed")
        raise RuntimeError(f"Airflow run has failed to start: {e}")

    logger.info(
        f"{schema} - successfully sent alliance collection {collection_id} segment generation task to Airflow!"
    )


def get_alliance(alliance_schema):
    for glob in run_query(
        f"SELECT cc.id, cc.status FROM {alliance_schema}.collection cc "
        f"WHERE cc.source = 'GLOBAL ANALYTICS' ORDER BY cc.id DESC LIMIT 1"
    ):
        status = glob[1]
        break
    else:
        status = ""

    return {"id": alliance_schema, "status": status}