from elasticsearch import Elasticsearch

# Initialize the Elasticsearch client
es = Elasticsearch(
    [
        "https://elastic:8oKIqy312EBsAPzWT64NUzji@scic-elasticsearch.es.us-central1.gcp.cloud.es.io:443"
    ]
)


def update_documents_by_user_and_org(index_name, org_name, new_token_value):
    # Construct the search query to find documents matching username and org name
    query = {"query": {"bool": {"must": [{"match": {"organization_name": org_name}}]}}}

    # Search for documents that match the criteria
    response = es.search(index=index_name, body=query)

    # Check if any documents were found
    if response["hits"]["total"]["value"] > 0:
        for hit in response["hits"]["hits"]:
            doc_id = hit["_id"]
            # Update the assigned_token field for each matching document
            update_body = {"doc": {"assigned_tokens": 100000}}

            # Update the document in Elasticsearch
            es.update(index=index_name, id=doc_id, body=update_body)
            print(
                f"Updated assigned_token for document ID {doc_id} in index '{index_name}'"
            )
    else:
        print(f"No documents found matching username and org name '{org_name}'")


# Usage example
update_documents_by_user_and_org(
    "organization_level_token_count", "Risk & Insurance Education Alliance", 100000
)