"""Milvus client helper for Airflow plugin endpoints.

Provides a lazy-initialized MilvusClient singleton and helper functions
for partial upsert and querying chunks by domain term.

NOTE: This module runs inside the Airflow api-server process. Heavy src.*
imports (DB engines, settings) are still forbidden, but ``src.resilience.breakers``
is import-time safe — it only declares pybreaker + prometheus_client objects.
"""

import logging
import os

from src.resilience.breakers import breaker_call, milvus_breaker

log = logging.getLogger(__name__)

_client = None

VERIFIED_TAG_NAMES_MAX_CAPACITY = 64


def project_verified_tag_names(tags, chunk_id: str | None = None) -> list[str]:
    """Project the flat list of names for tags whose status is 'verified'.

    Accepts list[dict] (HITL / dynamic-field path) with keys 'name' and 'status'.
    Any other shape (incl. legacy list[str] without status info) yields an empty
    list — those rows re-populate on the next HITL mutation or re-ingest.

    Input order is preserved. Cap at VERIFIED_TAG_NAMES_MAX_CAPACITY (64); on
    overflow, truncate and log a warning.
    """
    if not tags:
        return []
    verified: list[str] = []
    for t in tags:
        if not isinstance(t, dict):
            continue
        if t.get("status") != "verified":
            continue
        name = t.get("name")
        if not name:
            continue
        verified.append(name)
    if len(verified) > VERIFIED_TAG_NAMES_MAX_CAPACITY:
        dropped = len(verified) - VERIFIED_TAG_NAMES_MAX_CAPACITY
        log.warning(
            "verified_tag_names_truncated chunk_id=%s dropped=%d",
            chunk_id, dropped,
        )
        verified = verified[:VERIFIED_TAG_NAMES_MAX_CAPACITY]
    return verified


def get_collection() -> str:
    """Return the Milvus collection name from env."""
    return os.getenv("MILVUS_COLLECTION", "textiq_records")


def get_client():
    """Lazy-initialized MilvusClient singleton.

    Imports pymilvus at call time so the module can be loaded in
    environments where pymilvus is not installed (e.g. plugin discovery).
    """
    global _client
    if _client is None:
        from pymilvus import MilvusClient

        host = os.getenv("MILVUS_HOST", "localhost")
        port = os.getenv("MILVUS_PORT", "19530")
        user = os.getenv("MILVUS_USER", "")
        password = os.getenv("MILVUS_PASSWORD", "")
        uri = f"http://{host}:{port}"
        _client = MilvusClient(uri=uri, user=user, password=password)
        log.info("Milvus client initialized: %s:%s", host, port)
    return _client


def upsert_chunk_field(chunk_id: str, timeout: float = 10.0, **fields) -> None:
    """Partial upsert: update only the specified fields on a chunk.

    Uses Milvus merge mode (partial_update=True, requires server >= 2.6.2).
    Only the PK and fields to update are sent; all other fields are preserved.

    Skips the upsert if the chunk does not exist in Milvus (e.g. non-leaf
    nodes are stored in the docstore but not in the vector index).
    """
    client = get_client()
    collection = get_collection()
    existing = breaker_call(
        milvus_breaker,
        client.get,
        collection_name=collection,
        ids=[chunk_id],
        output_fields=["id"],
    )
    if not existing:
        log.info("Chunk %s not found in Milvus, skipping upsert", chunk_id)
        return
    batch_upsert_chunks([{"id": chunk_id, **fields}], timeout=timeout)


def batch_upsert_chunks(data: list[dict], timeout: float = 10.0) -> None:
    """Batch partial upsert: update fields on multiple chunks in one call.

    Each dict must include 'id' (PK) and the fields to update.
    All dicts must have the same set of fields.
    Requires pymilvus >= 2.6.11 and Milvus server >= 2.6.2.
    """
    if not data:
        return
    client = get_client()
    collection = get_collection()
    breaker_call(
        milvus_breaker,
        client.upsert,
        collection_name=collection,
        data=data,
        partial_update=True,
        timeout=timeout,
    )


def _escape_milvus_str(value: str) -> str:
    """Escape double quotes and backslashes for Milvus filter expressions."""
    return value.replace("\\", "\\\\").replace('"', '\\"')


def delete_chunks_by_document_id(
    document_id: str, tenant_id: str, node_id: str,
) -> None:
    """Delete all Milvus vectors for a document, scoped by tenant and node.

    Uses filter-based delete to avoid default query limit truncation.
    """
    client = get_client()
    collection = get_collection()
    safe_doc_id = _escape_milvus_str(document_id)
    safe_tenant = _escape_milvus_str(tenant_id)
    safe_node = _escape_milvus_str(node_id)
    filter_expr = (
        f'document_id == "{safe_doc_id}"'
        f' and tenant_id == "{safe_tenant}"'
        f' and node_id == "{safe_node}"'
    )
    breaker_call(
        milvus_breaker,
        client.delete,
        collection_name=collection,
        filter=filter_expr,
    )
    log.info(
        "Milvus vectors deleted: document_id=%s, tenant_id=%s, node_id=%s",
        document_id, tenant_id, node_id,
    )


def query_chunks_with_term(term: str, tenant_id: str) -> list[dict]:
    """Find chunks containing a domain term for cascading updates.

    Returns list of dicts with 'id' and 'domain_terms' fields.
    """
    client = get_client()
    collection = get_collection()
    safe_term = _escape_milvus_str(term)
    safe_tenant = _escape_milvus_str(tenant_id)
    results = breaker_call(
        milvus_breaker,
        client.query,
        collection_name=collection,
        filter=f'array_contains(domain_terms, "{safe_term}") and tenant_id == "{safe_tenant}"',
        output_fields=["id", "domain_terms"],
    )
    return results