import logging

import pandas as pd
from elasticsearch import Elasticsearch
from tqdm import tqdm


def create_index_elastic_chunks(
    df: pd.DataFrame,
    logger: logging.Logger | None,
):
    if logger is None:
        logger = logging.getLogger(__name__)

    # Подключение к Elasticsearch
    es = Elasticsearch(hosts='localhost:9200')

    INDEX_NAME = 'nmd_full_text2'

    # Удаление старого индекса, если он существует
    if es.indices.exists(index=INDEX_NAME):
        es.indices.delete(index=INDEX_NAME)

    mapping = {
        "mappings": {
            "properties": {
                "index": {"type": "keyword"},
                "text": {"type": "text", "analyzer": "standard"},
            }
        }
    }

    # Создание индекса с указанным маппингом
    es.indices.create(index=INDEX_NAME, body=mapping)

    # Индексация документов
    for ind, row in tqdm(df.iterrows()):
        document = {'index': ind, 'text': row['Text']}

        # Индексирование документа в Elasticsearch
        es.index(index=INDEX_NAME, id=ind, body=document)

    if es.indices.exists(index=INDEX_NAME):
        print(f"Index '{INDEX_NAME}' exists.")

    # # Подсчет количества документов в индексе
    count_response = es.count(index=INDEX_NAME)
    print(f"Total documents in '{INDEX_NAME}': {count_response['count']}")

    # Поиск документов, где поле "person_full_name" содержит определенное значение "Александров Д.В."
    query = {
        "query": {
            "multi_match": {
                "query": "4.1. Комиссия ГО имеет право: привлекать работников Компании (по согласованию с руководителями структурных подразделений) для подготовки проектов документов Комиссии ГО, в сроки, установленные Комиссией ГО, а также в целях выполнения других работ, необходимых для принятия решений Комиссии ГО; отклонять материалы, представленные для рассмотрения на заседания Комиссии ГО в случае, если материалы требуют доработки или не относятся к компетенции Комиссии ГО в соответствии с разделом 6 настоящего Положения; \uf02d запрашивать у руководителей структурных подразделений Компании информацию и документы для принятия решений в рамках компетенции Комиссии ГО в соответствии с разделом 6 настоящего Положения; приглашать на заседания Комиссии ГО работников Группы компаний «Норильский никель», представителей Комиссий Филиалов, а также внешних консультантов, экспертов",
                "fields": ["*"],
            }
        }
    }

    # Выполнение поиска в Elasticsearch
    response = es.search(index=INDEX_NAME, body=query, size=2)
    logger.info(f"Number of hits: {response['hits']['total']['value']}")

    # Вывод результата поиска
    for hit in response['hits']['hits']:
        logger.info(hit['_source'])


if __name__ == '__main__':
    df = pd.read_pickle(
        '/mnt/ntr_work/project/nmd800/data/db/dataset_local_tables2.pkl'
    )
    create_index_elastic_chunks(df)