"""
Elasticsearch sink for extractor.

Streams records into an ES index using bulk API.

Auth (env vars):
    EXTRACTOR_ES_URL      (default: http://localhost:9200)
    EXTRACTOR_ES_API_KEY  (Elastic Cloud / Elasticsearch API key)
    EXTRACTOR_ES_USER     (basic auth username)
    EXTRACTOR_ES_PASSWORD (basic auth password)
    EXTRACTOR_ES_CA_CERT  (path to CA certificate file)

Install: pip install 'coresdk-extractor[elasticsearch]'
"""

from __future__ import annotations

import logging
import os
from typing import Any

from .base import BaseSink
from .registry import register_sink

logger = logging.getLogger(__name__)

_INDEX_MAPPING = {
    "mappings": {
        "properties": {
            "id": {"type": "keyword"},
            "element_type": {"type": "keyword"},
            "text": {"type": "text", "analyzer": "standard"},
            "source_filename": {"type": "keyword"},
            "source_sha256": {"type": "keyword"},
            "source_uri": {"type": "keyword"},
            "source_media_type": {"type": "keyword"},
            "page": {"type": "integer"},
            "sequence_index": {"type": "integer"},
            "schema_version": {"type": "keyword"},
            "token_count": {"type": "integer"},
            "doc_category": {"type": "keyword"},
            "doc_category_score": {"type": "float"},
            "entities": {"type": "object", "enabled": False},
            "relations": {"type": "object", "enabled": False},
        }
    }
}


class ElasticsearchSink(BaseSink):
    def __init__(
        self,
        url: str | None = None,
        index: str = "extractor_elements",
        batch_size: int = 500,
    ) -> None:
        self.url = url or os.environ.get("EXTRACTOR_ES_URL", "http://localhost:9200")
        self.index = index
        self.batch_size = batch_size
        self._client = None
        self._buffer: list[dict[str, Any]] = []
        self._total_inserted = 0

    def connect(self) -> None:
        try:
            from elasticsearch import Elasticsearch
        except ImportError as e:
            raise ImportError(
                "elasticsearch not installed. Run: pip install 'coresdk-extractor[elasticsearch]'"
            ) from e
        kwargs: dict[str, Any] = {"hosts": [self.url]}
        api_key = os.environ.get("EXTRACTOR_ES_API_KEY")
        if api_key:
            kwargs["api_key"] = api_key
        else:
            user = os.environ.get("EXTRACTOR_ES_USER")
            password = os.environ.get("EXTRACTOR_ES_PASSWORD")
            if user and password:
                kwargs["basic_auth"] = (user, password)
        ca_cert = os.environ.get("EXTRACTOR_ES_CA_CERT")
        if ca_cert:
            kwargs["ca_certs"] = ca_cert
        self._client = Elasticsearch(**kwargs)
        if not self._client.indices.exists(index=self.index):
            self._client.indices.create(index=self.index, body=_INDEX_MAPPING)
        logger.info("ElasticsearchSink connected: %s/%s", self.url, self.index)

    def feed(self, record: dict[str, Any]) -> None:
        self._buffer.append(record)
        if len(self._buffer) >= self.batch_size:
            self.write_batch(self._buffer)
            self._buffer = []

    def write_batch(self, records: list[dict[str, Any]]) -> None:
        if not records or self._client is None:
            return
        from elasticsearch.helpers import bulk

        actions = [{"_index": self.index, "_id": r.get("id"), "_source": r} for r in records]
        success, _ = bulk(self._client, actions, raise_on_error=False)
        self._total_inserted += success

    def flush(self) -> None:
        if self._buffer:
            self.write_batch(self._buffer)
            self._buffer = []

    def close(self) -> None:
        if self._client:
            import contextlib

            with contextlib.suppress(Exception):
                self._client.close()
            self._client = None

    @property
    def total_inserted(self) -> int:
        return self._total_inserted


register_sink("elasticsearch", ElasticsearchSink)
register_sink("es", ElasticsearchSink)
