import boto3
import config

from opensearchpy import OpenSearch, RequestsHttpConnection, Search
from requests_aws4auth import AWS4Auth
from snowflake_executor import SFExecutor


def get_es_ids_from_snowflake():
    with SFExecutor(config.SF_PARAMS) as snowflake_executor:
        es_ids_tuples = snowflake_executor.get_es_ids()
        es_ids = [es_id[0] for es_id in es_ids_tuples]
        return es_ids

def get_opensearch_client():
    """Get OpenSearch client."""
    region = 'us-east-1'
    service = 'es'
    credentials = boto3.Session().get_credentials()
    awsauth = AWS4Auth(
        credentials.access_key, credentials.secret_key, region, service,
        session_token=credentials.token
    )

    return OpenSearch(
            [config.OPENSEARCH_HOST], use_ssl=config.OPENSEARCH_USE_SSL,
            port=config.OPENSEARCH_PORT, timeout=config.OPENSEARCH_TIMEOUT,
            max_retries=config.OPENSEARCH_MAX_RETRIES, retry_on_timeout=True,
            http_auth=awsauth, connection_class=RequestsHttpConnection)


def get_es_ids_from_opensearch():
    opensearch_client = get_opensearch_client()
    s = Search(using=opensearch_client, index=config.OPENSEARCH_INDEX)

    res = s.scan()
    es_ids = [h.meta.id for h in res]
    return es_ids

print('Fetching ES IDs from Snowflake...')
sf_es_ids = get_es_ids_from_snowflake()
print('Fetching ES IDs from OpenSearch...')
es_ids = get_es_ids_from_opensearch()

missing_from_opensearch = set(sf_es_ids) - set(es_ids)
missing_from_snowflake = set(es_ids) - set(sf_es_ids)

print('Documents missing from snowflake (', len(missing_from_snowflake), ')')
for es_id in missing_from_snowflake:
    print(es_id)
print('Documents missing from opensearch (', len(missing_from_opensearch), ')')
for es_id in missing_from_opensearch:
    print(es_id)