"""BTQL Agent - Main orchestrator for query generation and refinement."""

import asyncio
from collections import Counter
import httpx
import json
import re
from pathlib import Path
from typing import Any

from .models import AgentConfig, AgentResult, JudgeResponse, QueryStatus, RefinementAttempt, SchemaInfo
from .runner import QueryRunner
from .builder import QueryBuilder
from .refiner import QueryRefiner
from .client_factory import create_llm_client


# Default cache directory (tools/auto-btql/.btql-agent)
_DEFAULT_CACHE_DIR = Path(__file__).parent.parent.parent / ".btql-agent"
SCHEMA_CACHE_VERSION = 4


class BTQLAgent:
    """Agent that generates and refines BTQL queries to fetch experiment metrics."""

    def __init__(self, config: AgentConfig, cache_dir: Path | None = None) -> None:
        self.config = config
        self.runner = QueryRunner(
            api_key=config.braintrust_api_key,
            timeout=config.timeout,
            braintrust_org=config.braintrust_org,
        )
        self.builder = QueryBuilder(default_filters=config.default_filters)

        # Create LLM client using factory based on model_client_type
        client, model = create_llm_client(config)
        self.refiner = QueryRefiner(client=client, model=model)

        self._cache_dir = cache_dir or _DEFAULT_CACHE_DIR
        self._cache_file = self._cache_dir / "schema_cache.json"
        self._schema_cache: dict[str, SchemaInfo] = self._load_schema_cache()

    def _load_schema_cache(self) -> dict[str, SchemaInfo]:
        """Load schema cache from disk."""
        if not self._cache_file.exists():
            return {}

        try:
            with open(self._cache_file, "r", encoding="utf-8") as f:
                data = json.load(f)

            cache: dict[str, SchemaInfo] = {}
            for exp_id, schema_dict in data.items():
                version = int(schema_dict.get("_schema_version", 1))
                metric_types = schema_dict.get("metric_types", {})
                if not isinstance(metric_types, dict):
                    metric_types = {}
                metric_groups = schema_dict.get("metric_groups", [])
                if not isinstance(metric_groups, list):
                    metric_groups = []
                # Force template regeneration for stale cache entries.
                stale_template = version < SCHEMA_CACHE_VERSION
                cache[exp_id] = SchemaInfo(
                    scores_path=schema_dict.get("scores_path", "scores"),
                    available_metrics=schema_dict.get("available_metrics", []),
                    metric_types=metric_types,
                    metric_groups=metric_groups,
                    explanation_paths=schema_dict.get("explanation_paths", []),
                    id_field=schema_dict.get("id_field", "root_span_id"),
                    raw_sample=None,  # Don't persist raw samples
                    scores_query_template=None if stale_template else schema_dict.get("scores_query_template"),
                    judge_query_template=None if stale_template else schema_dict.get("judge_query_template"),
                    trace_query_template=None if stale_template else schema_dict.get("trace_query_template"),
                    judge_query_templates=(
                        []
                        if stale_template
                        else (schema_dict.get("judge_query_templates", []) if isinstance(schema_dict.get("judge_query_templates", []), list) else [])
                    ),
                    # explanation_query_template is only used internally during discovery verification, not persisted
                )
            return cache
        except (json.JSONDecodeError, OSError):
            return {}

    def _save_schema_cache(self) -> None:
        """Save schema cache to disk."""
        try:
            self._cache_dir.mkdir(parents=True, exist_ok=True)

            data: dict[str, dict[str, Any]] = {}
            for exp_id, schema in self._schema_cache.items():
                data[exp_id] = {
                    "_schema_version": SCHEMA_CACHE_VERSION,
                    "scores_path": schema.scores_path,
                    "available_metrics": schema.available_metrics,
                    "metric_types": schema.metric_types,
                    "metric_groups": schema.metric_groups,
                    "explanation_paths": schema.explanation_paths,
                    "id_field": schema.id_field,
                    "scores_query_template": schema.scores_query_template,
                    "judge_query_template": schema.judge_query_template,
                    "trace_query_template": schema.trace_query_template,
                    "judge_query_templates": schema.judge_query_templates,
                    # explanation_query_template is only used internally during discovery verification, not persisted
                }

            with open(self._cache_file, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2)
        except OSError:
            pass  # Silently fail if we can't write cache

    def get_cached_schema(self, experiment_id: str) -> SchemaInfo | None:
        """Get cached schema for an experiment if available."""
        return self._schema_cache.get(experiment_id)

    async def _fetch_available_metrics(self, experiment_id: str, max_retries: int = 3) -> list[str]:
        """Fetch all available metric keys from the /summarize API.
        
        Args:
            experiment_id: The Braintrust experiment ID
            max_retries: Maximum number of retry attempts
            
        Returns:
            List of available metric keys from the experiment
        """
        # Determine the base URL based on braintrust_org
        base_urls = {
            "bt-azure": "https://prod-msft-api-endpoint-eackahcmezgzaybz.a03.azurefd.net",
            "non-pme": "https://api.braintrust.dev",
        }
        base_url = base_urls.get(self.config.braintrust_org, base_urls["non-pme"])
        url = f"{base_url}/v1/experiment/{experiment_id}/summarize"
        
        headers = {
            "Authorization": f"Bearer {self.config.braintrust_api_key}",
        }
        params = {"summarize_scores": "true"}
        
        for attempt in range(max_retries):
            try:
                async with httpx.AsyncClient(timeout=60.0) as client:
                    response = await client.get(url, headers=headers, params=params)
                    
                    if response.status_code == 200:
                        response_json = response.json()
                        scores = response_json.get('scores', {})
                        if scores and isinstance(scores, dict):
                            metric_keys = list(scores.keys())
                            print(f"Successfully fetched {len(metric_keys)} available metrics from /summarize API")
                            return metric_keys
                    elif response.status_code == 429:  # Rate limit
                        if attempt < max_retries - 1:
                            wait_time = 2 ** attempt  # Exponential backoff: 1s, 2s, 4s
                            print(f"Rate limited, retrying in {wait_time}s... (attempt {attempt + 1}/{max_retries})")
                            await asyncio.sleep(wait_time)
                            continue
                    else:
                        print(f"Warning: /summarize API returned status {response.status_code}")
                        if attempt < max_retries - 1:
                            await asyncio.sleep(1)
                            continue
                        
            except httpx.TimeoutException:
                if attempt < max_retries - 1:
                    print(f"Timeout fetching metrics, retrying... (attempt {attempt + 1}/{max_retries})")
                    await asyncio.sleep(1)
                    continue
                else:
                    print("Warning: Timeout fetching available metrics after all retries")
            except Exception as e:
                if attempt < max_retries - 1:
                    print(f"Error fetching metrics: {e}, retrying... (attempt {attempt + 1}/{max_retries})")
                    await asyncio.sleep(1)
                    continue
                else:
                    print(f"Warning: Failed to fetch available metrics after {max_retries} attempts: {e}")
        
        return []

    @staticmethod
    def _metadata_is_non_empty(metadata: Any) -> bool:
        """Return True when metadata contains non-empty content."""
        if metadata is None:
            return False
        if isinstance(metadata, str):
            text = metadata.strip()
            return text not in {"", "{}", "null", "None"}
        if isinstance(metadata, dict):
            return len(metadata) > 0
        if isinstance(metadata, list):
            return len(metadata) > 0
        return True

    @staticmethod
    def _metadata_has_explanation_like_content(metadata: Any) -> bool:
        """Heuristic check for explanation/rationale-like metadata content."""
        if not BTQLAgent._metadata_is_non_empty(metadata):
            return False

        keywords = (
            "rationale",
            "reason",
            "reasoning",
            "feedback",
            "explanation",
            "judge",
            "analysis",
            "improvement",
        )
        queue: list[Any] = [metadata]
        visited = 0
        max_visit = 200
        while queue and visited < max_visit:
            visited += 1
            item = queue.pop(0)
            if isinstance(item, dict):
                for k, v in item.items():
                    key = str(k).lower()
                    if any(tok in key for tok in keywords):
                        return True
                    queue.append(v)
            elif isinstance(item, list):
                queue.extend(item[:10])
            elif isinstance(item, str):
                text = item.lower()
                if any(tok in text for tok in keywords):
                    return True
        return False

    def _classify_metrics_from_sample_rows(
        self,
        rows: list[dict[str, Any]],
        available_metrics: list[str],
    ) -> tuple[dict[str, str], list[str]]:
        """Classify metrics into llm_judged/rule_based/hybrid from sampled scorer rows."""
        stats: dict[str, Counter] = {m: Counter() for m in available_metrics}
        available_set = set(available_metrics)

        for row in rows:
            scores = row.get("scores")
            if not isinstance(scores, dict):
                continue
            metadata = row.get("metadata")
            has_meta = self._metadata_is_non_empty(metadata)
            has_expl = self._metadata_has_explanation_like_content(metadata)
            for metric in scores.keys():
                metric_name = str(metric)
                if metric_name not in available_set:
                    continue
                stats[metric_name]["rows"] += 1
                if has_meta:
                    stats[metric_name]["meta_rows"] += 1
                if has_expl:
                    stats[metric_name]["expl_rows"] += 1

        metric_types: dict[str, str] = {}
        explanation_targets: list[str] = []
        for metric in available_metrics:
            row_cnt = stats[metric]["rows"]
            if row_cnt == 0:
                metric_types[metric] = "hybrid"
                explanation_targets.append(metric)
                continue
            meta_ratio = stats[metric]["meta_rows"] / row_cnt
            expl_ratio = stats[metric]["expl_rows"] / row_cnt

            if meta_ratio < 0.2 or expl_ratio < 0.1:
                metric_types[metric] = "rule_based"
            elif meta_ratio >= 0.6 and expl_ratio >= 0.3:
                metric_types[metric] = "llm_judged"
                explanation_targets.append(metric)
            else:
                metric_types[metric] = "hybrid"
                explanation_targets.append(metric)

        return metric_types, explanation_targets

    def _build_metric_groups_from_sample_rows(
        self,
        rows: list[dict[str, Any]],
        available_metrics: list[str],
    ) -> list[dict[str, Any]]:
        """Build scorer-like metric groups from sampled rows.

        Group key is the set of score keys present in a row. This captures cases where
        one scorer row emits multiple metrics together.
        """
        available_set = set(available_metrics)
        grouped: dict[str, dict[str, Any]] = {}
        ordered_keys: list[str] = []

        for row in rows:
            scores = row.get("scores")
            if not isinstance(scores, dict):
                continue

            row_metrics = sorted([str(k) for k in scores.keys() if str(k) in available_set])
            if not row_metrics:
                continue
            key = "|".join(row_metrics)

            if key not in grouped:
                grouped[key] = {
                    "id": f"group_{len(ordered_keys) + 1}",
                    "metrics": row_metrics,
                    "row_count": 0,
                    "metadata_non_empty_rows": 0,
                    "explanation_like_rows": 0,
                }
                ordered_keys.append(key)

            group = grouped[key]
            group["row_count"] += 1
            metadata = row.get("metadata")
            if self._metadata_is_non_empty(metadata):
                group["metadata_non_empty_rows"] += 1
            if self._metadata_has_explanation_like_content(metadata):
                group["explanation_like_rows"] += 1

        out: list[dict[str, Any]] = []
        for key in ordered_keys:
            group = grouped[key]
            if group["explanation_like_rows"] > 0:
                group_type = "llm_judged"
            elif group["metadata_non_empty_rows"] > 0:
                group_type = "metadata_non_empty"
            else:
                group_type = "rule_based"
            group["group_type"] = group_type
            out.append(group)

        return out

    async def discover_schema(
        self,
        cache_key: str,
        force_refresh: bool = False,
        experiment_id: str | None = None,
        max_attempts: int = 3,
    ) -> SchemaInfo:
        """Discover the schema of an experiment by sampling data and using LLM analysis.

        Discovers metrics and their explanations through iterative LLM analysis
        and verification, then generates query templates.

        All available metrics are automatically discovered and all explanations are extracted.

        Args:
            cache_key: The key used for caching (typically eval name or experiment ID)
            force_refresh: If True, ignore cache and re-discover schema
            experiment_id: The Braintrust experiment ID (uses cache_key if not provided)
            max_attempts: Maximum number of schema discovery attempts

        Returns:
            SchemaInfo with discovered paths to scores and explanations
        """
        # Check cache first (unless force refresh)
        if not force_refresh and cache_key in self._schema_cache:
            cached = self._schema_cache[cache_key]
            if not cached.metric_types and cached.available_metrics:
                cached.metric_types = {m: "hybrid" for m in cached.available_metrics}
            if not (cached.scores_query_template and cached.judge_query_template and cached.trace_query_template):
                self._generate_query_templates(cached)
                self._schema_cache[cache_key] = cached
                self._save_schema_cache()
            return cached

        # Use experiment_id if provided, otherwise use cache_key
        exp_id = experiment_id if experiment_id else cache_key

        # Step 1: Fetch all available metrics from /summarize API
        available_metric_keys = await self._fetch_available_metrics(exp_id)
        if not available_metric_keys:
            print("Warning: Could not fetch available metrics from /summarize API, falling back to basic sampling")
            print("Warning: No metrics available, returning default schema")
            return SchemaInfo()

        # Step 1.5: Metric typing from lightweight scorer sample.
        # This lets us skip expensive explanation discovery for fully rule-based experiments.
        metric_profile_query = (
            f"from: experiment('{exp_id}') | "
            "select: root_span_id, scores, metadata | "
            "filter: scores IS NOT NULL | "
            "sort: root_span_id asc | "
            "limit: 100"
        )
        metric_profile_result = await self.runner.execute(metric_profile_query)
        metric_types = {m: "hybrid" for m in available_metric_keys}
        metric_groups: list[dict[str, Any]] = []
        explanation_target_metrics = available_metric_keys[:]
        if metric_profile_result.is_success and metric_profile_result.data:
            metric_types, explanation_target_metrics = self._classify_metrics_from_sample_rows(
                rows=metric_profile_result.data,
                available_metrics=available_metric_keys,
            )
            metric_groups = self._build_metric_groups_from_sample_rows(
                rows=metric_profile_result.data,
                available_metrics=available_metric_keys,
            )

        def _build_and_cache_basic_schema() -> SchemaInfo:
            basic_schema = SchemaInfo(
                scores_path="scores",
                available_metrics=available_metric_keys,
                metric_types=metric_types,
                metric_groups=metric_groups,
                explanation_paths=[],
                id_field="root_span_id",
            )
            self._generate_query_templates(basic_schema)
            self._schema_cache[cache_key] = basic_schema
            self._save_schema_cache()
            return basic_schema

        rule_based_count = sum(1 for t in metric_types.values() if t == "rule_based")
        print(
            "Metric typing summary: total=%s rule_based=%s explanation_targets=%s"
            % (len(available_metric_keys), rule_based_count, len(explanation_target_metrics))
        )
        if not explanation_target_metrics:
            print("All discovered metrics look rule-based; skipping explanation path discovery.")
            return _build_and_cache_basic_schema()

        # Step 2: Find a root_span_id with complete scores
        # Query a small sample first, then expand only if metric coverage is low.
        # IMPORTANT: coverage for schema-discovery sampling should be based on
        # explanation-target metrics only (LLM-judged/hybrid), not all metrics.
        required_keys = set(explanation_target_metrics)
        initial_sample_limit = 100
        fallback_sample_limit = 300
        min_metric_coverage_ratio = 0.8

        def _build_root_span_query(limit: int) -> str:
            return (
                f"from: experiment('{exp_id}') | "
                "select: root_span_id, scores | "
                'filter: scores IS NOT NULL AND metadata != "" | '
                "sort: root_span_id asc | "
                f"limit: {limit}"
            )

        root_span_query = _build_root_span_query(initial_sample_limit)
        root_span_result = await self.runner.execute(root_span_query)
        if not root_span_result.is_success or not root_span_result.data:
            # Return default schema if sampling fails
            return _build_and_cache_basic_schema()

        sampled_metric_keys: set[str] = set()
        for row in root_span_result.data:
            scores = row.get("scores")
            if isinstance(scores, dict):
                sampled_metric_keys.update(scores.keys())

        metric_coverage = (len(sampled_metric_keys) / len(required_keys)) if required_keys else 1.0
        if metric_coverage < min_metric_coverage_ratio and len(root_span_result.data) >= initial_sample_limit:
            print(
                "Warning: Initial explanation-target metric coverage %.1f%% is low, expanding to %s rows"
                % (metric_coverage * 100, fallback_sample_limit)
            )
            fallback_query = _build_root_span_query(fallback_sample_limit)
            fallback_result = await self.runner.execute(fallback_query)
            if fallback_result.is_success and fallback_result.data:
                root_span_result = fallback_result

        # Step 3: Iterate through root_span_ids and find one with complete scores
        first_root_span_id = None
        
        # Group rows by root_span_id and collect all score keys
        root_span_scores: dict[str, set[str]] = {}
        for row in root_span_result.data:
            root_span_id = row.get('root_span_id')
            scores = row.get('scores')
            
            if root_span_id and scores and isinstance(scores, dict):
                if root_span_id not in root_span_scores:
                    root_span_scores[root_span_id] = set()
                root_span_scores[root_span_id].update(scores.keys())
        
        # Find the first root_span_id that has all available metric keys
        if available_metric_keys:
            for root_span_id, score_keys in root_span_scores.items():
                if required_keys.issubset(score_keys):
                    first_root_span_id = root_span_id
                    print(f"Found root_span_id with complete scores: {first_root_span_id}")
                    break
        
        # If no complete root_span_id found, use the one with the most scores
        if not first_root_span_id and root_span_scores:
            first_root_span_id = max(root_span_scores.items(), key=lambda x: len(x[1]))[0]
            max_score_count = len(root_span_scores[first_root_span_id])
            print(
                "Warning: No root_span_id with complete explanation-target metrics found, "
                f"using one with most scores: {first_root_span_id} ({max_score_count}/{len(required_keys)} metrics)"
            )
        
        if not first_root_span_id:
            return _build_and_cache_basic_schema()

        # Now get all rows with this root_span_id
        sample_query = (
            f"from: experiment('{exp_id}') | "
            "select: root_span_id, metadata, scores | "
            f"filter: root_span_id = '{first_root_span_id}' AND scores IS NOT NULL | "
            "limit: 100"
        )
        sample_result = await self.runner.execute(sample_query)

        if not sample_result.is_success or not sample_result.data:
            # Return default schema if sampling fails
            return _build_and_cache_basic_schema()

        sample_data = sample_result.data

        # ===== TASK 1: Discover and verify explanation paths =====
        print("\n=== Task 1: Discovering explanation paths ===")
        target_metric_set = set(explanation_target_metrics)
        group_analysis_specs: list[dict[str, Any]] = []
        for group in metric_groups:
            raw_group_metrics = group.get("metrics", [])
            if not isinstance(raw_group_metrics, list):
                continue
            group_metrics = [str(m) for m in raw_group_metrics if str(m) in target_metric_set]
            if not group_metrics:
                continue
            if str(group.get("group_type", "")).strip() == "rule_based":
                continue
            group_analysis_specs.append(
                {
                    "group_id": str(group.get("id", "")).strip() or f"group_{len(group_analysis_specs) + 1}",
                    "metrics": group_metrics,
                }
            )

        # Fallback to a single group if scorer grouping is unavailable.
        if not group_analysis_specs:
            group_analysis_specs.append({"group_id": "group_all", "metrics": explanation_target_metrics[:]})

        merged_explanation_paths: list[dict[str, Any]] = []
        merged_explanation_keys: set[tuple[str, str, str, str]] = set()
        final_scores_path = "scores"
        final_id_field = "root_span_id"
        any_group_succeeded = False

        for group_idx, group_spec in enumerate(group_analysis_specs, start=1):
            group_id = group_spec["group_id"]
            group_metrics = group_spec["metrics"]
            print(
                f"Explanation group {group_idx}/{len(group_analysis_specs)}: "
                f"{group_id} metrics={group_metrics}"
            )

            explanation_attempt_history: list[dict[str, str]] = []
            group_explanation_schema: dict[str, Any] | None = None

            for attempt in range(max_attempts):
                print(f"  Attempt {attempt + 1}/{max_attempts}")

                group_explanation_schema = await self.refiner.analyze_explanations(
                    sample_data=sample_data,
                    metrics=group_metrics,
                    attempt_history=explanation_attempt_history,
                )

                if not group_explanation_schema or not group_explanation_schema.get("explanation_paths"):
                    error_msg = "Failed to find explanation paths"
                    explanation_attempt_history.append({
                        "query": "N/A",
                        "error": error_msg,
                    })
                    print(f"    {error_msg}")
                    continue

                temp_schema = SchemaInfo(
                    scores_path=group_explanation_schema.get("scores_path", "scores"),
                    available_metrics=available_metric_keys,
                    metric_types=metric_types,
                    metric_groups=metric_groups,
                    explanation_paths=group_explanation_schema.get("explanation_paths", []),
                    id_field=group_explanation_schema.get("id_field", "root_span_id"),
                )

                self._generate_explanation_query(temp_schema)
                if not temp_schema.explanation_query_template:
                    error_msg = "Failed to generate explanation query template"
                    explanation_attempt_history.append({
                        "query": "N/A",
                        "error": error_msg,
                    })
                    print(f"    {error_msg}")
                    continue

                test_query = temp_schema.explanation_query_template.replace("{experiment_id}", exp_id)
                offset = attempt * 3
                if offset > 0:
                    test_query += f" | skip: {offset}"
                test_query += " | limit: 3"

                query_result = await self.runner.execute(test_query)
                if not query_result.is_success or not query_result.data:
                    error_msg = query_result.error or "Query returned no data"
                    explanation_attempt_history.append({
                        "query": temp_schema.explanation_query_template,
                        "error": f"Query execution failed: {error_msg}",
                    })
                    print(f"    Query execution failed: {error_msg}")
                    continue

                first_row = query_result.data[0]
                verification = await self.refiner.verify_explanation_query(
                    sample_result=first_row,
                    available_metrics=group_metrics,
                )

                if verification.get("is_valid", False):
                    print("    Explanation verification passed!")
                    break

                explanation = verification.get("explanation", "Unknown verification failure")
                explanation_attempt_history.append({
                    "query": temp_schema.explanation_query_template,
                    "error": f"Verification failed: {explanation}",
                })
                print(f"    Verification failed: {explanation}")
                if attempt == max_attempts - 1:
                    print("    Max attempts reached for this group")

            if not group_explanation_schema:
                print(f"  Skip group {group_id}: no usable explanation schema")
                continue

            any_group_succeeded = True
            final_scores_path = group_explanation_schema.get("scores_path", final_scores_path)
            final_id_field = group_explanation_schema.get("id_field", final_id_field)
            for exp_info in group_explanation_schema.get("explanation_paths", []):
                if not isinstance(exp_info, dict):
                    continue
                metric_name = str(exp_info.get("metric", "")).strip()
                if metric_name and metric_name not in group_metrics:
                    continue
                key = (
                    metric_name,
                    str(exp_info.get("path", "")),
                    str(exp_info.get("type", "")),
                    str(exp_info.get("json_key", "")),
                )
                if key in merged_explanation_keys:
                    continue
                merged_explanation_keys.add(key)
                merged_explanation_paths.append(exp_info)

        if not any_group_succeeded:
            print("Failed to discover explanation paths, using basic schema without judge paths")
            return _build_and_cache_basic_schema()

        # ===== Generate final query template =====
        print("\n=== Generating final query template ===")
        final_schema = SchemaInfo(
            scores_path=final_scores_path,
            available_metrics=available_metric_keys,
            metric_types=metric_types,
            metric_groups=metric_groups,
            explanation_paths=merged_explanation_paths,
            id_field=final_id_field,
            raw_sample=sample_data[0] if sample_data else None,
        )

        # Generate query templates for explanations
        self._generate_query_templates(final_schema)

        # Cache the result (both in-memory and on disk) using cache_key
        self._schema_cache[cache_key] = final_schema
        self._save_schema_cache()
        
        print("Schema discovery completed successfully!")
        return final_schema

    async def fetch_metrics(
        self,
        experiment_id: str,
        metrics: list[str],
        custom_filters: list[str] | None = None,
        use_llm_for_initial: bool = False,
        force_refresh_schema: bool = False,
        eval_name: str | None = None,
        max_schema_attempts: int = 3,
    ) -> AgentResult:
        """Fetch metrics from an experiment using BTQL.

        Always uses schema discovery (from cache or LLM) to find explanations.

        Args:
            experiment_id: The Braintrust experiment ID
            metrics: List of metric names to fetch
            custom_filters: Additional filter conditions (beyond defaults)
            use_llm_for_initial: Whether to use LLM to generate the initial query
            force_refresh_schema: Whether to force refresh the schema cache
            eval_name: Evaluation name for schema caching (defaults to experiment_id)
            max_schema_attempts: Maximum number of schema discovery attempts

        Returns:
            AgentResult with the fetched metrics or error information
        """
        attempts: list[RefinementAttempt] = []
        all_filters = list(self.config.default_filters)
        if custom_filters:
            all_filters.extend(custom_filters)

        # Use eval_name as cache key if provided, otherwise use experiment_id
        cache_key = eval_name if eval_name else experiment_id

        # Always discover schema (uses cache if available)
        schema = await self.discover_schema(cache_key, force_refresh=force_refresh_schema, experiment_id=experiment_id, max_attempts=max_schema_attempts)

        # Generate initial query
        if use_llm_for_initial:
            current_query = await self.refiner.generate_initial_query(
                experiment_id=experiment_id,
                metrics=metrics,
                filters=all_filters,
            )
        else:
            current_query = self.builder.build_query(
                experiment_id=experiment_id,
                metrics=metrics,
                custom_filters=custom_filters,
            )

        # Execute and refine loop
        for iteration in range(self.config.max_iterations):
            result = await self.runner.execute(current_query)

            attempt = RefinementAttempt(
                iteration=iteration + 1,
                query=current_query,
                result=result,
            )
            attempts.append(attempt)

            if result.is_success:
                # Parse and aggregate the metrics
                aggregated = self._aggregate_scores(result.data, metrics, schema)

                # Check if we actually found any of the requested metrics
                if aggregated:
                    # Extract judge responses using schema
                    judge_responses = self._extract_judge_responses(result.data, metrics, schema)
                    return AgentResult(
                        success=True,
                        metrics=aggregated,
                        judge_responses=judge_responses,
                        final_query=current_query,
                        iterations=iteration + 1,
                        history=self._format_history(attempts),
                    )
                else:
                    # Query succeeded but no matching metrics found - treat as needing refinement
                    result.status = QueryStatus.ERROR
                    result.error = f"No data found for metrics: {', '.join(metrics)}. The query returned {len(result.data or [])} rows but none contained the requested score fields."

            # Query failed or no metrics found - try to refine
            if iteration < self.config.max_iterations - 1:
                attempt_history = [
                    {"query": a.query, "error": a.result.error}
                    for a in attempts
                ]
                current_query = await self.refiner.refine_query(
                    experiment_id=experiment_id,
                    metrics=metrics,
                    previous_query=current_query,
                    error_message=result.error or "Unknown error",
                    attempt_history=attempt_history,
                    filters=all_filters,
                )

        # Max iterations reached
        return AgentResult(
            success=False,
            error=f"Failed after {self.config.max_iterations} iterations. Last error: {attempts[-1].result.error}",
            final_query=current_query,
            iterations=self.config.max_iterations,
            history=self._format_history(attempts),
        )

    def _aggregate_scores(
        self,
        data: list[dict[str, Any]] | None,
        metrics: list[str],
        schema: SchemaInfo,
    ) -> dict[str, float]:
        """Aggregate scores from query results using schema.

        Args:
            data: Raw query result data
            metrics: List of metrics to aggregate
            schema: Discovered schema with path to scores

        Returns:
            Dictionary mapping metric names to their average values
        """
        if not data:
            return {}

        # Collect all score values per metric
        metric_values: dict[str, list[float]] = {m: [] for m in metrics}

        for row in data:
            scores = self._extract_by_path(row, schema.scores_path)
            if not scores or not isinstance(scores, dict):
                continue

            for metric in metrics:
                # Try exact match first
                value = scores.get(metric)
                # Try case-insensitive match
                if value is None:
                    for key, val in scores.items():
                        if key.lower() == metric.lower():
                            value = val
                            break

                if value is not None and isinstance(value, (int, float)):
                    metric_values[metric].append(float(value))

        # Compute averages
        aggregated = {}
        for metric, values in metric_values.items():
            if values:
                aggregated[metric] = sum(values) / len(values)

        return aggregated

    def _extract_judge_responses(
        self,
        data: list[dict[str, Any]] | None,
        metrics: list[str],
        schema: SchemaInfo,
    ) -> list[JudgeResponse]:
        """Extract judge responses using discovered schema paths.

        Args:
            data: Raw query result data
            metrics: List of metrics to extract explanations for
            schema: Discovered schema with paths to explanations

        Returns:
            List of JudgeResponse objects
        """
        if not data:
            return []

        judge_responses: list[JudgeResponse] = []

        for row in data:
            root_span_id = self._extract_by_path(row, schema.id_field)
            scores = self._extract_by_path(row, schema.scores_path)

            if not scores or not isinstance(scores, dict):
                continue

            # Build explanations dict using schema paths
            explanations: dict[str, str] = {}
            for exp_info in schema.explanation_paths:
                if not isinstance(exp_info, dict):
                    continue

                metric_name = exp_info.get("metric", "").lower()
                path = exp_info.get("path", "")
                exp_type = exp_info.get("type", "direct")
                json_key = exp_info.get("json_key", "")

                if not path:
                    continue

                value = self._extract_by_path(row, path)
                if not value:
                    continue

                if exp_type == "direct" and isinstance(value, str):
                    explanations[metric_name] = value
                elif exp_type == "json" and isinstance(value, str):
                    parsed = self._parse_json(value)
                    if parsed and json_key:
                        # Navigate json_key path (e.g., "metric.feedback")
                        exp_value = self._extract_by_path(parsed, json_key)
                        if exp_value and isinstance(exp_value, str):
                            explanations[metric_name] = exp_value

            # Extract explanations for requested metrics
            for metric in metrics:
                # Get score value (case-insensitive)
                score_value = None
                actual_metric_key = metric
                for key, val in scores.items():
                    if key.lower() == metric.lower():
                        score_value = val
                        actual_metric_key = key
                        break

                if score_value is None:
                    continue

                # Get explanation for this metric
                explanation = explanations.get(metric.lower()) or explanations.get(actual_metric_key.lower(), "")

                if explanation:
                    judge_responses.append(JudgeResponse(
                        metric=metric,
                        score=float(score_value) if isinstance(score_value, (int, float)) else 0.0,
                        explanation=explanation,
                        root_span_id=root_span_id,
                    ))

        return judge_responses

    def _extract_by_path(self, data: dict[str, Any] | list[Any] | None, path: str) -> Any:
        """Extract value using dot-notation path.

        Supports:
        - Dot notation: "scores.completeness"
        - Array indices: "input.output.0.judge_response"
        - Negative indices: "array.-1" (last element)

        Args:
            data: The data structure to extract from
            path: Dot-notation path (e.g., "input.output.0.text")

        Returns:
            The extracted value or None if not found
        """
        if data is None:
            return None

        if not path:
            return data

        parts = path.split(".")
        current: Any = data

        for part in parts:
            if current is None:
                return None

            if isinstance(current, dict):
                current = current.get(part)
            elif isinstance(current, list):
                try:
                    idx = int(part)
                    if -len(current) <= idx < len(current):
                        current = current[idx]
                    else:
                        return None
                except ValueError:
                    return None
            else:
                return None

        return current

    def _parse_json(self, content: str) -> dict[str, Any] | None:
        """Parse JSON from content string."""
        if not content:
            return None

        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass

        # Try to extract JSON from content
        json_match = re.search(r'\{[\s\S]*\}', content)
        if json_match:
            try:
                return json.loads(json_match.group(0))
            except json.JSONDecodeError:
                pass

        return None

    def _format_history(
        self,
        attempts: list[RefinementAttempt],
    ) -> list[dict[str, Any]]:
        """Format attempt history for output."""
        return [
            {
                "iteration": a.iteration,
                "query": a.query,
                "status": a.result.status.value,
                "error": a.result.error,
                "data_count": len(a.result.data) if a.result.data else 0,
            }
            for a in attempts
        ]

    def _generate_explanation_query(self, schema: SchemaInfo) -> None:
        """Generate query template for testing explanation extraction only.
        
        Templates use {experiment_id} as placeholder.
        This method generates a simplified query that only fetches explanations,
        used during the explanation discovery phase.
        """
        # Import helper functions from __main__ module
        from . import __main__ as main_module
        
        # Build dimension fields for explanation query
        dimension_fields = ["root_span_id", "input.input AS input"]
        direct_explanations = []
        json_explanations = []
        json_paths = set()
        
        if schema.explanation_paths:
            for exp_info in schema.explanation_paths:
                if isinstance(exp_info, dict):
                    path = exp_info.get("path", "")
                    exp_type = exp_info.get("type", "direct")
                    metric = exp_info.get("metric", "")
                    
                    if not path or exp_type == "not_found":
                        continue
                    btql_path = main_module._convert_path_to_btql(path)
                    if exp_type == "direct":
                        direct_explanations.append({"metric": metric, "path": btql_path})
                    elif exp_type == "json" and path not in json_paths:
                        json_paths.add(path)
                        alias = main_module._create_alias_from_path(path)
                        json_explanations.append({"metric": alias, "path": btql_path})
            
            # Combine direct and json explanations into judge_response
            all_explanation_parts = []
            
            for exp in direct_explanations:
                safe_metric = exp["metric"].replace('"', '\\"')  # Escape quotes in metric name
                all_explanation_parts.append(
                    f'"\\"", "{safe_metric}", "\\": \\"", coalesce(to_string({exp["path"]}), \'\'), "\\""'
                )
            
            for exp in json_explanations:
                safe_metric = exp["metric"].replace('"', '\\"')  # Escape quotes in metric name
                all_explanation_parts.append(
                    f'"\\"", "{safe_metric}", "\\": \\"", coalesce(to_string({exp["path"]}), \'\'), "\\""'
                )
            
            if all_explanation_parts:
                concat_expr = "concat('{', " + ", ', ', ".join(all_explanation_parts) + ", '}'" + ")"
                dimension_fields.append(f"{concat_expr} AS judge_response")
            else:
                dimension_fields.append('"{}" AS judge_response')
        else:
            dimension_fields.append('"{}" AS judge_response')
        dimension_fields.append('scores AS scores')
        
        dimensions_str = ", ".join(dimension_fields)
        
        # Build filter conditions
        filter_conditions = ["scores IS NOT NULL"]
        
        if direct_explanations or json_explanations:
            explanation_checks = []
            for exp in direct_explanations:
                explanation_checks.append(f"{exp['path']} IS NOT NULL")
            for exp in json_explanations:
                explanation_checks.append(f"{exp['path']} IS NOT NULL")
            if explanation_checks:
                filter_conditions.append("(" + " OR ".join(explanation_checks) + ")")
        
        filter_str = " AND ".join(filter_conditions)
        
        # Generate explanation query template
        schema.explanation_query_template = (
            f"from: experiment('{{experiment_id}}') | "
            f"select: root_span_id | "
            f"dimensions: {dimensions_str} | "
            f"filter: {filter_str}"
        )

    def _generate_query_templates(self, schema: SchemaInfo) -> None:
        """Generate and store query templates in the schema.
        
        Templates use {experiment_id} and {root_span_id} as placeholders.
        """
        # Import helper functions from __main__ module
        from . import __main__ as main_module
        
        # Generate scores query template
        schema.scores_query_template = (
            "from: experiment('{experiment_id}') | "
            "select: root_span_id, scores | "
            "dimensions: root_span_id, scores | "
            'filter: scores IS NOT NULL'
        )
        
        # Build dimension fields for judge query
        dimension_fields = ["root_span_id", "input.input AS input"]
        direct_explanations = []
        json_explanations = []
        json_paths = set()
        
        if schema.explanation_paths:
            for exp_info in schema.explanation_paths:
                if isinstance(exp_info, dict):
                    path = exp_info.get("path", "")
                    exp_type = exp_info.get("type", "direct")
                    metric = exp_info.get("metric", "")
                    
                    if not path or exp_type == "not_found":
                        continue
                    btql_path = main_module._convert_path_to_btql(path)
                    if exp_type == "direct":
                        direct_explanations.append({"metric": metric, "path": btql_path})
                    elif exp_type == "json" and path not in json_paths:
                        json_paths.add(path)
                        alias = main_module._create_alias_from_path(path)
                        json_explanations.append({"metric": alias, "path": btql_path})
            
            # Combine direct and json explanations into judge_response
            all_explanation_parts = []
            
            for exp in direct_explanations:
                safe_metric = exp["metric"].replace('"', '\\"')  # Escape quotes in metric name
                all_explanation_parts.append(
                    f'"\\"", "{safe_metric}", "\\": \\"", coalesce(to_string({exp["path"]}), \'\'), "\\""'
                )
            
            for exp in json_explanations:
                safe_metric = exp["metric"].replace('"', '\\"')  # Escape quotes in metric name
                all_explanation_parts.append(
                    f'"\\"", "{safe_metric}", "\\": \\"", coalesce(to_string({exp["path"]}), \'\'), "\\""'
                )
            
            if all_explanation_parts:
                concat_expr = "concat('{', " + ", ', ', ".join(all_explanation_parts) + ", '}'" + ")"
                dimension_fields.append(f"{concat_expr} AS judge_response")
            else:
                dimension_fields.append('"{}" AS judge_response')
        else:
            dimension_fields.append('"{}" AS judge_response')
        
        dimensions_str = ", ".join(dimension_fields)
        
        filter_str = "scores IS NOT NULL"
        
        # Generate judge query template
        schema.judge_query_template = (
            f"from: experiment('{{experiment_id}}') | "
            f"select: root_span_id | "
            f"dimensions: {dimensions_str} | "
            f"filter: {filter_str}"
        )
        
        # Keep trace query aligned with scores coverage. Judge response may be empty for rule-based metrics.
        trace_filter_str = "root_span_id = '{{root_span_id}}' AND scores IS NOT NULL"
        
        # Generate trace query template
        schema.trace_query_template = (
            f"from: experiment('{{experiment_id}}') | "
            f"select: root_span_id | "
            f"dimensions: {dimensions_str} | "
            f"filter: {trace_filter_str}"
        )

        # Generate grouped judge/trace templates (one per scorer-like metric group).
        grouped_templates: list[dict[str, Any]] = []
        metric_to_explanations: dict[str, list[dict[str, Any]]] = {}
        for exp_info in (schema.explanation_paths or []):
            if not isinstance(exp_info, dict):
                continue
            metric_name = str(exp_info.get("metric", "")).strip()
            path = exp_info.get("path")
            exp_type = str(exp_info.get("type", "direct")).strip()
            if not metric_name or not isinstance(path, str) or not path.strip() or exp_type == "not_found":
                continue
            metric_to_explanations.setdefault(metric_name, []).append(exp_info)

        for idx, group in enumerate(schema.metric_groups or [], start=1):
            raw_metrics = group.get("metrics", [])
            if not isinstance(raw_metrics, list):
                continue
            group_metrics = [str(m) for m in raw_metrics if str(m)]
            if not group_metrics:
                continue

            group_items: list[dict[str, Any]] = []
            seen_items: set[tuple[str, str, str, str]] = set()
            for metric_name in group_metrics:
                for exp_info in metric_to_explanations.get(metric_name, []):
                    key = (
                        str(exp_info.get("metric", "")),
                        str(exp_info.get("path", "")),
                        str(exp_info.get("type", "")),
                        str(exp_info.get("json_key", "")),
                    )
                    if key in seen_items:
                        continue
                    seen_items.add(key)
                    group_items.append(exp_info)

            judge_q = main_module._build_judge_query(
                "{experiment_id}",
                schema,
                group_items,
                metrics_for_filter=group_metrics,
            )
            trace_q = main_module._build_trace_query(
                "{experiment_id}",
                schema,
                group_items,
                metrics_for_filter=group_metrics,
            )

            grouped_templates.append(
                {
                    "group_id": str(group.get("id", "")).strip() or f"group_{idx}",
                    "group_type": str(group.get("group_type", "")).strip() or ("rule_based" if not group_items else "llm_judged"),
                    "metrics": group_metrics,
                    "judge_query_template": judge_q,
                    "trace_query_template": trace_q,
                }
            )

        schema.judge_query_templates = grouped_templates
