from typing import Dict
from typing import Any
import json


def get_es_query(
    search_keys: Dict[str, Any],
    model_name: str,
    es_batch: int = 1000
) -> str:
    """
    Builds the Elasticsearch query to be used when retrieving data.

    Args
    ----
      args: NamedTuple
        args.search_keys: Dict[str, Any]
            Search query sent by the customer as well as other variables that sets its
            context, such as region, favorite brand and so on.
        args.model_name: str
            Name of RankLib model saved on Elasticsearch
        args.index: str
            Index on Elasticsearch where to retrieve documents
        args.es_batch: int
            How many documents to retrieve

    Returns
    -------
      query: str
          String representation of final query
    """
    # it's expected that a ES query will be available at:
    # ./queries/{model_name}/es_query.json
    query = open(f'queries/{model_name}/es_query.json').read()
    query = json.loads(query.replace('{query}', search_keys['search_term']))
    # We just want to retrieve the id of the document to evaluate the ranks between
    # customers purchases and the retrieve list result
    query['_source'] = '_id'
    query['size'] = es_batch
    query['rescore']['window_size'] = 50  # Hardcoded to optimize first 50 skus
    query['rescore']['query']['rescore_query']['sltr']['params'] = search_keys
    query['rescore']['query']['rescore_query']['sltr']['model'] = model_name
    return query