"""Utilities for accumulating RAG source citations across multiple searches."""

import json
from typing import Any, Dict


def _citation_item_key(item: Any) -> str:
    """Build a stable identity for one RAG citation chunk."""
    if not isinstance(item, dict):
        return json.dumps(item, ensure_ascii=False, sort_keys=True, default=str)

    chunk_id = item.get("id")
    if chunk_id not in (None, ""):
        return f"id:{chunk_id}"

    location = {
        "file_id": item.get("file_id"),
        "page_idx": item.get("page_idx"),
        "positions": item.get("positions"),
    }
    if any(value not in (None, "", []) for value in location.values()):
        return "location:" + json.dumps(
            location, ensure_ascii=False, sort_keys=True, default=str
        )

    # index/score may differ between rewritten queries and do not identify a source.
    fallback = {
        key: value
        for key, value in item.items()
        if key not in {"index", "score", "text"}
    }
    return "fallback:" + json.dumps(
        fallback, ensure_ascii=False, sort_keys=True, default=str
    )


def merge_rag_source_citations(
    accumulated: Dict[str, Any], incoming: Any
) -> Dict[str, Any]:
    """
    Append a RAG call's citations without losing earlier calls or duplicating chunks.

    Citation text is intentionally removed here because the stream response has never
    exposed it; the full text remains available inside the workflow for answer generation.
    """
    if not isinstance(incoming, dict) or not incoming:
        return accumulated

    for doc_name, chunks in incoming.items():
        if isinstance(chunks, list):
            target = accumulated.setdefault(doc_name, [])
            if not isinstance(target, list):
                # Be defensive about malformed/legacy data while preserving both values.
                target = [target]
                accumulated[doc_name] = target

            seen = {_citation_item_key(item) for item in target}
            for item in chunks:
                sanitized_item = (
                    {key: value for key, value in item.items() if key != "text"}
                    if isinstance(item, dict)
                    else item
                )
                item_key = _citation_item_key(sanitized_item)
                if item_key not in seen:
                    target.append(sanitized_item)
                    seen.add(item_key)
        elif doc_name not in accumulated:
            accumulated[doc_name] = chunks

    return accumulated