from elasticsearch_client import elasticsearch_client
import concurrent.futures
from snowflake_client import snowflake_client


def playlist_positions_over_time(args):
    isrc = args.get('isrc')
    playlist_id = args.get('playlist_id')
    elasticsearch_client = args.get('es_client')

    query = {
        'from': 0,
        'size': 200,
        'query': {
            'bool': {
                'filter': [
                    {
                        'bool': {
                            'must': [
                                {
                                    'bool': {
                                        'must': [
                                            {
                                                'term': {
                                                    'isrc': {
                                                        'value': isrc,
                                                        'boost': 1
                                                    }
                                                }
                                            },
                                            {
                                                'term': {
                                                    'playlist_id': {
                                                        'value': playlist_id,
                                                        'boost': 1
                                                    }
                                                }
                                            }
                                        ],
                                        'adjust_pure_negative': True,
                                        'boost': 1
                                    }
                                }
                            ],
                            'adjust_pure_negative': True,
                            'boost': 1
                        }
                    }
                ],
                'adjust_pure_negative': True,
                'boost': 1
            }
        },
        '_source': {
            'includes': [
                'position',
                '@timestamp'
            ],
            'excludes': []
        },
        'sort': [
            {
                '@timestamp': {
                    'order': 'desc'
                }
            }
        ]
    }

    results = elasticsearch_client.search(
        index='placements_timeseries', body=query)
    return results


def current_placements_in_playlist(args):
    playlist_id = args.get('playlist_id')
    elasticsearch_client = args.get('es_client')

    query = {
        'size': 0,
        'query': {
            'bool': {
                'filter': [
                    {
                        'bool': {
                            'must': [
                                {
                                    'term': {
                                        'playlist_id': {
                                            'value': playlist_id,
                                            'boost': 1
                                        }
                                    }
                                }
                            ],
                            'adjust_pure_negative': True,
                            'boost': 1
                        }
                    }
                ],
                'adjust_pure_negative': True,
                'boost': 1
            }
        },
        '_source': {
            'includes': [
                'playlist_name',
                'track_name',
                'position',
                '@timestamp'
            ],
            'excludes': []
        },
        'aggs': {
            'current_placements': {
                'composite': {
                    'sources': [
                        {'positions': {'terms': {'field': 'position'}}}
                    ],
                    'size': 100
                },
                'aggs': {
                    'current_placement': {
                        'top_hits': {
                            'sort': [
                                {
                                    '@timestamp': {
                                        'order': 'desc'
                                    }
                                }
                            ],
                            '_source': {
                                'includes': ['track_name', 'position', '@timestamp']
                            },
                            'size': 1
                        }
                    }
                }
            }
        }
    }

    results = elasticsearch_client.search(
        index='placements_timeseries', body=query)
    return results


def current_placements_in_playlist_with_streams(args):
    playlist_id = args.get('playlist_id')
    elasticsearch_client = args.get('es_client')

    placements_query =\
        {
            'size': 0,
            'query': {
                'bool': {
                    'filter': [
                        {
                            'bool': {
                                'must': [
                                    {
                                        'term': {
                                            'playlist_id': {
                                                'value': playlist_id,
                                                'boost': 1,
                                            }
                                        }
                                    }
                                ],
                                'adjust_pure_negative': True,
                                'boost': 1,
                            }
                        }
                    ],
                    'adjust_pure_negative': True,
                    'boost': 1,
                }
            },
            '_source': {
                'includes': ['playlist_name', 'track_name', 'position', '@timestamp'],
                'excludes': [],
            },
            'aggs': {
                'current_placements': {
                    'composite': {
                        'sources': [{'positions': {'terms': {'field': 'position'}}}],
                        'size': 100,
                    },
                    'aggs': {
                        'current_placement': {
                            'top_hits': {
                                'sort': [{'@timestamp': {'order': 'desc'}}],
                                '_source': {
                                    'includes': ['track_name', 'position', '@timestamp', 'isrc']
                                },
                                'size': 1,
                            }
                        }
                    },
                }
            }
        }

    results1 = elasticsearch_client.search(
        index='placements_timeseries', body=placements_query)

    isrcs = [result['current_placement']['hits']['hits'][0]['_source']['isrc']
             for result in results1['aggregations']['current_placements']['buckets']]

    streams_query =\
        {
            'size': 0,
            'query': {
                'bool': {
                    'filter': [
                        {
                            'bool': {
                                'must': [
                                    {
                                        'bool': {
                                            'must': [
                                                {
                                                    'term': {
                                                        'playlist_id': {
                                                            'value': playlist_id,
                                                            'boost': 1,
                                                        }
                                                    }
                                                }
                                            ],
                                            'adjust_pure_negative': True,
                                            'boost': 1,
                                        }
                                    },
                                    {
                                        'bool': {
                                            'must': [
                                                {
                                                    'terms': {
                                                        'isrc': isrcs
                                                    }
                                                }
                                            ],
                                            'adjust_pure_negative': True,
                                            'boost': 1,
                                        }
                                    },
                                ],
                                'adjust_pure_negative': True,
                                'boost': 1,
                            }
                        }
                    ],
                    'adjust_pure_negative': True,
                    'boost': 1,
                }
            },
            '_source': {
                'includes': ['@timestamp', 'country_code', 'streams', 'playlist_id', 'isrc'],
                'excludes': [],
            },
            'sort': [{'@timestamp': {'order': 'desc'}}],
            'aggs': {
                'streams_per_isrc': {
                    'composite': {
                        'size': 65535,
                        'sources': [{'isrc': {'terms': {'field': 'isrc'}}}],
                    },
                    'aggregations': {'total_streams': {'sum': {'field': 'streams'}}},
                }
            }
        }

    results2 = elasticsearch_client.search(
        index='placements_streams', body=streams_query)

    return results1, results2


def fetch_streams(cs, isrc, playlist_id):
    cs.execute("""
            SELECT SUM(daily.streams), download_activity_date
            FROM
            facts.prod.streams_by_track_playlist_country_feed_distributor_daily
                AS daily
            WHERE isrc = '%s' AND store_playlist_id = '%s'
            GROUP BY daily.download_activity_date
            ORDER BY daily.download_activity_date DESC;
        """ % (isrc, playlist_id))
    streams = cs.fetchall()
    return streams


def fetch_streams_multiple(cs, isrcs, playlist_id):
    isrcs_str = ','.join(['\'{}\''.format(isrc) for isrc in isrcs])
    cs.execute("""
            SELECT SUM(daily.streams), isrc
            FROM
            facts.prod.streams_by_track_playlist_country_feed_distributor_daily
                AS daily
            WHERE isrc IN (%s) AND store_playlist_id = '%s'
            GROUP BY isrc;
        """ % (isrcs_str, playlist_id))
    streams = cs.fetchall()
    return streams


def current_placements_in_playlist_with_streams_hybrid(args):
    playlist_id = args.get('playlist_id')
    elasticsearch_client = args.get('es_client')
    snowflake_client = args.get('sf_client')

    placements_query =\
        {
            'size': 0,
            'query': {
                'bool': {
                    'filter': [
                        {
                            'bool': {
                                'must': [
                                    {
                                        'term': {
                                            'playlist_id': {
                                                'value': playlist_id,
                                                'boost': 1,
                                            }
                                        }
                                    }
                                ],
                                'adjust_pure_negative': True,
                                'boost': 1,
                            }
                        }
                    ],
                    'adjust_pure_negative': True,
                    'boost': 1,
                }
            },
            '_source': {
                'includes': ['playlist_name', 'track_name', 'position', '@timestamp'],
                'excludes': [],
            },
            'aggs': {
                'current_placements': {
                    'composite': {
                        'sources': [{'positions': {'terms': {'field': 'position'}}}],
                        'size': 100,
                    },
                    'aggs': {
                        'current_placement': {
                            'top_hits': {
                                'sort': [{'@timestamp': {'order': 'desc'}}],
                                '_source': {
                                    'includes': ['track_name', 'position', '@timestamp', 'isrc']
                                },
                                'size': 1,
                            }
                        }
                    },
                }
            }
        }

    results1 = elasticsearch_client.search(
        index='placements_timeseries', body=placements_query)

    isrcs = [result['current_placement']['hits']['hits'][0]['_source']['isrc']
             for result in results1['aggregations']['current_placements']['buckets']]

    snowflake_client = args.get('sf_client')
    cs = snowflake_client.cursor()
    results2 = fetch_streams_multiple(cs, isrcs, playlist_id)

    return results1, results2


def placements_and_streams_overtime_hybrid(args):
    isrc = args.get('isrc')
    playlist_id = args.get('playlist_id')
    snowflake_client = args.get('sf_client')

    cs = snowflake_client.cursor()

    try:
        futures = {}
        results = {}
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            future1 = executor.submit(
                playlist_positions_over_time, args)
            futures[future1] = 'positions'
            future2 = executor.submit(fetch_streams, cs, isrc, playlist_id)
            futures[future2] = 'streams'
            for future in concurrent.futures.as_completed(futures):
                what_completed = futures[future]
                try:
                    results[what_completed] = future.result()
                except Exception as exc:
                    print('%r generated an exception: %s' %
                          (what_completed, exc))

        return results
    finally:
        cs.close()


def placements_and_streams_overtime_es(args):
    isrc = args.get('isrc')
    playlist_id = args.get('playlist_id')
    elasticsearch_client = args.get('es_client')

    query1 = {
        'from': 0,
        'size': 200,
        'query': {
            'bool': {
                'filter': [
                    {
                        'bool': {
                            'must': [
                                {
                                    'bool': {
                                        'must': [
                                            {
                                                'term': {
                                                    'isrc': {
                                                        'value': isrc,
                                                        'boost': 1
                                                    }
                                                }
                                            },
                                            {
                                                'term': {
                                                    'playlist_id': {
                                                        'value': playlist_id,
                                                        'boost': 1
                                                    }
                                                }
                                            }
                                        ],
                                        'adjust_pure_negative': True,
                                        'boost': 1
                                    }
                                }
                            ],
                            'adjust_pure_negative': True,
                            'boost': 1
                        }
                    }
                ],
                'adjust_pure_negative': True,
                'boost': 1
            }
        },
        '_source': {
            'includes': [
                'position',
                '@timestamp'
            ],
            'excludes': []
        },
        'sort': [
            {
                '@timestamp': {
                    'order': 'desc'
                }
            }
        ]
    }

    query2 = {
        'size': 0,
        'query': {
            'bool': {
                'filter': [
                    {
                        'bool': {
                            'must': [
                                {
                                    'bool': {
                                        'must': [
                                            {
                                                'term': {
                                                    'isrc': {
                                                        'value': isrc,
                                                        'boost': 1
                                                    }
                                                }
                                            },
                                            {
                                                'term': {
                                                    'playlist_id': {
                                                        'value': playlist_id,
                                                        'boost': 1
                                                    }
                                                }
                                            }
                                        ],
                                        'adjust_pure_negative': True,
                                        'boost': 1
                                    }
                                }
                            ],
                            'adjust_pure_negative': True,
                            'boost': 1
                        }
                    }
                ],
                'adjust_pure_negative': True,
                'boost': 1
            }
        },
        '_source': {
            'includes': [
                '@timestamp',
                'country_code',
                'streams'
            ],
            'excludes': []
        },
        'sort': [
            {
                '@timestamp': {
                    'order': 'desc'
                }
            }
        ],
        'aggs': {
            'streams_over_time': {
                'composite': {
                    'size': 65535,
                    'sources': [
                        {
                            'date': {
                                'terms': {
                                    'field': '@timestamp'
                                }
                            }
                        }
                    ]
                },
                'aggregations': {
                    'total_streams': {
                        'sum': {
                            'field': 'streams'
                        }
                    }
                }
            }
        }
    }

    body = []
    body.extend([{'index': 'placements_timeseries'}, query1])
    body.extend([{'index': 'placements_streams'}, query2])
    results = elasticsearch_client.msearch(body=body)
    return results

def streams_for_isrc_playlist(isrc, playlist_id):
    streams_query =\
        {
            'size': 0,
            'query': {
                "bool": {
                    "must": [
                        {
                            "term": {
                                "playlist_id": playlist_id
                            }
                        },
                        {
                            "term": {
                                "isrc": isrc
                            }
                        }
                    ]
                }
            },
            '_source': {
                'includes': ['@timestamp', 'country_code', 'streams', 'playlist_id', 'isrc'],
                'excludes': [],
            },
            'sort': [{'@timestamp': {'order': 'desc'}}],
            'aggs': {
                'streams_per_isrc': {
                    'composite': {
                        'size': 65535,
                        'sources': [
                            {'isrc': {'terms': {'field': 'isrc'}}},
                            {'playlist_id': {'terms': {'field': 'playlist_id'}}}
                        ],
                    },
                    'aggregations': {'total_streams': {'sum': {'field': 'streams'}}},
                }
            }
        }

    return streams_query

def latest_placement_at_position_query(playlist_id, position):
    query = {
        'size': 1,
        'query': {
            'bool': {
                'must': [
                    {
                        'term': {
                            'playlist_id': playlist_id
                        }
                    },
                    {
                        'term': {
                            'position': position
                        }
                    }
                ]
            }
        },
        'sort': [
            {
                '@timestamp': {
                    'order': 'desc'
                }
            }
        ],
        '_source': {
            'includes': [
                'playlist_id',
                'track_name',
                'position',
                '@timestamp',
                'isrc'
            ]
        }
    }
    return query


def current_placements_given_isrc(args):
    isrc = args.get('isrc')
    es_client = args.get('es_client')

    query = {
       'size':0,
       'query': {
           'bool': {
               'must': {
                   'term': {
                       'isrc': isrc
                   }
               }
           }
       },
       'aggs': {
          'byPlaylistIdPosition': {
             'composite': {
                'sources': [
                   {
                      'playlist': {
                         'terms': {
                            'field': 'playlist_id'
                         }
                      }
                   },
                   {
                      'position': {
                         'terms': {
                            'field': 'position'
                         }
                      }
                   }
                ],
                'size': 7500
             },
             'aggs': {
                'latest_placement':{
                   'top_hits': {
                      'sort': [
                         {
                            '@timestamp': {
                               'order': 'desc'
                            }
                         }
                      ],
                      '_source': {
                         'includes': [
                            'playlist_id',
                            'position',
                            '@timestamp',
                         ]
                      },
                      'size': 1
                   }
                }
             }
          }
       }
    }

    results = es_client.search(index='placements_timeseries', body=query)
    queries = []
    for result in results['aggregations']['byPlaylistIdPosition']['buckets']:
        source = result['latest_placement']['hits']['hits'][0]['_source']
        queries.extend([{'index': 'placements_timeseries'}, latest_placement_at_position_query(source['playlist_id'], source['position'])])

    results = es_client.msearch(body=queries)

    responses = []
    for result in results['responses']:
        source = result['hits']['hits'][0]['_source']
        if source['isrc'] == isrc:
            responses.append(source)

    queries2 = []
    for response in responses:
        queries2.extend([{'index': 'placements_streams'}, streams_for_isrc_playlist(response['isrc'], response['playlist_id'])])

    streams_results = es_client.msearch(body=queries2)
    return responses, streams_results

def current_placements_given_isrc_hybrid(args):
    isrc = args.get('isrc')
    es_client = args.get('es_client')

    query = {
       'size':0,
       'query': {
           'bool': {
               'must': {
                   'term': {
                       'isrc': isrc
                   }
               }
           }
       },
       'aggs': {
          'byPlaylistIdPosition': {
             'composite': {
                'sources': [
                   {
                      'playlist': {
                         'terms': {
                            'field': 'playlist_id'
                         }
                      }
                   },
                   {
                      'position': {
                         'terms': {
                            'field': 'position'
                         }
                      }
                   }
                ],
                'size': 7500
             },
             'aggs': {
                'latest_placement':{
                   'top_hits': {
                      'sort': [
                         {
                            '@timestamp': {
                               'order': 'desc'
                            }
                         }
                      ],
                      '_source': {
                         'includes': [
                            'playlist_id',
                            'position',
                            '@timestamp',
                         ]
                      },
                      'size': 1
                   }
                }
             }
          }
       }
    }

    results = es_client.search(index='placements_timeseries', body=query)

    queries = []
    for result in results['aggregations']['byPlaylistIdPosition']['buckets']:
        source = result['latest_placement']['hits']['hits'][0]['_source']
        queries.extend([{'index': 'placements_timeseries'}, latest_placement_at_position_query(source['playlist_id'], source['position'])])

    results = es_client.msearch(body=queries)

    responses = []
    playlist_ids = []
    for result in results['responses']:
        source = result['hits']['hits'][0]['_source']
        if source['isrc'] == isrc:
            playlist_ids.append(source['playlist_id'])
            responses.append(source)

    playlist_ids_str = ','.join(['\'{}\''.format(pid) for pid in playlist_ids])

    sf_client = args.get('sf_client')
    results = sf_client.cursor().execute(
    f"""
        SELECT store_playlist_id, SUM(streams) as total_streams
        FROM FACTS.PROD.STREAMS_BY_TRACK_PLAYLIST_COUNTRY_FEED_DISTRIBUTOR_DAILY
        WHERE store_playlist_id IN ({playlist_ids_str}) AND isrc = '{isrc}'
        GROUP BY store_playlist_id;
    """
    )
    streams_results = results.fetchall()

    return responses, streams_results


if __name__ == '__main__':
    es = elasticsearch_client()
    sf = snowflake_client()
    res = playlist_positions_over_time(
        {
            'isrc': 'QMFME2004132',
            'playlist_id': '37i9dQZF1DXcBWIGoYBM5M',
            'es_client': es
        })

    print('=' * 500)

    res = current_placements_in_playlist(
        {
            'playlist_id': '37i9dQZF1DXcBWIGoYBM5M',
            'es_client': es
        })

    print('=' * 500)

    res = placements_and_streams_overtime_es(
        {
            'isrc': 'QMFME2004132',
            'playlist_id': '37i9dQZF1DXcBWIGoYBM5M',
            'es_client': es
        })

    print('=' * 500)

    res = current_placements_in_playlist_with_streams(
        {
            'playlist_id': '37i9dQZF1DXcBWIGoYBM5M',
            'es_client': es
        })

    print('=' * 500)

    res = current_placements_in_playlist_with_streams_hybrid(
        {
            'playlist_id': '37i9dQZF1DXcBWIGoYBM5M',
            'es_client': es,
            'sf_client': sf
        })

    print('=' * 500)

    res = current_placements_given_isrc({
        'isrc': 'QMFME2004132',
        'es_client': es
    })

    print('=' * 500)

    res = current_placements_given_isrc_hybrid({
        'isrc': 'QMFME2004132',
        'es_client': es,
        'sf_client': sf
    })
    print('=' * 500)