"""Query builder for generating initial BTQL queries."""


class QueryBuilder:
    """Generates initial BTQL queries based on experiment ID and metrics."""

    # Metrics that need quoting due to special characters
    SPECIAL_CHAR_METRICS = {
        "clarity & readability",
        "factual accuracy",
    }

    def __init__(self, default_filters: list[str] | None = None) -> None:
        self.default_filters = default_filters or []

    def _quote_metric_name(self, metric: str) -> str:
        """Quote metric names that contain special characters."""
        metric_lower = metric.lower()
        if metric_lower in self.SPECIAL_CHAR_METRICS or " " in metric or "&" in metric:
            return f'"{metric}"'
        return metric

    def _build_filter_clause(self, custom_filters: list[str] | None = None) -> str:
        """Build the filter clause combining default and custom filters."""
        all_filters = list(self.default_filters)
        if custom_filters:
            all_filters.extend(custom_filters)

        if not all_filters:
            return ""

        return " and ".join(all_filters)

    def build_query(
        self,
        experiment_id: str,
        metrics: list[str],
        custom_filters: list[str] | None = None,
    ) -> str:
        """Build a BTQL query to fetch specified metrics from an experiment.

        Args:
            experiment_id: The Braintrust experiment ID
            metrics: List of metric names to fetch
            custom_filters: Additional filter conditions

        Returns:
            A BTQL query string
        """
        # Build filter clause
        filter_clause = self._build_filter_clause(custom_filters)

        # Build the query - fetch scores, metadata, input, output for judge explanations
        query_parts = [
            f"from: experiment('{experiment_id}')",
            "select: root_span_id, scores, metadata, input, output",
        ]

        if filter_clause:
            query_parts.append(f"filter: {filter_clause}")

        return " | ".join(query_parts)

    def build_aggregation_query(
        self,
        experiment_id: str,
        metrics: list[str],
        custom_filters: list[str] | None = None,
    ) -> str:
        """Build a BTQL query that aggregates metrics.

        Args:
            experiment_id: The Braintrust experiment ID
            metrics: List of metric names to fetch
            custom_filters: Additional filter conditions

        Returns:
            A BTQL query string with aggregation
        """
        # Build filter clause
        filter_clause = self._build_filter_clause(custom_filters)

        # Build measure expressions for each metric
        measures = []
        for metric in metrics:
            quoted_name = self._quote_metric_name(metric)
            safe_alias = metric.replace(" ", "_").replace("&", "and")
            measures.append(f"avg(scores.{quoted_name}) as avg_{safe_alias}")

        measures_str = ", ".join(measures)

        # Build the aggregation query
        query_parts = [
            f"from: experiment('{experiment_id}') summary",
            f"measures: {measures_str}, count(1) as total_count",
        ]

        if filter_clause:
            query_parts.append(f"filter: {filter_clause}")

        return " | ".join(query_parts)
