"""OpenSearch document indexer implementation."""

from __future__ import annotations

import asyncio
import hashlib
import logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Sequence

# Note: Callable is still needed for on_batch_complete callback type

from gnosisllm_knowledge.backends.opensearch.config import OpenSearchConfig
from gnosisllm_knowledge.backends.opensearch.mappings import (
    get_knowledge_index_mappings,
    get_knowledge_index_settings,
)
from gnosisllm_knowledge.core.domain.document import Document
from gnosisllm_knowledge.core.domain.result import BatchResult, IndexResult
from gnosisllm_knowledge.core.exceptions import IndexError

if TYPE_CHECKING:
    from opensearchpy import AsyncOpenSearch

logger = logging.getLogger(__name__)


class OpenSearchIndexer:
    """OpenSearch document indexer.

    Implements the IDocumentIndexer protocol for indexing documents
    into OpenSearch with k-NN vector embeddings.

    Example:
        ```python
        config = OpenSearchConfig.from_env()
        client = AsyncOpenSearch(hosts=[config.url])
        indexer = OpenSearchIndexer(client, config)

        # Index a single document
        result = await indexer.index(document, "my-index")

        # Bulk index documents
        result = await indexer.bulk_index(documents, "my-index")
        ```
    """

    def __init__(
        self,
        client: AsyncOpenSearch,
        config: OpenSearchConfig,
    ) -> None:
        """Initialize the indexer.

        Args:
            client: OpenSearch async client.
            config: OpenSearch configuration.

        Note:
            Embeddings are generated automatically by OpenSearch ingest pipeline.
            Run 'gnosisllm-knowledge setup' to configure the ML model and pipeline.
        """
        self._client = client
        self._config = config

    async def index(
        self,
        document: Document,
        index_name: str,
        **options: Any,
    ) -> IndexResult:
        """Index a single document.

        Args:
            document: Document to index.
            index_name: Target index name.
            **options: Additional options (e.g., refresh).

        Returns:
            Index result with success/failure info.

        Raises:
            IndexError: If indexing fails.
        """
        try:
            # Prepare document for indexing
            # Embeddings are generated by OpenSearch ingest pipeline
            doc_body = self._prepare_document(document)

            # Index the document
            refresh = options.get("refresh", False)
            await self._client.index(
                index=index_name,
                id=document.doc_id,
                body=doc_body,
                refresh=refresh,
            )

            return IndexResult(
                success=True,
                document_id=document.doc_id,
                index_name=index_name,
                indexed_count=1,
                failed_count=0,
            )
        except Exception as e:
            logger.error(f"Failed to index document {document.doc_id}: {e}")
            raise IndexError(
                message=f"Failed to index document: {e}",
                details={"document_id": document.doc_id},
                cause=e,
            ) from e

    async def bulk_index(
        self,
        documents: Sequence[Document],
        index_name: str,
        **options: Any,
    ) -> IndexResult:
        """Index multiple documents efficiently.

        Args:
            documents: Documents to index.
            index_name: Target index name.
            **options: Additional options (batch_size, refresh).

        Returns:
            Index result with counts.
        """
        if not documents:
            return IndexResult(success=True, index_name=index_name, indexed_count=0, failed_count=0)

        batch_size = options.get("batch_size", self._config.bulk_batch_size)
        refresh = options.get("refresh", False)

        total_indexed = 0
        total_failed = 0
        errors: list[str] = []

        # Process in batches
        for i in range(0, len(documents), batch_size):
            batch = documents[i : i + batch_size]
            batch_result = await self._index_batch(batch, index_name)

            total_indexed += batch_result.indexed_count
            total_failed += batch_result.failed_count
            if batch_result.errors:
                errors.extend(batch_result.errors)

        # Final refresh if requested
        if refresh:
            await self._client.indices.refresh(index=index_name)

        return IndexResult(
            success=total_failed == 0,
            index_name=index_name,
            indexed_count=total_indexed,
            failed_count=total_failed,
            errors=errors if errors else None,
        )

    async def bulk_index_streaming(
        self,
        documents: AsyncIterator[Document],
        index_name: str,
        *,
        batch_size: int = 500,
        max_concurrent_batches: int = 3,
        on_batch_complete: Callable[[BatchResult], None] | None = None,
    ) -> IndexResult:
        """Stream-index documents with backpressure handling.

        Args:
            documents: Async iterator of documents.
            index_name: Target index name.
            batch_size: Documents per batch.
            max_concurrent_batches: Max concurrent batch operations.
            on_batch_complete: Callback for batch completion.

        Returns:
            Index result with totals.
        """
        semaphore = asyncio.Semaphore(max_concurrent_batches)
        total_indexed = 0
        total_failed = 0
        batch_count = 0
        errors: list[str] = []

        current_batch: list[Document] = []

        async def process_batch(batch: list[Document], batch_num: int) -> IndexResult:
            async with semaphore:
                result = await self._index_batch(batch, index_name)

                if on_batch_complete:
                    batch_result = BatchResult(
                        total=len(batch),
                        succeeded=result.indexed_count,
                        failed=result.failed_count,
                        errors=result.errors or [],
                    )
                    on_batch_complete(batch_result)

                return result

        tasks: list[asyncio.Task[IndexResult]] = []

        async for doc in documents:
            current_batch.append(doc)

            if len(current_batch) >= batch_size:
                batch_count += 1
                task = asyncio.create_task(
                    process_batch(current_batch.copy(), batch_count)
                )
                tasks.append(task)
                current_batch = []

        # Process remaining documents
        if current_batch:
            batch_count += 1
            task = asyncio.create_task(
                process_batch(current_batch.copy(), batch_count)
            )
            tasks.append(task)

        # Wait for all tasks
        results = await asyncio.gather(*tasks)
        for result in results:
            total_indexed += result.indexed_count
            total_failed += result.failed_count
            if result.errors:
                errors.extend(result.errors)

        return IndexResult(
            success=total_failed == 0,
            index_name=index_name,
            indexed_count=total_indexed,
            failed_count=total_failed,
            errors=errors if errors else None,
        )

    async def upsert(
        self,
        document: Document,
        index_name: str,
    ) -> IndexResult:
        """Upsert (update or insert) a document.

        Args:
            document: Document to upsert.
            index_name: Target index name.

        Returns:
            Index result.
        """
        # Embeddings are generated by OpenSearch ingest pipeline
        doc_body = self._prepare_document(document)

        await self._client.index(
            index=index_name,
            id=document.doc_id,
            body=doc_body,
            refresh=False,
        )

        return IndexResult(
            success=True,
            document_id=document.doc_id,
            index_name=index_name,
            indexed_count=1,
            failed_count=0,
        )

    async def delete(
        self,
        doc_id: str,
        index_name: str,
    ) -> bool:
        """Delete a document by ID.

        Args:
            doc_id: Document ID to delete.
            index_name: Index name.

        Returns:
            True if deleted, False if not found.
        """
        try:
            await self._client.delete(index=index_name, id=doc_id)
            return True
        except Exception as e:
            if "not_found" in str(e).lower():
                return False
            raise IndexError(
                message=f"Failed to delete document: {e}",
                details={"document_id": doc_id},
                cause=e,
            ) from e

    async def bulk_delete(
        self,
        doc_ids: Sequence[str],
        index_name: str,
    ) -> int:
        """Delete multiple documents.

        Args:
            doc_ids: Document IDs to delete.
            index_name: Index name.

        Returns:
            Count of deleted documents.
        """
        if not doc_ids:
            return 0

        actions = [
            {"delete": {"_index": index_name, "_id": doc_id}} for doc_id in doc_ids
        ]

        response = await self._client.bulk(body=actions)
        deleted = sum(
            1
            for item in response.get("items", [])
            if item.get("delete", {}).get("result") == "deleted"
        )

        return deleted

    async def delete_by_query(
        self,
        query: dict[str, Any],
        index_name: str,
    ) -> int:
        """Delete documents matching query.

        Args:
            query: OpenSearch query dictionary.
            index_name: Index name.

        Returns:
            Count of deleted documents.
        """
        response = await self._client.delete_by_query(
            index=index_name,
            body=query,
            refresh=True,
        )
        return response.get("deleted", 0)

    async def ensure_index(
        self,
        index_name: str,
        **options: Any,
    ) -> bool:
        """Ensure index exists with correct mapping.

        Args:
            index_name: Index name to create.
            **options: Override settings/mappings.

        Returns:
            True if created, False if already exists.
        """
        exists = await self._client.indices.exists(index=index_name)
        if exists:
            return False

        settings = options.get("settings", get_knowledge_index_settings(self._config))
        mappings = options.get("mappings", get_knowledge_index_mappings(self._config))

        await self._client.indices.create(
            index=index_name,
            body={
                "settings": settings,
                "mappings": mappings,
            },
        )

        logger.info(f"Created index: {index_name}")
        return True

    async def delete_index(self, index_name: str) -> bool:
        """Delete an index.

        Args:
            index_name: Index to delete.

        Returns:
            True if deleted, False if not found.
        """
        try:
            await self._client.indices.delete(index=index_name)
            return True
        except Exception as e:
            if "index_not_found" in str(e).lower():
                return False
            raise

    async def refresh_index(self, index_name: str) -> bool:
        """Refresh index to make documents searchable.

        Args:
            index_name: Index to refresh.

        Returns:
            True on success.
        """
        await self._client.indices.refresh(index=index_name)
        return True

    async def _index_batch(
        self,
        documents: Sequence[Document],
        index_name: str,
    ) -> IndexResult:
        """Index a batch of documents.

        Args:
            documents: Documents to index.
            index_name: Target index.

        Returns:
            Batch index result.
        """
        # Embeddings are generated by OpenSearch ingest pipeline
        actions = []
        for doc in documents:
            doc_body = self._prepare_document(doc)
            actions.append({"index": {"_index": index_name, "_id": doc.doc_id}})
            actions.append(doc_body)

        if not actions:
            return IndexResult(success=True, index_name=index_name, indexed_count=0, failed_count=0)

        response = await self._client.bulk(body=actions)

        indexed = 0
        failed = 0
        errors: list[str] = []

        for item in response.get("items", []):
            index_result = item.get("index", {})
            if index_result.get("status") in (200, 201):
                indexed += 1
            else:
                failed += 1
                error = index_result.get("error", {})
                errors.append(f"{index_result.get('_id')}: {error.get('reason', 'Unknown error')}")

        return IndexResult(
            success=failed == 0,
            index_name=index_name,
            indexed_count=indexed,
            failed_count=failed,
            errors=errors if errors else None,
        )

    def _prepare_document(self, document: Document) -> dict[str, Any]:
        """Prepare document for indexing.

        Args:
            document: Document to prepare.

        Returns:
            Dictionary for indexing.
        """
        # Calculate content hash if not present
        content_hash = document.content_hash
        if not content_hash:
            content_hash = hashlib.sha256(document.content.encode()).hexdigest()

        now = datetime.now(timezone.utc)

        return {
            "id": document.doc_id,
            "content": document.content,
            "url": document.url,
            "title": document.title,
            "source": document.source,
            "account_id": document.account_id,
            "collection_id": document.collection_id,
            "source_id": document.source_id,
            "chunk_index": document.chunk_index,
            "total_chunks": document.total_chunks,
            "parent_doc_id": document.parent_doc_id,
            "quality_score": document.quality_score,
            "language": document.language,
            "content_hash": content_hash,
            "word_count": document.word_count or len(document.content.split()),
            "status": document.status.value,
            "pii_detected": document.pii_detected,
            "pii_redacted": document.pii_redacted,
            "metadata": document.metadata,
            "created_at": document.created_at.isoformat() if document.created_at else now.isoformat(),
            "updated_at": document.updated_at.isoformat() if document.updated_at else None,
            "indexed_at": now.isoformat(),
        }
