from intranet.dispenser.utils.common import utils
import argparse
import json
import logging
from psycopg2.extras import execute_values


def insert_mappings(mapping_file, env, dry_run):
    file = open(mapping_file)

    connection = dispenser_db_conn(env)
    cursor = utils.dict_cursor(connection)
    providers = load_dispenser_providers(cursor)
    resource_keys = load_dispenser_resources(cursor)
    segments = load_dispenser_segments(cursor)

    d_resources_by_provider = {}

    mapping_bodies = []

    mappings_by_campaign = json.load(file)
    for campaign_mappings in mappings_by_campaign:
        campaign_id = campaign_mappings['campaignId']
        for mapping in campaign_mappings['mappings']:
            provider_key = mapping["serviceKey"]
            resource_key = mapping["resourceKey"]
            segment_keys = mapping["segmentKeys"]

            provider_id = providers[provider_key]
            resource_id = resource_keys[provider_id][resource_key]
            segment_ids = sorted([segments[key] for key in segment_keys])
            if not mapping.get("targets"):
                mapping_bodies.append({
                    "campaign_id": campaign_id,
                    "resource_id": resource_id,
                    "segment_ids": segment_ids,
                    "key": None,
                    "external_resource_key": None,
                    "external_provider_key": None,
                    "external_resource_unit_key": None,
                    "numerator": None,
                    "denominator": None,
                    "skip": True
                })

            for target in mapping["targets"]:
                external_provider_key = target["externalProviderKey"]
                external_resource_key = target['externalResourceKey']
                external_unit = target['externalUnit']
                target_key = target.get('key')
                numerator = target['numerator']
                denominator = target['denominator']

                if d_resources_by_provider.get(external_provider_key) is None:
                    d_resources_by_provider[external_provider_key] = set()
                d_resources_by_provider[external_provider_key].add(external_resource_key)

                mapping_bodies.append({
                    "campaign_id": campaign_id,
                    "resource_id": resource_id,
                    "segment_ids": segment_ids,
                    "key": target_key,
                    "external_resource_key": external_resource_key,
                    "external_provider_key": external_provider_key,
                    "external_resource_unit_key": external_unit,
                    "numerator": numerator,
                    "denominator": denominator,
                    "skip": False
                })
    provider_id_by_key = {}

    d_service = utils.D(env == 'prod')
    all_providers = d_service.get_all_providers().body
    for provider in all_providers['items']:
        key = provider['key']
        id = provider['id']
        provider_id_by_key[key] = id

    d_resources = {}
    for (provider_key, resource_keys) in d_resources_by_provider.items():
        d_resources[provider_key] = {}
        provider_id = provider_id_by_key[provider_key]
        provider_resources = d_service.get_provider_resources(provider_id).body
        for resource in provider_resources['items']:
            if resource['key'] not in resource_keys:
                continue
            d_resources[provider_key][resource['key']] = resource

    for body in mapping_bodies:
        if body['skip']:
            continue
        provider_key = body['external_provider_key']
        resource_key = body['external_resource_key']
        resource_id = d_resources[provider_key][resource_key]['id']
        body['external_resource_id'] = resource_id

    logging.debug(f"Parsed mappings ({len(mapping_bodies)}): {json.dumps(mapping_bodies, indent=4)}")
    if not dry_run:
        insert_mappings_db(cursor, mapping_bodies)
        connection.commit()


def dispenser_db_conn(env):
    connection = utils.common_prod_connection() if env == 'prod' else utils.common_testing_connection()
    return connection


def insert_mappings_db(cursor, mappings):
    values = [(m['resource_id'], m['segment_ids'], m['campaign_id'], m.get('numerator'), m.get('denominator'),
               m.get('external_resource_id'), m.get('external_resource_unit_key'), m.get('key'), m['skip'])
              for m in mappings]
    logging.debug(f"INSERTING VALUES {values}")
    execute_values(cursor, "INSERT INTO resources_model_mapping(resource_id, segment_ids, campaign_id, numerator, "
                           "denominator, external_resource_id, external_resource_unit_key, key, skip) VALUES %s "
                           "ON CONFLICT DO NOTHING", values)


def load_dispenser_providers(cursor):
    cursor.execute("SELECT id, key FROM service")
    provider_id_by_key = {}
    for row in cursor:
        id = int(row['id'])
        key = row['key']
        provider_id_by_key[key] = id
    return provider_id_by_key


def load_dispenser_resources(cursor):
    cursor.execute("SELECT id, key, service_id FROM resource")
    result = {}
    for row in cursor:
        id = int(row['id'])
        service_id = int(row['service_id'])
        key = row['key']
        if result.get(service_id) is None:
            result[service_id] = {}
        result[service_id][key] = id
    return result


def load_dispenser_segments(cursor):
    cursor.execute("SELECT id, key FROM segment")
    result = {}
    for row in cursor:
        id = int(row['id'])
        key = row['key']
        result[key] = id
    return result


def prepare_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-e", "--env", required=True, choices=['prod', 'test'], type=str)
    parser.add_argument("--dry-run", required=False, default=False, nargs='?', const=True)
    parser.add_argument("-f", "--file", required=True, type=str, help="JSON с маппингами в корректном формате")
    return parser.parse_args()

"""
    Пока реализована только функциональность вставки маппингов. В будущем можно добавить команды, в т.ч. вывод хранимых
    маппингов в более читаемом виде.
"""
if __name__ == '__main__':
    args = prepare_args()
    utils.configure_logging("resource-mapping-tool.log")
    insert_mappings(args.file, args.env, args.dry_run)
