import asyncio
import re
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from time import time
from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Callable, Optional, Tuple

from pydantic import BaseModel, Field, Secret, field_validator

from unstructured_ingest.data_types.file_data import (
    BatchFileData,
    BatchItem,
    FileData,
    FileDataSourceMetadata,
)
from unstructured_ingest.error import (
    DestinationConnectionError,
    SourceConnectionError,
)
from unstructured_ingest.interfaces import (
    AccessConfig,
    ConnectionConfig,
)
from unstructured_ingest.interfaces.downloader import download_responses
from unstructured_ingest.logger import logger
from unstructured_ingest.processes.connector_registry import (
    DestinationRegistryEntry,
    SourceRegistryEntry,
)
from unstructured_ingest.processes.connectors.elasticsearch.elasticsearch import (
    ElasticsearchBatchFileData,
    ElasticsearchDownloader,
    ElasticsearchDownloaderConfig,
    ElasticsearchIndexer,
    ElasticsearchIndexerConfig,
    ElasticsearchUploader,
    ElasticsearchUploaderConfig,
    ElasticsearchUploadStager,
    ElasticsearchUploadStagerConfig,
    ElastisearchAdditionalMetadata,
)
from unstructured_ingest.utils.data_prep import batch_generator, generator_batching_wbytes
from unstructured_ingest.utils.dep_check import requires_dependencies

if TYPE_CHECKING:
    from opensearchpy import OpenSearch

CONNECTOR_TYPE = "opensearch"

"""OpenSearch connector - inherits from Elasticsearch connector (OpenSearch is an ES fork)."""

# Precompiled regex patterns for AWS hostname detection (GovCloud, China, standard, FIPS)
_ES_PATTERN = re.compile(r"\.([a-z]{2}(?:-[a-z]+)+-\d+)\.es(?:-fips)?\.amazonaws\.com$")
_AOSS_PATTERN = re.compile(
    r"^[a-z0-9]+\.([a-z]{2}(?:-[a-z]+)+-\d+)\.aoss(?:-fips)?\.amazonaws\.com$"
)


def _run_coroutine(fn: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
    """Run an async function from sync context, handling existing event loops."""
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(fn(*args, **kwargs))

    with ThreadPoolExecutor(thread_name_prefix="opensearch") as pool:
        return pool.submit(lambda: asyncio.run(fn(*args, **kwargs))).result()


class OpenSearchAccessConfig(AccessConfig):
    password: Optional[str] = Field(default=None, description="password when using basic auth")
    aws_access_key_id: Optional[str] = Field(
        default=None,
        description="AWS access key ID. When provided (with secret), IAM authentication is used. "
        "Region and service type are auto-detected from the host URL.",
    )
    aws_secret_access_key: Optional[str] = Field(
        default=None,
        description="AWS secret access key. Required when aws_access_key_id is provided.",
    )
    aws_session_token: Optional[str] = Field(
        default=None, description="AWS session token for temporary credentials (optional)"
    )


def detect_aws_opensearch_config(host: str) -> Optional[Tuple[str, str]]:
    """Auto-detect AWS region and service from OpenSearch hostname."""
    clean_host = host.replace("https://", "").replace("http://", "")
    clean_host = clean_host.split(":")[0]

    match = _ES_PATTERN.search(clean_host)
    if match:
        return (match.group(1), "es")

    match = _AOSS_PATTERN.search(clean_host)
    if match:
        return (match.group(1), "aoss")

    return None


class OpenSearchClientInput(BaseModel):
    http_auth: Secret[Optional[tuple[str, str]]] = None
    hosts: Optional[list[str]] = None
    use_ssl: bool = False
    verify_certs: bool = False
    ssl_show_warn: bool = False
    ca_certs: Optional[str] = None
    client_cert: Optional[str] = None
    client_key: Optional[str] = None


class OpenSearchConnectionConfig(ConnectionConfig):
    hosts: list[str] = Field(
        ...,
        min_length=1,
        description="List of the OpenSearch hosts to connect",
        examples=["http://localhost:9200"],
    )
    username: Optional[str] = Field(default=None, description="username when using basic auth")
    use_ssl: bool = Field(default=False, description="use ssl for the connection")
    verify_certs: bool = Field(default=False, description="whether to verify SSL certificates")
    ssl_show_warn: bool = Field(
        default=False, description="show warning when verify certs is disabled"
    )
    ca_certs: Optional[Path] = Field(default=None, description="path to CA bundle")
    client_cert: Optional[Path] = Field(
        default=None,
        description="path to the file containing the private key and the certificate,"
        " or cert only if using client_key",
    )
    client_key: Optional[Path] = Field(
        default=None,
        description="path to the file containing the private key"
        " if using separate cert and key files",
    )

    access_config: Secret[OpenSearchAccessConfig]

    @field_validator("hosts", mode="before")
    @classmethod
    def validate_hosts(cls, value):
        if isinstance(value, str):
            value = [value]
        if not value:
            raise ValueError("At least one OpenSearch host must be provided. ")
        for host in value:
            if not host or not host.strip():
                raise ValueError("Host URL cannot be empty")
        return value

    def _has_aws_credentials(self) -> bool:
        """Check if AWS IAM credentials are provided."""
        access_config = self.access_config.get_secret_value()
        has_access_key = access_config.aws_access_key_id is not None
        has_secret_key = access_config.aws_secret_access_key is not None

        # Validate: Either both credentials or neither - partial credentials are invalid
        if has_access_key != has_secret_key:  # XOR: exactly one is set
            raise ValueError(
                "AWS IAM authentication requires BOTH aws_access_key_id and aws_secret_access_key. "
                f"Currently provided: aws_access_key_id={'set' if has_access_key else 'not set'}, "
                f"aws_secret_access_key={'set' if has_secret_key else 'not set'}"
            )

        return has_access_key and has_secret_key

    def _detect_and_validate_aws_config(self) -> Tuple[str, str]:
        """Auto-detect AWS region and service from host URL."""
        if not self.hosts:
            raise ValueError("Host is required for AWS OpenSearch connection")

        detected = detect_aws_opensearch_config(self.hosts[0])

        if not detected:
            raise ValueError(
                f"Could not auto-detect AWS region and service from host: {self.hosts[0]}. "
                f"Ensure your host URL follows AWS OpenSearch format: "
                f"https://search-domain-xxx.REGION.es.amazonaws.com (for OpenSearch Service), "
                f"https://search-domain-xxx.REGION.es-fips.amazonaws.com (for Service with FIPS), "
                f"https://xxx.REGION.aoss.amazonaws.com (for OpenSearch Serverless), or "
                f"https://xxx.REGION.aoss-fips.amazonaws.com (for Serverless with FIPS)"
            )

        region, service = detected
        logger.debug(
            f"Auto-detected AWS configuration from host: region={region}, service={service}"
        )
        return region, service

    @requires_dependencies(["opensearchpy", "boto3"], extras="opensearch")
    async def _get_async_aws_auth(self):
        """Create AWS SigV4 authentication handler for async clients."""
        import boto3
        from opensearchpy import AWSV4SignerAsyncAuth

        access_config = self.access_config.get_secret_value()

        session = boto3.Session(
            aws_access_key_id=access_config.aws_access_key_id,
            aws_secret_access_key=access_config.aws_secret_access_key,
            aws_session_token=access_config.aws_session_token,
        )
        credentials = session.get_credentials()

        if not credentials:
            raise ValueError("Failed to obtain AWS credentials from provided keys")

        return AWSV4SignerAsyncAuth(credentials, *self._detect_and_validate_aws_config())

    @requires_dependencies(["opensearchpy"], extras="opensearch")
    async def get_async_client_kwargs(self) -> dict:
        """Build AsyncOpenSearch client config (auto-detects IAM, basic auth, or cert auth)."""
        access_config = self.access_config.get_secret_value()
        client_input_kwargs = {}

        if self.hosts:
            client_input_kwargs["hosts"] = self.hosts
        if self.use_ssl:
            client_input_kwargs["use_ssl"] = self.use_ssl
        if self.verify_certs:
            client_input_kwargs["verify_certs"] = self.verify_certs
        if self.ssl_show_warn:
            client_input_kwargs["ssl_show_warn"] = self.ssl_show_warn
        if self.ca_certs:
            client_input_kwargs["ca_certs"] = str(self.ca_certs)
        if self.client_cert:
            client_input_kwargs["client_cert"] = str(self.client_cert)
        if self.client_key:
            client_input_kwargs["client_key"] = str(self.client_key)

        if self._has_aws_credentials():
            logger.debug("Using AWS IAM authentication")

            # Must use http_async.AsyncHttpConnection for IAM auth handlers
            from opensearchpy.connection.http_async import AsyncHttpConnection

            client_input = OpenSearchClientInput(**client_input_kwargs)
            client_kwargs = client_input.model_dump()
            client_kwargs["http_auth"] = await self._get_async_aws_auth()
            client_kwargs["connection_class"] = AsyncHttpConnection

        elif self.username and access_config.password:
            logger.debug("Using basic HTTP authentication")
            client_input_kwargs["http_auth"] = (self.username, access_config.password)

            client_input = OpenSearchClientInput(**client_input_kwargs)
            client_kwargs = client_input.model_dump()
            if client_input.http_auth:
                client_kwargs["http_auth"] = client_input.http_auth.get_secret_value()

        elif self.client_cert:
            logger.debug("Using certificate-based authentication")
            client_input = OpenSearchClientInput(**client_input_kwargs)
            client_kwargs = client_input.model_dump()

        else:
            logger.warning("No authentication configured - connecting without credentials")
            client_input = OpenSearchClientInput(**client_input_kwargs)
            client_kwargs = client_input.model_dump()

        # Retry and timeout configuration for resilience against transient errors
        client_kwargs["max_retries"] = 3
        client_kwargs["retry_on_status"] = [429, 502, 503]
        client_kwargs["retry_on_timeout"] = True
        client_kwargs["timeout"] = 60

        return {k: v for k, v in client_kwargs.items() if v is not None}


class OpenSearchIndexerConfig(ElasticsearchIndexerConfig):
    pass


@dataclass
class OpenSearchIndexer(ElasticsearchIndexer):
    connection_config: OpenSearchConnectionConfig
    index_config: OpenSearchIndexerConfig
    client: "OpenSearch" = field(init=False)

    def is_async(self) -> bool:
        """Signal pipeline to use async execution."""
        return True

    @requires_dependencies(["opensearchpy"], extras="opensearch")
    def precheck(self) -> None:
        """Validate connection and index (sync wrapper required by pipeline framework)."""

        async def _async_precheck():
            from opensearchpy import AsyncOpenSearch

            try:
                async with AsyncOpenSearch(
                    **await self.connection_config.get_async_client_kwargs()
                ) as client:
                    # Use get_alias (GET) instead of exists (HEAD) - HEAD has IAM signing issues
                    # Also respects AWS FGAC by checking only the specific index
                    await client.indices.get_alias(index=self.index_config.index_name)
            except Exception as e:
                logger.error(f"failed to validate connection: {e}", exc_info=True)
                raise SourceConnectionError(f"failed to validate connection: {e}")

        _run_coroutine(_async_precheck)

    @requires_dependencies(["opensearchpy"], extras="opensearch")
    async def run_async(self, **kwargs: Any) -> AsyncGenerator[ElasticsearchBatchFileData, None]:
        """Async indexing for all authentication types."""
        ids = list(await self._get_doc_ids_async())
        for batch in batch_generator(ids, self.index_config.batch_size):
            batch_items = [BatchItem(identifier=b) for b in batch]
            url = f"{self.connection_config.hosts[0]}/{self.index_config.index_name}"
            display_name = (
                f"url={url}, batch_size={len(batch_items)} "
                f"ids={batch_items[0].identifier}..{batch_items[-1].identifier}"
            )
            yield ElasticsearchBatchFileData(
                connector_type=CONNECTOR_TYPE,
                metadata=FileDataSourceMetadata(
                    url=url,
                    date_processed=str(time()),
                ),
                additional_metadata=ElastisearchAdditionalMetadata(
                    index_name=self.index_config.index_name,
                ),
                batch_items=batch_items,
                display_name=display_name,
            )

    @requires_dependencies(["opensearchpy"], extras="opensearch")
    async def _get_doc_ids_async(self) -> set[str]:
        """Fetch all document IDs, trying PIT + search_after first with scroll fallback.

        PIT is required for OpenSearch Serverless (AOSS) and preferred for
        OpenSearch Service. Falls back to scroll if PIT creation fails due to
        missing permissions (403) or unsupported version (400/404).
        """
        from opensearchpy import AsyncOpenSearch
        from opensearchpy.exceptions import TransportError

        async with AsyncOpenSearch(
            **await self.connection_config.get_async_client_kwargs()
        ) as client:
            try:
                pit = await client.create_pit(
                    index=self.index_config.index_name, params={"keep_alive": "5m"}
                )
            except TransportError as e:
                if e.status_code in (400, 403, 404):
                    logger.warning(
                        f"PIT creation failed (HTTP {e.status_code}), "
                        "falling back to scroll. Note: scroll is not supported on AOSS."
                    )
                    return await self._get_doc_ids_scroll(client)
                raise
            return await self._get_doc_ids_pit(client, pit["pit_id"])

    async def _get_doc_ids_pit(self, client: Any, pit_id: str) -> set[str]:
        """Paginate through all document IDs using an existing PIT context."""
        try:
            doc_ids: set[str] = set()
            search_after = None
            while True:
                body: dict[str, Any] = {
                    "stored_fields": [],
                    "_source": False,
                    "track_total_hits": False,
                    "query": {"match_all": {}},
                    "pit": {"id": pit_id, "keep_alive": "5m"},
                    # _shard_doc is not available on all OpenSearch versions; _id works
                    # on OpenSearch 2.x, AWS OpenSearch Service, and AOSS.
                    "sort": [{"_id": "asc"}],
                    "size": 1000,
                }
                if search_after:
                    body["search_after"] = search_after
                resp = await client.search(body=body)
                if "pit_id" in resp:
                    pit_id = resp["pit_id"]
                hits = resp["hits"]["hits"]
                if not hits:
                    break
                for hit in hits:
                    doc_ids.add(hit["_id"])
                search_after = hits[-1]["sort"]
            return doc_ids
        finally:
            try:
                await client.delete_pit(body={"pit_id": [pit_id]})
            except Exception:
                logger.warning("Failed to delete PIT, it will expire automatically")

    async def _get_doc_ids_scroll(self, client: Any) -> set[str]:
        """Fetch document IDs using scroll (fallback when PIT is unavailable)."""
        from opensearchpy.helpers import async_scan

        doc_ids: set[str] = set()
        async for hit in async_scan(
            client,
            query={"stored_fields": [], "query": {"match_all": {}}},
            scroll="1m",
            index=self.index_config.index_name,
        ):
            doc_ids.add(hit["_id"])
        return doc_ids


class OpenSearchDownloaderConfig(ElasticsearchDownloaderConfig):
    pass


@dataclass
class OpenSearchDownloader(ElasticsearchDownloader):
    connection_config: OpenSearchConnectionConfig
    download_config: OpenSearchDownloaderConfig
    connector_type: str = CONNECTOR_TYPE

    @requires_dependencies(["opensearchpy"], extras="opensearch")
    async def run_async(self, file_data: BatchFileData, **kwargs: Any) -> download_responses:
        """Download documents from OpenSearch.

        Uses a direct search by IDs instead of scroll, since the batch
        is already bounded by the indexer's batch_size (typically 100).
        """
        from opensearchpy import AsyncOpenSearch

        elasticsearch_filedata = ElasticsearchBatchFileData.cast(file_data=file_data)

        index_name: str = elasticsearch_filedata.additional_metadata.index_name
        ids: list[str] = [item.identifier for item in elasticsearch_filedata.batch_items]

        search_body: dict[str, Any] = {
            "version": True,
            "query": {"ids": {"values": ids}},
            "size": len(ids),
        }

        if self.download_config.fields:
            search_body["_source"] = self.download_config.fields

        download_responses = []
        async with AsyncOpenSearch(
            **await self.connection_config.get_async_client_kwargs()
        ) as client:
            resp = await client.search(body=search_body, index=index_name)
            for result in resp["hits"]["hits"]:
                download_responses.append(
                    self.generate_download_response(
                        result=result, index_name=index_name, file_data=elasticsearch_filedata
                    )
                )
        return download_responses


class OpenSearchUploaderConfig(ElasticsearchUploaderConfig):
    batch_size_bytes: int = Field(
        default=5_000_000,
        description="Size limit (in bytes) for each batch of items to be uploaded. "
        "Default is 5MB, lower than Elasticsearch default to accommodate "
        "AWS OpenSearch cluster rate limits.",
    )


@dataclass
class OpenSearchUploader(ElasticsearchUploader):
    connection_config: OpenSearchConnectionConfig
    upload_config: OpenSearchUploaderConfig
    connector_type: str = CONNECTOR_TYPE

    def is_async(self) -> bool:
        """Signal pipeline to use async execution."""
        return True

    @requires_dependencies(["opensearchpy"], extras="opensearch")
    def precheck(self) -> None:
        """Validate connection and index (sync wrapper required by pipeline framework)."""

        async def _async_precheck():
            from opensearchpy import AsyncOpenSearch

            try:
                async with AsyncOpenSearch(
                    **await self.connection_config.get_async_client_kwargs()
                ) as client:
                    # Use get_alias (GET) instead of exists (HEAD) - HEAD has IAM signing issues
                    # Also respects AWS FGAC by checking only the specific index
                    await client.indices.get_alias(index=self.upload_config.index_name)
            except Exception as e:
                logger.error(f"failed to validate connection: {e}", exc_info=True)
                raise DestinationConnectionError(f"failed to validate connection: {e}")

        _run_coroutine(_async_precheck)

    @requires_dependencies(["opensearchpy"], extras="opensearch")
    async def run_data_async(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
        """Upload data to OpenSearch using async_bulk."""
        from opensearchpy import AsyncOpenSearch
        from opensearchpy.exceptions import TransportError
        from opensearchpy.helpers import async_bulk

        logger.info(
            f"writing {len(data)} elements via document batches to destination "
            f"index named {self.upload_config.index_name} at {self.connection_config.hosts} "
            f"with batch size (in bytes) {self.upload_config.batch_size_bytes}"
        )

        async with AsyncOpenSearch(
            **await self.connection_config.get_async_client_kwargs()
        ) as client:
            for batch in generator_batching_wbytes(
                data, batch_size_limit_bytes=self.upload_config.batch_size_bytes
            ):
                # Retry with delay for rate limiting (429 errors)
                max_attempts = 3
                for attempt in range(max_attempts):
                    try:
                        success, failed = await async_bulk(
                            client=client,
                            actions=batch,
                            chunk_size=len(batch),
                            max_chunk_bytes=self.upload_config.batch_size_bytes,
                            raise_on_error=False,
                        )
                        break
                    except Exception as e:
                        # Check for rate limiting: precise type check, then string fallback
                        is_rate_limited = (
                            (isinstance(e, TransportError) and e.status_code == 429)
                            or "429" in str(e)
                            or "too many requests" in str(e).lower()
                        )

                        if attempt < max_attempts - 1 and is_rate_limited:
                            logger.warning(
                                f"Rate limited (attempt {attempt + 1}/{max_attempts}), "
                                f"waiting 5s before retry: {e}"
                            )
                            await asyncio.sleep(5)
                        else:
                            logger.error(f"Batch upload failed: {e}")
                            raise DestinationConnectionError(str(e))

                # Check for document failures (outside try to avoid catching our own exception)
                if failed:
                    logger.error(
                        f"Batch upload had {len(failed)} failures out of {len(batch)}. "
                        f"Failed items: {failed[:5]}"
                    )
                    raise DestinationConnectionError(
                        f"Failed to upload {len(failed)} out of {len(batch)} documents"
                    )

                logger.info(
                    f"uploaded batch of {len(batch)} elements to index "
                    f"{self.upload_config.index_name}"
                )


class OpenSearchUploadStagerConfig(ElasticsearchUploadStagerConfig):
    pass


@dataclass
class OpenSearchUploadStager(ElasticsearchUploadStager):
    upload_stager_config: OpenSearchUploadStagerConfig


opensearch_source_entry = SourceRegistryEntry(
    connection_config=OpenSearchConnectionConfig,
    indexer=OpenSearchIndexer,
    indexer_config=OpenSearchIndexerConfig,
    downloader=OpenSearchDownloader,
    downloader_config=OpenSearchDownloaderConfig,
)


opensearch_destination_entry = DestinationRegistryEntry(
    connection_config=OpenSearchConnectionConfig,
    upload_stager_config=OpenSearchUploadStagerConfig,
    upload_stager=OpenSearchUploadStager,
    uploader_config=OpenSearchUploaderConfig,
    uploader=OpenSearchUploader,
)
