from typing import Any

from django.db.models import Q

from apps.wpm.models import BatchLog, BatchSt
from mcp_server.context import authenticated_factory_user
from mcp_server.tools.common import json_safe_result, validate_result_size


def _base_queryset():
    return BatchSt.objects.select_related("material_start")


def _batch_summary(batch_stat: BatchSt) -> dict[str, Any]:
    material = batch_stat.material_start
    return {
        "id": str(batch_stat.id),
        "batch": batch_stat.batch,
        "version": batch_stat.version,
        "zt_batch": batch_stat.zt_batch,
        "first_time": batch_stat.first_time,
        "last_time": batch_stat.last_time,
        "material_start": (
            {
                "id": str(material.id),
                "name": material.name,
                "model": material.model,
                "specification": material.specification,
            }
            if material
            else None
        ),
        "data_keys": sorted((batch_stat.data or {}).keys()),
        "update_time": batch_stat.update_time,
    }


def search_batch_stats(
    query: str = "",
    zt_batch: str = "",
    material_id: str | None = None,
    version: int | None = 1,
    limit: int = 20,
) -> dict[str, Any]:
    """搜索批次统计；摘要仅返回数据分组名称，不返回完整统计数据。"""
    authenticated_factory_user()
    safe_limit = max(1, min(limit, 100))
    queryset = _base_queryset()
    if query.strip():
        queryset = queryset.filter(batch__icontains=query.strip())
    if zt_batch.strip():
        queryset = queryset.filter(zt_batch=zt_batch.strip())
    if material_id:
        queryset = queryset.filter(material_start_id=material_id)
    if version is not None:
        queryset = queryset.filter(version=version)
    items = [
        _batch_summary(item)
        for item in queryset.order_by("batch", "version")[:safe_limit]
    ]
    result = json_safe_result({"items": items, "limit": safe_limit})
    validate_result_size(result)
    return result


def get_batch_stat(
    batch: str,
    version: int = 1,
    include_relations: bool = True,
) -> dict[str, Any]:
    """读取指定批次版本的完整统计数据，并可附带直接拆合批关系。"""
    authenticated_factory_user()
    try:
        batch_stat = _base_queryset().get(batch=batch, version=version)
    except BatchSt.DoesNotExist as exc:
        raise ValueError(f"未找到批次统计：{batch} v{version}") from exc

    result = _batch_summary(batch_stat)
    result["data"] = batch_stat.data
    if include_relations:
        result["relations"] = list(
            BatchLog.objects.filter(Q(source=batch_stat) | Q(target=batch_stat))
            .select_related("source", "target")
            .values(
                "id",
                "relation_type",
                "source_id",
                "source__batch",
                "source__version",
                "target_id",
                "target__batch",
                "target__version",
                "handover_id",
                "mlog_id",
            )
        )
    result = json_safe_result(result)
    validate_result_size(result)
    return result


def register_batch_stat_tools(server) -> None:
    server.tool()(search_batch_stats)
    server.tool()(get_batch_stat)