src/wandb_mcp_server/weave_api/query_builder.py

"""
Query builder for the Weave API.

This module provides utilities for constructing query expressions for the Weave API.
It separates the complex query construction logic from the API client.
"""

import calendar
from datetime import datetime
from typing import Any, Dict, Optional, Tuple, Union

# Import the query models for building complex queries
from weave.trace_server.interface.query import (
    AndOperation,
    ContainsOperation,
    ContainsSpec,
    ConvertOperation,
    EqOperation,
    GetFieldOperator,
    GteOperation,
    GtOperation,
    LiteralOperation,
    NotOperation,
    Query,
)

from wandb_mcp_server.weave_api.models import (
    FilterOperator,
    QueryFilter,
    QueryParams,
)
from wandb_mcp_server.utils import get_rich_logger

logger = get_rich_logger(__name__)


class QueryBuilder:
    """Builds query expressions for the Weave API."""

    # Define cost fields as a class constant
    COST_FIELDS = {"total_cost", "completion_cost", "prompt_cost"}

    # Define synthetic fields that need special handling
    SYNTHETIC_FIELDS = {"status", "latency_ms"}

    @staticmethod
    def datetime_to_timestamp(dt_str: str) -> int:
        """Convert an ISO format datetime string to Unix timestamp.

        Args:
            dt_str: The ISO format datetime string. Handles 'Z' suffix for UTC.

        Returns:
            The Unix timestamp (seconds since epoch). Returns 0 if the input is empty or parsing fails.
        """
        if not dt_str:
            return 0

        # Handle 'Z' suffix for UTC
        dt_str = dt_str.replace("Z", "+00:00")

        try:
            dt = datetime.fromisoformat(dt_str)
            return int(calendar.timegm(dt.utctimetuple()))
        except ValueError:
            # If parsing fails, return 0 (beginning of epoch)
            logger.warning(f"Failed to parse datetime string: {dt_str}")
            return 0

    @classmethod
    def create_comparison_operation(
        cls, field_name: str, operator: FilterOperator, value: Any
    ) -> Optional[Union[EqOperation, GtOperation, GteOperation, NotOperation]]:
        """Create a comparison operation for a field.

        Args:
            field_name: The name of the field to compare.
            operator: The comparison operator.
            value: The value to compare against.

        Returns:
            A comparison operation, or None if the operation could not be created.
        """
        try:
            field_op_base = GetFieldOperator(**{"$getField": field_name})
            field_op = field_op_base  # Default to no $convert

            # Apply $convert selectively as it was found to be needed for some attributes
            if field_name.startswith("attributes."):
                # For general attributes, convert to double if comparing with a number
                if isinstance(value, (int, float)):
                    field_op = ConvertOperation(
                        **{"$convert": {"input": field_op_base, "to": "double"}}
                    )

            literal_op = LiteralOperation(**{"$literal": value})
        except Exception as e:
            logger.warning(
                f"Invalid value for {field_name} comparison {operator}: {value}. Error: {e}"
            )
            return None

        if operator == FilterOperator.GREATER_THAN:
            return GtOperation(**{"$gt": (field_op, literal_op)})
        elif operator == FilterOperator.GREATER_THAN_EQUAL:
            return GteOperation(**{"$gte": (field_op, literal_op)})
        elif operator == FilterOperator.EQUALS:
            return EqOperation(**{"$eq": (field_op, literal_op)})
        elif operator == FilterOperator.LESS_THAN:  # Implement $lt as $not($gte)
            gte_op = GteOperation(**{"$gte": (field_op, literal_op)})
            return NotOperation(**{"$not": [gte_op]})
        elif operator == FilterOperator.LESS_THAN_EQUAL:  # Implement $lte as $not($gt)
            gt_op = GtOperation(**{"$gt": (field_op, literal_op)})
            return NotOperation(**{"$not": [gt_op]})
        else:
            logger.warning(
                f"Unsupported comparison operator '{operator}' for {field_name}"
            )
            return None

    @classmethod
    def create_contains_operation(
        cls, field_name: str, substring: str, case_insensitive: bool = True
    ) -> ContainsOperation:
        """Create a contains operation for a field.

        Args:
            field_name: The name of the field to check.
            substring: The substring to look for.
            case_insensitive: Whether the comparison should be case-insensitive.

        Returns:
            A contains operation.
        """
        return ContainsOperation(
            **{
                "$contains": ContainsSpec(
                    input=GetFieldOperator(**{"$getField": field_name}),
                    substr=LiteralOperation(**{"$literal": substring}),
                    case_insensitive=case_insensitive,
                )
            }
        )

    @classmethod
    def build_query_expression(cls, filters: Dict[str, Any]) -> Optional[Query]:
        """Build a Query expression from the filter dictionary.

        Args:
            filters: Dictionary of filter conditions.

        Returns:
            The constructed Query object, or None if no valid filters were provided.
        """
        operations = []

        # Handle op_name filter (regex or string)
        if "op_name" in filters:
            op_name = filters["op_name"]
            if isinstance(op_name, str):
                # If it's a string with wildcard pattern, treat as contains
                if "*" in op_name or ".*" in op_name:
                    # Extract the part between wildcards
                    pattern = op_name.replace("*", "").replace(".*", "")
                    operations.append(cls.create_contains_operation("op_name", pattern))
                else:
                    # Exact match
                    operations.append(
                        EqOperation(
                            **{
                                "$eq": (
                                    GetFieldOperator(**{"$getField": "op_name"}),
                                    LiteralOperation(**{"$literal": op_name}),
                                )
                            }
                        )
                    )
            elif hasattr(op_name, "pattern"):  # Regex pattern
                operations.append(
                    cls.create_contains_operation("op_name", op_name.pattern)
                )

        # Handle op_name_contains custom filter (for simple substring matching)
        if "op_name_contains" in filters:
            substring = filters["op_name_contains"]
            operations.append(cls.create_contains_operation("op_name", substring))

        # Handle display_name filter (regex or string)
        if "display_name" in filters:
            display_name = filters["display_name"]
            if isinstance(display_name, str):
                # If it's a string with wildcard pattern, treat as contains
                if "*" in display_name or ".*" in display_name:
                    # Extract the part between wildcards
                    pattern = display_name.replace("*", "").replace(".*", "")
                    operations.append(
                        cls.create_contains_operation("display_name", pattern)
                    )
                else:
                    # Exact match
                    operations.append(
                        EqOperation(
                            **{
                                "$eq": (
                                    GetFieldOperator(**{"$getField": "display_name"}),
                                    LiteralOperation(**{"$literal": display_name}),
                                )
                            }
                        )
                    )
            elif hasattr(display_name, "pattern"):  # Regex pattern
                operations.append(
                    cls.create_contains_operation("display_name", display_name.pattern)
                )

        # Handle display_name_contains custom filter (for simple substring matching)
        if "display_name_contains" in filters:
            substring = filters["display_name_contains"]
            operations.append(cls.create_contains_operation("display_name", substring))

        # Handle status filter based on summary.weave.status using dot notation
        if "status" in filters:
            target_status = filters["status"]
            if isinstance(target_status, str):
                comp_op = cls.create_comparison_operation(
                    "summary.weave.status", FilterOperator.EQUALS, target_status.lower()
                )
                if comp_op:
                    operations.append(comp_op)
            else:
                logger.warning(
                    f"Invalid status filter value: {target_status}. Expected a string."
                )

        # Handle time range filter (convert ISO datetime strings to Unix seconds)
        if "time_range" in filters:
            time_range = filters["time_range"]

            # >= start
            if "start" in time_range and time_range["start"]:
                start_ts = cls.datetime_to_timestamp(time_range["start"])
                if start_ts > 0:
                    comp_op = cls.create_comparison_operation(
                        "started_at", FilterOperator.GREATER_THAN_EQUAL, start_ts
                    )
                    if comp_op:
                        operations.append(comp_op)

            # < end (i.e. started_at strictly before end_ts)
            if "end" in time_range and time_range["end"]:
                end_ts = cls.datetime_to_timestamp(time_range["end"])
                if end_ts > 0:
                    comp_op = cls.create_comparison_operation(
                        "started_at", FilterOperator.LESS_THAN, end_ts
                    )
                    if comp_op:
                        operations.append(comp_op)

        # Handle wb_run_id filter (top-level)
        if "wb_run_id" in filters:
            run_id = filters["wb_run_id"]
            # This filter expects a string for wb_run_id and uses $contains or $eq.
            if isinstance(run_id, str):
                if (
                    "$contains" in run_id or "*" in run_id
                ):  # Simple check for contains style
                    pattern = run_id.replace("$contains:", "").replace(
                        "*", ""
                    )  # Basic cleanup
                    operations.append(
                        cls.create_contains_operation("wb_run_id", pattern.strip())
                    )
                else:
                    operations.append(
                        EqOperation(
                            **{
                                "$eq": (
                                    GetFieldOperator(**{"$getField": "wb_run_id"}),
                                    LiteralOperation(**{"$literal": run_id}),
                                )
                            }
                        )
                    )
            elif (
                isinstance(run_id, dict) and "$contains" in run_id
            ):  # wb_run_id: {"$contains": "foo"}
                pattern = run_id["$contains"]
                if isinstance(pattern, str):
                    operations.append(
                        cls.create_contains_operation("wb_run_id", pattern)
                    )
                else:
                    logger.warning(
                        f"Invalid $contains value for wb_run_id: {pattern}. Expected string."
                    )
            else:
                logger.warning(
                    f"Invalid wb_run_id filter value: {run_id}. Expected a string or dict with $contains."
                )

        # Handle latency filter based on summary.weave.latency_ms
        if "latency" in filters:
            latency_filter = filters["latency"]
            # Extract the operator and value
            if isinstance(latency_filter, dict) and len(latency_filter) == 1:
                op_key, value = next(iter(latency_filter.items()))
                try:
                    op = FilterOperator(op_key)
                    comp_op = cls.create_comparison_operation(
                        "summary.weave.latency_ms", op, value
                    )
                    if comp_op:
                        operations.append(comp_op)
                except (ValueError, KeyError):
                    logger.warning(f"Invalid operator for latency filter: {op_key}")
            else:
                logger.warning(
                    f"Invalid format for latency filter: {latency_filter}. Expected a dict with one operator key."
                )

        # Handle attributes filter using dot notation AND supporting comparison operators
        if "attributes" in filters:
            attributes_filters = filters["attributes"]
            if isinstance(attributes_filters, dict):
                for attr_path, attr_value_or_op in attributes_filters.items():
                    full_attr_path = f"attributes.{attr_path}"

                    # Check if the value is a comparison operator dict or a literal
                    if (
                        isinstance(attr_value_or_op, dict)
                        and len(attr_value_or_op) == 1
                        and next(iter(attr_value_or_op.keys()))
                        in [op.value for op in FilterOperator]
                    ):
                        # It's a comparison operation
                        op_key, value = next(iter(attr_value_or_op.items()))
                        try:
                            op = FilterOperator(op_key)
                            comp_op = cls.create_comparison_operation(
                                full_attr_path, op, value
                            )
                            if comp_op:
                                operations.append(comp_op)
                        except (ValueError, KeyError):
                            logger.warning(
                                f"Invalid operator for attribute filter: {op_key}"
                            )
                    elif (
                        isinstance(attr_value_or_op, dict)
                        and "$contains" in attr_value_or_op
                    ):
                        # It's a contains operation
                        if isinstance(attr_value_or_op["$contains"], str):
                            operations.append(
                                cls.create_contains_operation(
                                    full_attr_path, attr_value_or_op["$contains"]
                                )
                            )
                        else:
                            logger.warning(
                                f"Invalid value for $contains on {full_attr_path}: {attr_value_or_op['$contains']}. Expected string."
                            )
                    else:
                        # Assume literal equality
                        comp_op = cls.create_comparison_operation(
                            full_attr_path, FilterOperator.EQUALS, attr_value_or_op
                        )
                        if comp_op:
                            operations.append(comp_op)
            else:
                logger.warning(
                    f"Invalid format for 'attributes' filter: {attributes_filters}. Expected a dictionary."
                )

        # Handle has_exception filter (checking top-level exception field)
        if "has_exception" in filters:
            has_exception = filters["has_exception"]
            # Skip filtering if has_exception is None (show everything)
            if has_exception is not None:
                # Create base operation that checks if exception is None (no exception case)
                base_op = EqOperation(
                    **{
                        "$eq": (
                            GetFieldOperator(**{"$getField": "exception"}),
                            LiteralOperation(**{"$literal": None}),
                        )
                    }
                )

                if has_exception:
                    # For has_exception=True: Negate the operation to get NOT NULL
                    operations.append(NotOperation(**{"$not": [base_op]}))
                else:
                    # For has_exception=False: Use the operation as is
                    operations.append(base_op)

        # Combine all operations with AND
        if operations:
            if len(operations) == 1:
                # Wrap the single operation in the Query model structure
                return Query(**{"$expr": operations[0]})
            else:
                # Wrap the AndOperation in the Query model structure
                and_op = AndOperation(**{"$and": operations})
                return Query(**{"$expr": and_op})

        return None  # No complex filters, so no Query object needed

    @classmethod
    def separate_filters(
        cls, filters: Union[Dict[str, Any], QueryFilter]
    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        """Separate filters into direct and complex filters.

        Args:
            filters: Dictionary or QueryFilter of filter conditions.

        Returns:
            Tuple of (direct_filters, complex_filters).
        """
        if not filters:
            return {}, {}

        # Convert QueryFilter to dict if needed
        if isinstance(filters, QueryFilter):
            filters_dict = {k: v for k, v in filters.__dict__.items() if v is not None}
        else:
            filters_dict = filters

        # Simple filters for CallsFilter object
        direct_filters = {}
        complex_filters = {}

        # Simple filter keys that go directly to the CallsFilter object
        simple_filter_keys = [
            "trace_roots_only",
            "op_names",
            "op_names_prefix",
            "trace_ids",
            "trace_parent_ids",
            "parent_ids",
            "call_ids",
        ]

        for key in simple_filter_keys:
            if key in filters_dict:
                # Ensure op_names, trace_ids, etc. are lists as expected by CallsFilter
                if key in [
                    "op_names",
                    "op_names_prefix",
                    "trace_ids",
                    "trace_parent_ids",
                    "parent_ids",
                    "call_ids",
                ] and not isinstance(filters_dict[key], list):
                    direct_filters[key] = [str(filters_dict[key])]
                else:
                    direct_filters[key] = filters_dict[key]

                # Ensure call_ids are strings
                if key == "call_ids" and isinstance(direct_filters[key], list):
                    direct_filters[key] = [
                        str(call_id) for call_id in direct_filters[key]
                    ]

        # Handle individual op_name and trace_id if op_names/trace_ids not already set
        if "op_name" in filters_dict and "op_names" not in direct_filters:
            # Only add if it's a simple name, not a pattern (patterns go to complex)
            if (
                isinstance(filters_dict["op_name"], str)
                and "*" not in filters_dict["op_name"]
                and ".*" not in filters_dict["op_name"]
            ):
                direct_filters["op_names"] = [filters_dict["op_name"]]
            else:
                # It's a pattern or complex, send to complex_filters
                complex_filters["op_name"] = filters_dict["op_name"]
        elif "op_name" in filters_dict and "op_names" in direct_filters:
            # If op_names is already set, and op_name is a pattern, it needs to go to complex
            if isinstance(filters_dict["op_name"], str) and (
                "*" in filters_dict["op_name"] or ".*" in filters_dict["op_name"]
            ):
                complex_filters["op_name"] = filters_dict["op_name"]

        if "trace_id" in filters_dict and "trace_ids" not in direct_filters:
            direct_filters["trace_ids"] = [str(filters_dict["trace_id"])]

        # All other keys from the original `filters` dict go to complex_filters
        all_handled_direct_keys = set(direct_filters.keys())
        # Add op_name/trace_id to handled if they were processed into op_names/trace_ids
        if (
            "op_names" in direct_filters
            and "op_name" in filters_dict
            and filters_dict["op_name"] in direct_filters["op_names"]
        ):
            all_handled_direct_keys.add("op_name")
        if (
            "trace_ids" in direct_filters
            and "trace_id" in filters_dict
            and filters_dict["trace_id"] in direct_filters["trace_ids"]
        ):
            all_handled_direct_keys.add("trace_id")

        for key, value in filters_dict.items():
            if key not in all_handled_direct_keys:
                # Exception: if op_name was simple and handled, but is also in filters, it shouldn't be added again
                if (
                    key == "op_name"
                    and "op_names" in direct_filters
                    and value in direct_filters["op_names"]
                    and not (isinstance(value, str) and ("*" in value or ".*" in value))
                ):
                    continue
                complex_filters[key] = value

        return direct_filters, complex_filters

    @classmethod
    def prepare_query_params(
        cls, params: Union[QueryParams, Dict[str, Any]]
    ) -> Dict[str, Any]:
        """Prepare query parameters for the Weave API.

        Args:
            params: Query parameters, either as a QueryParams object or a dictionary.

        Returns:
            Dictionary of query parameters ready for the Weave API.
        """
        # Convert QueryParams to dict if needed
        if isinstance(params, QueryParams):
            # Convert to dict (simplified for illustration)
            raw_params = {
                "entity_name": params.entity_name,
                "project_name": params.project_name,
                "filters": params.filters,
                "sort_by": params.sort_by,
                "sort_direction": params.sort_direction,
                "limit": params.limit,
                "offset": params.offset,
                "include_costs": params.include_costs,
                "include_feedback": params.include_feedback,
                "columns": params.columns,
                "expand_columns": params.expand_columns,
            }
        else:
            raw_params = params.copy()

        # Extract filters
        filters = raw_params.get("filters", {})

        # Separate filters
        direct_filters, complex_filters = cls.separate_filters(filters)

        # Prepare request body
        request_body = {
            "project_id": f"{raw_params['entity_name']}/{raw_params['project_name']}",
            "include_costs": raw_params.get("include_costs", True),
            "include_feedback": raw_params.get("include_feedback", True),
        }

        # Add filter if present
        if direct_filters:
            request_body["filter"] = direct_filters

        # Add query expression if present
        query_expression = cls.build_query_expression(complex_filters)
        if query_expression:
            request_body["query"] = query_expression.model_dump(by_alias=True)

        # Add sort criteria if present, but never send cost fields to server
        sort_by = raw_params.get("sort_by")
        if sort_by and sort_by not in cls.COST_FIELDS:
            request_body["sort_by"] = [
                {
                    "field": sort_by,
                    "direction": raw_params.get("sort_direction", "desc"),
                }
            ]

        # Add pagination parameters if present
        if "limit" in raw_params and raw_params["limit"] is not None:
            request_body["limit"] = raw_params["limit"]
        if "offset" in raw_params and raw_params["offset"] is not None:
            request_body["offset"] = raw_params["offset"]

        # Process columns, filtering out synthetic fields
        if "columns" in raw_params and raw_params["columns"]:
            original_columns = raw_params["columns"]

            # Store synthetic fields we need to generate later
            synthetic_fields_to_add = []

            # Filter out synthetic fields from the API request
            filtered_columns = []
            for col in original_columns:
                if col in cls.SYNTHETIC_FIELDS:
                    synthetic_fields_to_add.append(col)
                else:
                    filtered_columns.append(col)

            # If we need to synthesize status, ensure we request summary field
            if (
                "status" in synthetic_fields_to_add
                and "summary" not in filtered_columns
            ):
                filtered_columns.append("summary")

            # If we need to synthesize latency_ms, ensure we request summary field
            if (
                "latency_ms" in synthetic_fields_to_add
                and "summary" not in filtered_columns
            ):
                filtered_columns.append("summary")

            # Only send filtered columns to the API
            if filtered_columns:
                request_body["columns"] = filtered_columns

            # Store the synthetic fields that need to be processed
            if synthetic_fields_to_add:
                request_body["_synthetic_fields"] = synthetic_fields_to_add

        if "expand_columns" in raw_params and raw_params["expand_columns"]:
            request_body["expand_columns"] = raw_params["expand_columns"]

        return request_body