"""构造并 upsert NL2SQL 元数据行。"""
from __future__ import annotations

import time
from collections.abc import Awaitable, Callable
from typing import Any

from nl2sql.embedding import embed_texts
from nl2sql.metadata import build_metadata_chunks
from nl2sql.milvus_collections import NL2SQL_COLLECTION


def _escape_filter_value(value: str) -> str:
    return value.replace("\\", "\\\\").replace('"', '\\"')


async def prepare_metadata_rows(
    chunks: list[dict[str, Any]],
    *,
    embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
    timestamp: int | None = None,
) -> list[dict[str, Any]]:
    if not chunks:
        return []
    vectors = await embedder([chunk["text"] for chunk in chunks])
    if len(vectors) != len(chunks):
        raise ValueError("Embedding result count must match metadata chunk count")
    created_at = int(time.time()) if timestamp is None else timestamp
    return [
        {**chunk, "vector": vector, "created_at": created_at}
        for chunk, vector in zip(chunks, vectors)
    ]


async def sync_metadata(
    milvus_client,
    table_rows: list[dict[str, Any]],
    column_rows: list[dict[str, Any]],
    *,
    embedder: Callable[[list[str]], Awaitable[list[list[float]]]] = embed_texts,
    timestamp: int | None = None,
) -> int:
    chunks = build_metadata_chunks(table_rows, column_rows)
    chunks_by_table: dict[str, list[dict[str, Any]]] = {}
    for chunk in chunks:
        chunks_by_table.setdefault(chunk["table_name"], []).append(chunk)

    updated_count = 0
    for table_name, table_chunks in chunks_by_table.items():
        escaped_name = _escape_filter_value(table_name)
        existing_rows = await milvus_client.query(
            collection_name=NL2SQL_COLLECTION,
            filter=f'table_name == "{escaped_name}"',
            output_fields=["id", "content_hash", "is_valid", "is_deprecated"],
        )
        desired_signature = {
            (chunk["id"], chunk["content_hash"], chunk["is_valid"], chunk["is_deprecated"])
            for chunk in table_chunks
        }
        existing_signature = {
            (
                row.get("id"),
                row.get("content_hash"),
                row.get("is_valid", True),
                row.get("is_deprecated", False),
            )
            for row in existing_rows or []
        }
        if desired_signature == existing_signature:
            continue

        if existing_rows:
            await milvus_client.delete(
                collection_name=NL2SQL_COLLECTION,
                filter=f'table_name == "{escaped_name}"',
            )
        rows = await prepare_metadata_rows(
            table_chunks, embedder=embedder, timestamp=timestamp
        )
        if rows:
            await milvus_client.upsert(collection_name=NL2SQL_COLLECTION, data=rows)
            updated_count += len(rows)
    return updated_count