from typing import Any, List

from airflow.models import BaseOperator
from airflow.providers.elasticsearch.hooks.elasticsearch import ElasticsearchPythonHook
from airflow.utils.context import Context


__all__ = ["SwitchElasticsearchDataSource"]


class SwitchElasticsearchDataSource(BaseOperator):
    template_fields = ("index_name", "base_index_name")

    def __init__(self, hosts: List[str], index_name: str, base_index_name: str, **kwargs):
        super().__init__(**kwargs)
        self.hosts = hosts
        self.index_name = index_name
        self.base_index_name = base_index_name

    def execute(self, context: Context) -> Any:
        hook = ElasticsearchPythonHook(hosts=self.hosts).get_conn

        try:
            hook.indices.refresh(index=self.index_name)
            hook.indices.put_alias(index=self.index_name, name=self.base_index_name)
        except Exception:
            raise
        else:
            hook.indices.delete(index=f"{self.base_index_name}*,-{self.index_name}", ignore_unavailable=True)