"""
Weaviate sink for extractor.

Upserts chunk records as Weaviate objects.
Like QdrantSink, this stores text payload — embeddings populated separately
or via Weaviate's vectorizer module (configure in Weaviate schema).

Auth (env vars):
    EXTRACTOR_WEAVIATE_URL      (default: http://localhost:8080)
    EXTRACTOR_WEAVIATE_API_KEY  (Weaviate Cloud / WCS API key)

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

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__)


class WeaviateSink(BaseSink):
    def __init__(
        self,
        url: str | None = None,
        class_name: str = "ExtractorChunk",
        batch_size: int = 100,
    ) -> None:
        self.url = url or os.environ.get("EXTRACTOR_WEAVIATE_URL", "http://localhost:8080")
        self.class_name = class_name
        self.batch_size = batch_size
        self._client = None
        self._buffer: list[dict[str, Any]] = []
        self._total_inserted = 0

    def connect(self) -> None:
        try:
            import weaviate
        except ImportError as e:
            raise ImportError(
                "weaviate-client not installed. Run: pip install 'coresdk-extractor[weaviate]'"
            ) from e
        api_key = os.environ.get("EXTRACTOR_WEAVIATE_API_KEY")
        auth = weaviate.auth.AuthApiKey(api_key) if api_key else None
        self._client = weaviate.Client(self.url, auth_client_secret=auth)
        # Create class if not exists
        schema = self._client.schema.get()
        existing = {c["class"] for c in schema.get("classes", [])}
        if self.class_name not in existing:
            self._client.schema.create_class(
                {
                    "class": self.class_name,
                    "properties": [
                        {"name": "extractor_id", "dataType": ["text"]},
                        {"name": "element_type", "dataType": ["text"]},
                        {"name": "text", "dataType": ["text"]},
                        {"name": "source_filename", "dataType": ["text"]},
                        {"name": "source_sha256", "dataType": ["text"]},
                        {"name": "source_uri", "dataType": ["text"]},
                        {"name": "doc_category", "dataType": ["text"]},
                        {"name": "doc_category_score", "dataType": ["number"]},
                        {"name": "page", "dataType": ["int"]},
                        {"name": "token_count", "dataType": ["int"]},
                        {"name": "entity_types", "dataType": ["text[]"]},
                        {"name": "relation_subjects", "dataType": ["text[]"]},
                        {"name": "relation_objects", "dataType": ["text[]"]},
                    ],
                }
            )
        logger.info("WeaviateSink connected: %s class=%s", self.url, self.class_name)

    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
        with self._client.batch as batch:
            batch.batch_size = len(records)
            for r in records:
                props = {
                    "extractor_id": r.get("id", ""),
                    "element_type": r.get("element_type", ""),
                    "text": r.get("text", ""),
                    "source_filename": r.get("source_filename", ""),
                    "source_sha256": r.get("source_sha256", ""),
                    "source_uri": r.get("source_uri") or "",
                    "doc_category": r.get("doc_category") or "",
                    "doc_category_score": float(r.get("doc_category_score") or 0.0),
                    "page": int(r.get("page") or 0),
                    "token_count": int(r.get("token_count") or 0),
                    "entity_types": (r.get("chunk_metadata") or {}).get("entity_types") or [],
                    "relation_subjects": r.get("relation_subjects") or [],
                    "relation_objects": r.get("relation_objects") or [],
                }
                batch.add_data_object(props, self.class_name)
        self._total_inserted += len(records)

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

    def close(self) -> None:
        self._client = None

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


register_sink("weaviate", WeaviateSink)
