"""读取当前 MySQL 元数据并同步到 Milvus。"""
from __future__ import annotations

import asyncio
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from sqlalchemy import text

from config.database import milvus as milvus_db
from config.database import mysql
from config.database.milvus import client as milvus_client
from config.database.mysql import get_session_factory
from config.settings import settings
from nl2sql.metadata_sync import sync_metadata


TABLES_SQL = text(
    """
    SELECT TABLE_NAME, TABLE_COMMENT, TABLE_TYPE
    FROM information_schema.tables
    WHERE TABLE_SCHEMA = :database
    """
)

COLUMNS_SQL = text(
    """
    SELECT TABLE_NAME, COLUMN_NAME, COLUMN_COMMENT, DATA_TYPE,
           IS_NULLABLE, ORDINAL_POSITION
    FROM information_schema.columns
    WHERE TABLE_SCHEMA = :database
    ORDER BY TABLE_NAME, ORDINAL_POSITION
    """
)


async def load_information_schema() -> tuple[list[dict], list[dict]]:
    async with get_session_factory()() as session:
        tables = [
            dict(row)
            for row in (await session.execute(TABLES_SQL, {"database": settings.mysql.database})).mappings()
        ]
        columns = [
            dict(row)
            for row in (await session.execute(COLUMNS_SQL, {"database": settings.mysql.database})).mappings()
        ]
    return tables, columns


async def synchronize() -> int:
    try:
        tables, columns = await load_information_schema()
        return await sync_metadata(milvus_client(), tables, columns)
    finally:
        await mysql.dispose()
        await milvus_db.dispose()


if __name__ == "__main__":
    print(f"upserted {asyncio.run(synchronize())} NL2SQL metadata chunks")