"""投顾 Agent 的 NL2SQL 数据查询适配层。"""
from __future__ import annotations

from dataclasses import asdict
from typing import Any

from agent.advisor_agent.auth import ensure_customer_access
from agent.data_query.agent import DataQueryAgent
from config import database
from config.settings import settings
from nl2sql.contracts import DataQueryRequest, DataQueryResult
from nl2sql.retrieval import retrieve_metadata
from nl2sql.runtime_config import runtime_config
from nl2sql.schema import load_authoritative_schema
from service.nl2sql.permission_service import load_query_permission
from service.nl2sql.query_service import QueryServiceError
from tool.llm import llm as default_llm


async def execute_advisor_data_query(
    db,
    *,
    advisor_id: int,
    customer_id: int,
    question: str,
    trace_id: str,
    session_id: str | None = None,
    data_scope: dict[str, Any] | None = None,
    max_rows: int | None = None,
    page: int = 1,
    page_size: int = 100,
    sort_by: str | None = None,
    sort_order: str = "asc",
    milvus=None,
    redis=None,
    llm_client=None,
    query_agent=None,
) -> dict[str, Any]:
    """在当前投顾和选中客户范围内执行只读自然语言查询。

    ``data_scope`` 即使由调用方传入也不会被信任，服务端始终覆盖为当前
    ``customer_id``，避免投顾借助 NL2SQL 查询其他客户数据。
    """
    await ensure_customer_access(db, advisor_id=advisor_id, customer_id=customer_id)
    permission = await load_query_permission(db, advisor_id)
    if not permission.get("can_query", False):
        raise QueryServiceError("当前投顾没有 NL2SQL 查询权限")

    milvus = milvus or database.milvus.client()
    redis = redis or database.redis.client()
    llm_client = llm_client or default_llm
    request = DataQueryRequest(
        question=question,
        user_id=advisor_id,
        trace_id=trace_id,
        session_id=session_id,
        caller_agent="advisor_agent",
        data_scope={"customer_ids": [customer_id]},
        max_rows=min(max_rows or runtime_config.max_rows, runtime_config.max_rows),
        include_sql=False,
        page=page,
        page_size=page_size,
        sort_by=sort_by,
        sort_order=sort_order,
    )

    async def permission_loader(_user_id: int):
        return permission

    async def metadata_retriever(query: str):
        return await retrieve_metadata(
            query,
            milvus,
            top_k=runtime_config.retrieval_top_k,
        )

    async def schema_loader(table_names: set[str], _permission: dict):
        return await load_authoritative_schema(
            db,
            database=settings.mysql.database,
            candidate_tables=table_names,
        )

    result: DataQueryResult = await (query_agent or DataQueryAgent()).query(
        request,
        session=db,
        permission_loader=permission_loader,
        metadata_retriever=metadata_retriever,
        schema_loader=schema_loader,
        llm_client=llm_client,
        summary_llm=llm_client,
        masks=permission.get("masks"),
        redis=redis,
    )
    payload = asdict(result)
    payload["sql"] = None
    payload["customer_id"] = customer_id
    return payload