"""CLI entry point for BTQL Agent."""

import asyncio
import json
import os
import sys

import click
from rich.console import Console
from rich.table import Table


console = Console()


def get_config_from_env(
    max_iterations: int,
    braintrust_api_key: str | None = None,
    braintrust_org: str | None = None,
    model_type: str | None = None,
    model_endpoint: str | None = None,
    model_name: str | None = None,
):
    """Load configuration from CLI args, then env vars as fallback.

    Args:
        max_iterations: Maximum query refinement iterations
        braintrust_api_key: Braintrust API key (CLI override)
        braintrust_org: Braintrust organization (CLI override)
        model_type: Model client type: "azure", "copilot", or "local" (CLI override)
        model_endpoint: Model endpoint URL (CLI override)
        model_name: Model name/deployment (CLI override)
    """
    from dotenv import load_dotenv
    from .models import AgentConfig

    load_dotenv()

    # Braintrust configuration
    braintrust_org_value = braintrust_org or os.getenv("BRAINTRUST_ORG", "non-pme")
    if braintrust_api_key:
        braintrust_key = braintrust_api_key
    elif braintrust_org_value == "bt-azure":
        braintrust_key = os.getenv("BRAINTRUST_AZURE_API_KEY")
    else:
        braintrust_key = os.getenv("BRAINTRUST_API_KEY")

    if not braintrust_key:
        console.print("[red]Missing BRAINTRUST_API_KEY environment variable[/red]")
        sys.exit(1)

    # Model client type: CLI > env > default
    client_type = model_type or os.getenv("MODEL_CLIENT_TYPE", "azure")

    if client_type == "azure":
        return _build_azure_config(
            braintrust_key, braintrust_org_value, max_iterations, model_name
        )
    elif client_type == "papyrus":
        return _build_papyrus_config(
            braintrust_key, braintrust_org_value, max_iterations, model_endpoint, model_name
        )
    elif client_type == "copilot":
        return _build_copilot_config(
            braintrust_key, braintrust_org_value, max_iterations, model_endpoint, model_name
        )
    elif client_type == "local":
        return _build_local_config(
            braintrust_key, braintrust_org_value, max_iterations, model_endpoint, model_name
        )
    else:
        console.print(f"[red]Unknown model type: {client_type}[/red]")
        console.print("Valid options: azure, papyrus, copilot, local")
        sys.exit(1)


def _build_azure_config(
    braintrust_key: str,
    braintrust_org: str,
    max_iterations: int,
    model_name: str | None,
):
    """Build config for Azure OpenAI."""
    from azure.identity import DefaultAzureCredential, get_bearer_token_provider
    from .models import AgentConfig

    azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "https://cnmoderation-evaluation.openai.azure.com/")
    azure_deployment = model_name or os.getenv("AZURE_OPENAI_DEPLOYMENT", "gpt-5.1-chat-2")

    missing = []
    if not azure_endpoint:
        missing.append("AZURE_OPENAI_ENDPOINT")
    if not azure_deployment:
        missing.append("AZURE_OPENAI_DEPLOYMENT or --model-name")

    if missing:
        console.print(f"[red]Missing required for Azure: {', '.join(missing)}[/red]")
        console.print("\nAuthentication uses DefaultAzureCredential (Azure AD).")
        console.print("Ensure you are logged in via 'az login' or have appropriate credentials.")
        sys.exit(1)

    token_provider = get_bearer_token_provider(
        DefaultAzureCredential(),
        "https://cognitiveservices.azure.com/.default",
    )

    return AgentConfig(
        braintrust_api_key=braintrust_key,
        braintrust_org=braintrust_org,
        max_iterations=max_iterations,
        model_client_type="azure",
        azure_openai_endpoint=azure_endpoint,
        azure_openai_deployment=azure_deployment,
        token_provider=token_provider,
    )


def _build_papyrus_config(
    braintrust_key: str,
    braintrust_org: str,
    max_iterations: int,
    model_endpoint: str | None,
    model_name: str | None,
):
    """Build config for Papyrus."""
    from azure.identity import DefaultAzureCredential, get_bearer_token_provider
    from .models import AgentConfig

    papyrus_scope = os.getenv("PAPYRUS_VERIFY_SCOPES", "api://5fe538a8-15d5-4a84-961e-be66cd036687/.default")
    endpoint = model_endpoint or os.getenv("PAPYRUS_ENDPOINT", "https://hk.papyrus.binginternal.com/chat/completions")
    deployment = model_name or os.getenv("PAPYRUS_MODEL", "gpt-51-chat-2025-11-13-eval")

    missing = []
    if not endpoint:
        missing.append("PAPYRUS_ENDPOINT or --model-endpoint")
    if not deployment:
        missing.append("PAPYRUS_MODEL or --model-name")

    if missing:
        console.print(f"[red]Missing required for Papyrus: {', '.join(missing)}[/red]")
        console.print("\nAuthentication uses DefaultAzureCredential (Azure AD).")
        console.print("Ensure you are logged in via 'az login' or have appropriate credentials.")
        sys.exit(1)

    token_provider = get_bearer_token_provider(
        DefaultAzureCredential(),
        papyrus_scope,
    )

    return AgentConfig(
        braintrust_api_key=braintrust_key,
        braintrust_org=braintrust_org,
        max_iterations=max_iterations,
        model_client_type="papyrus",
        papyrus_endpoint=endpoint,
        papyrus_deployment=deployment,
        token_provider=token_provider,
    )


def _build_copilot_config(
    braintrust_key: str,
    braintrust_org: str,
    max_iterations: int,
    model_endpoint: str | None,
    model_name: str | None,
):
    """Build config for Copilot via VS Code extension proxy."""
    from .models import AgentConfig

    endpoint = model_endpoint or os.getenv("COPILOT_ENDPOINT", "http://localhost:8237")
    model = model_name or os.getenv("COPILOT_MODEL")

    if not model:
        console.print("[red]Copilot requires --model-name or COPILOT_MODEL environment variable[/red]")
        console.print("\nExample: btql-agent -e abc123 --eval picasso --model-type copilot --model-name gpt-4o")
        sys.exit(1)

    return AgentConfig(
        braintrust_api_key=braintrust_key,
        braintrust_org=braintrust_org,
        max_iterations=max_iterations,
        model_client_type="copilot",
        copilot_endpoint=endpoint,
        copilot_model=model,
    )


def _build_local_config(
    braintrust_key: str,
    braintrust_org: str,
    max_iterations: int,
    model_endpoint: str | None,
    model_name: str | None,
):
    """Build config for local OpenAI-compatible server."""
    from .models import AgentConfig

    endpoint = model_endpoint or os.getenv("LOCAL_ENDPOINT", "http://localhost:23333/api/openai/v1")
    model = model_name or os.getenv("LOCAL_MODEL")

    if not model:
        console.print("[red]Local requires --model-name or LOCAL_MODEL environment variable[/red]")
        console.print("\nExample: btql-agent -e abc123 --eval picasso --model-type local --model-name qwen-72b")
        sys.exit(1)

    return AgentConfig(
        braintrust_api_key=braintrust_key,
        braintrust_org=braintrust_org,
        max_iterations=max_iterations,
        model_client_type="local",
        local_endpoint=endpoint,
        local_model=model,
    )


def display_table_output(experiment_id: str, result: dict) -> None:
    """Display results as a formatted table."""
    table = Table(title=f"Experiment: {experiment_id}")
    table.add_column("Metric", style="cyan")
    table.add_column("Value", style="green", justify="right")

    metrics = result.get("metrics", {})
    for metric, value in metrics.items():
        table.add_row(metric, f"{value:.4f}")

    console.print(table)
    console.print(f"\n[dim]Iterations: {result.get('iterations', 0)}[/dim]")


def display_verbose_output(result: dict) -> None:
    """Display verbose output including attempt history."""
    console.print("\n[bold]Query History:[/bold]")
    for attempt in result.get("history", []):
        status_color = "green" if attempt["status"] == "success" else "red"
        console.print(f"\n[bold]Iteration {attempt['iteration']}:[/bold]")
        console.print(f"  Status: [{status_color}]{attempt['status']}[/{status_color}]")
        console.print(f"  Query: [dim]{attempt['query']}[/dim]")
        if attempt.get("error"):
            console.print(f"  Error: [red]{attempt['error']}[/red]")
        if attempt.get("data_count"):
            console.print(f"  Records: {attempt['data_count']}")

    if result.get("final_query"):
        console.print(f"\n[bold]Final Query:[/bold]\n[dim]{result['final_query']}[/dim]")


def display_judge_responses(judge_responses: list, limit: int = 5) -> None:
    """Display judge responses (explanations for scores)."""
    if not judge_responses:
        console.print("\n[dim]No judge responses found in data.[/dim]")
        return

    console.print(f"\n[bold]Judge Responses (showing {min(limit, len(judge_responses))} of {len(judge_responses)}):[/bold]")

    # Group by root_span_id for cleaner display
    displayed = 0
    seen_spans = set()

    for jr in judge_responses:
        if displayed >= limit:
            break
        if jr.root_span_id in seen_spans:
            continue
        seen_spans.add(jr.root_span_id)

        console.print(f"\n[cyan]Span:[/cyan] {jr.root_span_id}")

        # Find all responses for this span
        span_responses = [r for r in judge_responses if r.root_span_id == jr.root_span_id]
        for resp in span_responses:
            console.print(f"  [bold]{resp.metric}[/bold]: {resp.score:.2f}")
            # Wrap long explanations and sanitize for Windows console
            explanation = resp.explanation
            if len(explanation) > 200:
                explanation = explanation[:200] + "..."
            # Sanitize for Windows console compatibility using UTF-8 encoding
            explanation = explanation.encode("utf-8", errors="ignore").decode("utf-8")
            console.print(f"    [dim]{explanation}[/dim]")

        displayed += 1


def display_final_query_template(experiment_id: str, schema = None, format: str = "console", eval_name: str | None = None, metrics: list[str] | None = None) -> None:
    """Display the final query templates for fetching scores and judge responses.

    Args:
        experiment_id: The experiment ID
        schema: Optional schema information (SchemaInfo)
        format: Output format - "console" for formatted console output or "json" for JSON
        eval_name: Optional evaluation name
        metrics: Optional list of specific metrics to include (if None, includes all available metrics)
    """
    # Build the template data structure
    template_data = {
        "experiment_id": experiment_id,
        "eval": eval_name,
        "queries": {}
    }

    # Add metrics and schema info if available
    if schema:
        if schema.available_metrics:
            template_data["metrics"] = schema.available_metrics
        if schema.metric_types:
            template_data["metric_types"] = schema.metric_types
        template_data["schema_cached"] = bool(schema.scores_query_template and schema.judge_query_template and schema.trace_query_template)

    if format == "console":
        console.print("\n[bold]Final Query Templates:[/bold]")

    # Filter explanation paths by metrics if specified
    filtered_explanation_paths = schema.explanation_paths if schema else []
    if metrics and schema and schema.explanation_paths:
        # Convert metrics to lowercase for case-insensitive comparison
        metrics_lower = [m.lower() for m in metrics]
        filtered_explanation_paths = [
            exp for exp in schema.explanation_paths
            if isinstance(exp, dict) and exp.get("metric", "").lower() in metrics_lower
        ]

    def _build_grouped_dimension_queries() -> list[dict]:
        if not schema:
            return []

        def _bind_experiment_id(q: str | None) -> str | None:
            if not isinstance(q, str):
                return q
            return q.replace("{{experiment_id}}", experiment_id).replace("{experiment_id}", experiment_id)

        # Prefer prebuilt grouped templates from schema discovery.
        prebuilt = getattr(schema, "judge_query_templates", None)
        if isinstance(prebuilt, list) and prebuilt:
            selected_metrics = set(metrics[:] if metrics else (schema.available_metrics[:] if schema.available_metrics else []))
            out: list[dict] = []
            for item in prebuilt:
                if not isinstance(item, dict):
                    continue
                group_metrics = item.get("metrics", [])
                if not isinstance(group_metrics, list):
                    continue
                group_metrics = [str(m) for m in group_metrics if str(m)]
                if selected_metrics:
                    filtered_group_metrics = [m for m in group_metrics if m in selected_metrics]
                    if not filtered_group_metrics:
                        continue
                else:
                    filtered_group_metrics = group_metrics

                # If metrics are filtered and the group is only partially selected,
                # rebuild judge/trace queries from explanation_paths so only the
                # requested metrics' explanations appear in SELECT and WHERE.
                needs_rebuild = metrics and set(filtered_group_metrics) != set(group_metrics)
                if needs_rebuild:
                    # Collect explanation_paths for the filtered metrics only.
                    exp_by_metric: dict[str, list[dict]] = {}
                    for exp in (schema.explanation_paths or []):
                        if not isinstance(exp, dict):
                            continue
                        m = str(exp.get("metric", "")).strip()
                        p = exp.get("path")
                        t = str(exp.get("type", "direct")).strip()
                        if m and isinstance(p, str) and p.strip() and t != "not_found":
                            exp_by_metric.setdefault(m, []).append(exp)
                    group_items: list[dict] = []
                    seen: set[tuple] = set()
                    for m in filtered_group_metrics:
                        for exp in exp_by_metric.get(m, []):
                            key = (str(exp.get("metric","")), str(exp.get("path","")), str(exp.get("type","")), str(exp.get("json_key","")))
                            if key not in seen:
                                seen.add(key)
                                group_items.append(exp)
                    judge_q = _build_judge_query(
                        experiment_id=experiment_id,
                        schema=schema,
                        explanation_paths=group_items,
                        metrics_for_filter=filtered_group_metrics,
                    )
                    trace_q = _build_trace_query(
                        experiment_id=experiment_id,
                        schema=schema,
                        explanation_paths=group_items,
                        metrics_for_filter=filtered_group_metrics,
                    )
                else:
                    judge_q = _bind_experiment_id(item.get("judge_query_template"))
                    trace_q = _bind_experiment_id(item.get("trace_query_template"))

                out.append({
                    "group_id": item.get("group_id"),
                    "path": None,
                    "type": item.get("group_type"),
                    "metrics": filtered_group_metrics,
                    "get_judge_responses": judge_q,
                    "get_specific_trace": trace_q,
                })
            if out:
                return out

        selected_metrics = metrics[:] if metrics else (schema.available_metrics[:] if schema.available_metrics else [])
        selected_set = set(selected_metrics)

        # Index explanation paths by metric for quick lookup.
        explanation_by_metric: dict[str, list[dict]] = {}
        for exp in (filtered_explanation_paths or []):
            if not isinstance(exp, dict):
                continue
            metric_name = str(exp.get("metric", "")).strip()
            path = exp.get("path")
            exp_type = str(exp.get("type", "direct")).strip()
            if not metric_name or not isinstance(path, str) or not path.strip() or exp_type == "not_found":
                continue
            if selected_set and metric_name not in selected_set:
                continue
            explanation_by_metric.setdefault(metric_name, []).append(exp)

        grouped_queries: list[dict] = []
        covered_metrics: set[str] = set()
        metric_groups = getattr(schema, "metric_groups", None) if schema else None

        if isinstance(metric_groups, list) and metric_groups:
            for group in metric_groups:
                raw_metrics = group.get("metrics", [])
                if not isinstance(raw_metrics, list):
                    continue
                group_metrics = [str(m) for m in raw_metrics if str(m) in selected_set]
                if not group_metrics:
                    continue

                # Deduplicate explanation paths in this group.
                group_items: list[dict] = []
                seen_items: set[tuple[str, str, str, str]] = set()
                for metric_name in group_metrics:
                    for exp in explanation_by_metric.get(metric_name, []):
                        key = (
                            str(exp.get("metric", "")),
                            str(exp.get("path", "")),
                            str(exp.get("type", "")),
                            str(exp.get("json_key", "")),
                        )
                        if key in seen_items:
                            continue
                        seen_items.add(key)
                        group_items.append(exp)

                covered_metrics.update(group_metrics)
                path_values = sorted({str(item.get("path", "")).strip() for item in group_items if item.get("path")})
                group_type = str(group.get("group_type", "")).strip() or ("rule_based" if not group_items else "llm_judged")
                group_id = str(group.get("id", "")).strip() or f"group_{len(grouped_queries) + 1}"

                grouped_queries.append({
                    "group_id": group_id,
                    "path": ", ".join(path_values) if path_values else None,
                    "type": group_type,
                    "metrics": group_metrics,
                    "get_judge_responses": _build_judge_query(
                        experiment_id=experiment_id,
                        schema=schema,
                        explanation_paths=group_items,
                        metrics_for_filter=group_metrics,
                    ),
                    "get_specific_trace": _build_trace_query(
                        experiment_id=experiment_id,
                        schema=schema,
                        explanation_paths=group_items,
                        metrics_for_filter=group_metrics,
                    ),
                })
        else:
            groups = _build_explanation_groups(
                explanation_paths=filtered_explanation_paths or [],
                metric_types=schema.metric_types if getattr(schema, "metric_types", None) else None,
            )
            for group in groups:
                group_items = group.get("items", [])
                group_metrics = group.get("metrics", [])
                if isinstance(group_metrics, list):
                    covered_metrics.update([str(m) for m in group_metrics if str(m) in selected_set])
                grouped_queries.append({
                    "group_id": group.get("id"),
                    "path": group.get("path"),
                    "type": group.get("type"),
                    "metrics": group_metrics,
                    "get_judge_responses": _build_judge_query(
                        experiment_id=experiment_id,
                        schema=schema,
                        explanation_paths=group_items,
                        metrics_for_filter=group_metrics,
                    ),
                    "get_specific_trace": _build_trace_query(
                        experiment_id=experiment_id,
                        schema=schema,
                        explanation_paths=group_items,
                        metrics_for_filter=group_metrics,
                    ),
                })

        # Ensure rule/no-explanation metrics also get a grouped query.
        remaining_metrics = [m for m in selected_metrics if m not in covered_metrics]
        if remaining_metrics:
            grouped_queries.append({
                "group_id": f"group_{len(grouped_queries) + 1}",
                "path": None,
                "type": "rule_based",
                "metrics": remaining_metrics,
                "get_judge_responses": _build_judge_query(
                    experiment_id=experiment_id,
                    schema=schema,
                    explanation_paths=[],
                    metrics_for_filter=remaining_metrics,
                ),
                "get_specific_trace": _build_trace_query(
                    experiment_id=experiment_id,
                    schema=schema,
                    explanation_paths=[],
                    metrics_for_filter=remaining_metrics,
                ),
            })
        return grouped_queries

    # If schema has cached templates, use them directly.
    if schema and schema.scores_query_template and schema.judge_query_template and schema.trace_query_template:
        scores_query = schema.scores_query_template.replace("{experiment_id}", experiment_id)

        # Always rebuild judge/trace from the filtered paths so metric filters stay explicit.
        metrics_for_filter = metrics[:] if metrics else (schema.available_metrics[:] if schema.available_metrics else None)
        judge_query = _build_judge_query(
            experiment_id=experiment_id,
            schema=schema,
            explanation_paths=filtered_explanation_paths,
            metrics_for_filter=metrics_for_filter,
        )
        trace_query = _build_trace_query(
            experiment_id=experiment_id,
            schema=schema,
            explanation_paths=filtered_explanation_paths,
            metrics_for_filter=metrics_for_filter,
        )

        template_data["queries"] = {
            "get_scores": {
                "query": scores_query,
                "description": "Get scores for all traces in the experiment"
            },
            "get_judge_responses": {
                "query": judge_query,
                "description": "Get judge responses (explanations) for all traces"
            },
            "get_specific_trace": {
                "query": trace_query,
                "description": "Get specific trace by root_span_id",
                "note": "Replace {{root_span_id}} with the actual span ID"
            }
        }

        if filtered_explanation_paths:
            direct_count = sum(1 for exp in filtered_explanation_paths
                              if isinstance(exp, dict) and exp.get("type") == "direct" and exp.get("path"))
            if direct_count > 0:
                template_data["queries"]["get_judge_responses"]["note"] = "Direct-type explanations are combined into 'judge_response' field as a JSON string"

        grouped_queries = _build_grouped_dimension_queries()
        if grouped_queries:
            template_data["queries"]["get_judge_responses_grouped"] = grouped_queries

        if format == "json":
            print(json.dumps(template_data, separators=(',', ':')))
            return

        # Console format output
        console.print("\n[cyan]1. Get Scores:[/cyan]")
        print(scores_query)

        console.print("\n[cyan]2. Get Judge Responses:[/cyan]")
        print(judge_query)

        if filtered_explanation_paths:
            direct_count = sum(1 for exp in filtered_explanation_paths
                              if isinstance(exp, dict) and exp.get("type") == "direct" and exp.get("path"))
            if direct_count > 0:
                console.print("[dim]Note: Direct-type explanations are combined into 'judge_response' field as a JSON string.[/dim]")
        if grouped_queries:
            console.print(f"[dim]Grouped dimension queries: {len(grouped_queries)} (key: get_judge_responses_grouped).[/dim]")
            console.print("[dim]Below are per-group judge queries (recommended for execution).[/dim]")
            for idx, grouped in enumerate(grouped_queries, start=1):
                group_id = grouped.get("group_id")
                group_metrics = grouped.get("metrics")
                console.print(f"\n[cyan]2.{idx} Group {group_id} Judge Responses:[/cyan]")
                if group_metrics:
                    console.print(f"[dim]metrics: {group_metrics}[/dim]")
                print(grouped.get("get_judge_responses", ""))

        console.print("\n[cyan]3. Get Specific Trace:[/cyan]")
        print(trace_query)
        console.print("[dim]Replace {{root_span_id}} with the actual span ID.[/dim]")
        return

    # Otherwise, build queries from schema (legacy path)
    # Part 1: Get scores
    scores_query = f"from: experiment('{experiment_id}') | select: root_span_id, scores | dimensions: root_span_id, scores | filter: scores IS NOT NULL"

    # Part 2: Get judge responses (explanations)
    # Build dimension fields with aliases based on schema
    dimension_fields = ["root_span_id"]

    direct_explanations = []
    json_explanations = []
    # Use filtered explanation paths
    if filtered_explanation_paths:
        # Separate direct and json type explanations
        json_paths = set()
        
        for exp_info in filtered_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 = _convert_path_to_btql(path)
                if exp_type == "direct":
                    # Collect direct type explanations to concatenate
                    direct_explanations.append({
                        "metric": metric,
                        "path": btql_path
                    })
                elif exp_type == "json" and path not in json_paths:
                    # Keep json type paths separate
                    json_paths.add(path)
                    alias = _create_alias_from_path(path)
                    json_explanations.append({
                        "metric": alias,
                        "path": btql_path,
                    })
        
        # Combine direct and json explanations into a unified judge_response field
        all_explanation_parts = []
        
        # Add direct explanations
        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"]}), \'\'), "\\""'
            )
        
        # Add json explanations
        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"]}), \'\'), "\\""'
            )
        
        # Build the concatenated JSON string if we have any explanations
        if all_explanation_parts:
            # concat('{', "\"", "metric1", "\": \"", path1, "\"", ', ', "\"", "metric2", "\": \"", path2, "\"", '}')
            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"

    judge_query = f"from: experiment('{experiment_id}') | select: root_span_id | dimensions: {dimensions_str} | filter: {filter_str}"

    trace_filter_str = "root_span_id = '{root_span_id}' AND scores IS NOT NULL"

    trace_query = f"from: experiment('{experiment_id}') | select: root_span_id | dimensions: {dimensions_str} | filter: {trace_filter_str}"

    # Populate template_data with all queries
    template_data["queries"] = {
        "get_scores": {
            "query": scores_query,
            "description": "Get scores for all traces in the experiment"
        },
        "get_judge_responses": {
            "query": judge_query,
            "description": "Get judge responses (explanations) for all traces"
        },
        "get_specific_trace": {
            "query": trace_query,
            "description": "Get specific trace by root_span_id",
            "note": "Replace {{root_span_id}} with the actual span ID"
        }
    }

    # Add note about direct explanations if applicable
    if filtered_explanation_paths:
        direct_count = sum(1 for exp in filtered_explanation_paths
                          if isinstance(exp, dict) and exp.get("type") == "direct" and exp.get("path"))
        if direct_count > 0:
            template_data["queries"]["get_judge_responses"]["note"] = "Direct-type explanations are combined into 'judge_response' field as a JSON string"

    grouped_queries = _build_grouped_dimension_queries()
    if grouped_queries:
        template_data["queries"]["get_judge_responses_grouped"] = grouped_queries

    # Output based on format
    if format == "json":
        print(json.dumps(template_data, separators=(',', ':')))
    else:
        # Console format output
        console.print("\n[cyan]1. Get Scores:[/cyan]")
        print(scores_query)

        console.print("\n[cyan]2. Get Judge Responses:[/cyan]")
        print(judge_query)

        # Add helpful note if there are direct explanations
        if filtered_explanation_paths:
            direct_count = sum(1 for exp in filtered_explanation_paths
                              if isinstance(exp, dict) and exp.get("type") == "direct" and exp.get("path"))
            if direct_count > 0:
                console.print("[dim]Note: Direct-type explanations are combined into 'judge_response' field as a JSON string.[/dim]")
        if grouped_queries:
            console.print(f"[dim]Grouped dimension queries: {len(grouped_queries)} (key: get_judge_responses_grouped).[/dim]")
            console.print("[dim]Below are per-group judge queries (recommended for execution).[/dim]")
            for idx, grouped in enumerate(grouped_queries, start=1):
                group_id = grouped.get("group_id")
                group_metrics = grouped.get("metrics")
                console.print(f"\n[cyan]2.{idx} Group {group_id} Judge Responses:[/cyan]")
                if group_metrics:
                    console.print(f"[dim]metrics: {group_metrics}[/dim]")
                print(grouped.get("get_judge_responses", ""))

        console.print("\n[cyan]3. Get Specific Trace:[/cyan]")
        print(trace_query)
        console.print("[dim]Replace {{root_span_id}} with the actual span ID.[/dim]")


def _convert_path_to_btql(path: str) -> str:
    """Convert dot-notation path to BTQL format.

    e.g., "input.output.0.judge_response.text" -> "input.output[0].judge_response.text"
    """
    parts = path.split(".")
    result_parts = []
    i = 0
    while i < len(parts):
        part = parts[i]
        if part.isdigit():
            # This is an array index, append to previous part
            if result_parts:
                result_parts[-1] = f"{result_parts[-1]}[{part}]"
        else:
            result_parts.append(part)
        i += 1
    return ".".join(result_parts)


def _create_alias_from_path(path: str) -> str:
    """Create a readable alias from a path.

    e.g., "input.output.0.judge_response.text" -> "judge_response"
    e.g., "input.output[0].judge_response.text" -> "judge_response"
    e.g., "metadata.completeness.explanation" -> "completeness_explanation"
    """
    import re
    
    # Remove any brackets and their contents (e.g., "[0]", "[1]")
    path_clean = re.sub(r'\[\d+\]', '', path)
    
    # Split by dots
    parts = path_clean.split(".")
    
    # Remove numeric parts, empty parts, and common prefixes
    meaningful_parts = [
        p for p in parts 
        if p and not p.isdigit() and p not in ("input", "output", "metadata", "text")
    ]

    if not meaningful_parts:
        return "response"

    # Take last 1-2 meaningful parts
    if len(meaningful_parts) >= 2:
        return f"{meaningful_parts[-2]}_{meaningful_parts[-1]}"
    return meaningful_parts[-1]


def _build_metric_filter_expr(metrics: list[str] | None) -> str:
    """Build BTQL score presence filter for a metric list."""
    if not metrics:
        return ""
    checks = []
    for metric in metrics:
        safe_metric = metric.replace('"', '\\"')
        checks.append(f'scores."{safe_metric}" IS NOT NULL')
    if not checks:
        return ""
    return "(" + " OR ".join(checks) + ")"


def _build_explanation_groups(
    explanation_paths: list,
    metric_types: dict[str, str] | None = None,
) -> list[dict]:
    """Group explanation paths by source path/type for LLM-judged style metrics.

    This allows generating multiple dimension queries when different judge scorers
    store explanations in different structures.
    """
    metric_types = metric_types or {}
    allowed_types = {"llm_judged", "hybrid"}
    groups: dict[str, dict] = {}
    ordered_keys: list[str] = []

    for exp_info in explanation_paths:
        if not isinstance(exp_info, dict):
            continue
        metric = str(exp_info.get("metric", "")).strip()
        raw_path = exp_info.get("path")
        path = raw_path.strip() if isinstance(raw_path, str) else ""
        exp_type = str(exp_info.get("type", "direct")).strip()
        if not metric or not path or exp_type == "not_found":
            continue

        metric_type = metric_types.get(metric, "hybrid")
        if metric_type not in allowed_types:
            continue

        key = f"{exp_type}::{path}"
        if key not in groups:
            groups[key] = {
                "id": f"group_{len(ordered_keys) + 1}",
                "path": path,
                "type": exp_type,
                "items": [],
                "metrics": [],
            }
            ordered_keys.append(key)

        groups[key]["items"].append(exp_info)
        if metric not in groups[key]["metrics"]:
            groups[key]["metrics"].append(metric)

    return [groups[k] for k in ordered_keys]


def _build_judge_query(experiment_id: str, schema, explanation_paths: list, metrics_for_filter: list[str] | None = None) -> str:
    """Build judge query with filtered explanation paths.
    
    Args:
        experiment_id: The experiment ID
        schema: Schema information
        explanation_paths: Filtered list of explanation paths
    
    Returns:
        BTQL query string for fetching judge responses
    """
    dimension_fields = ["root_span_id", "input.input AS input"]
    direct_explanations = []
    json_explanations = []
    json_paths = set()
    
    for exp_info in 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 = _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 = _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('"', '\\"')
        all_explanation_parts.append(
            f'"\\"", "{safe_metric}", "\\": \\"", coalesce(to_string({exp["path"]}), \'\'), "\\""'
        )
    
    for exp in json_explanations:
        safe_metric = exp["metric"].replace('"', '\\"')
        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')
    
    dimensions_str = ", ".join(dimension_fields)
    
    filter_parts = ["scores IS NOT NULL"]
    metric_filter = _build_metric_filter_expr(metrics_for_filter)
    if metric_filter:
        filter_parts.append(metric_filter)
    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_parts.append("(" + " OR ".join(explanation_checks) + ")")
    filter_str = " AND ".join(filter_parts)

    return (
        f"from: experiment('{experiment_id}') | "
        f"select: root_span_id | "
        f"dimensions: {dimensions_str} | "
        f"filter: {filter_str}"
    )


def _build_trace_query(experiment_id: str, schema, explanation_paths: list, metrics_for_filter: list[str] | None = None) -> str:
    """Build trace query with filtered explanation paths.
    
    Args:
        experiment_id: The experiment ID (will use placeholder in result)
        schema: Schema information
        explanation_paths: Filtered list of explanation paths
    
    Returns:
        BTQL query string for fetching specific trace
    """
    dimension_fields = ["root_span_id", "input.input AS input"]
    direct_explanations = []
    json_explanations = []
    json_paths = set()
    
    for exp_info in 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 = _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 = _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('"', '\\"')
        all_explanation_parts.append(
            f'"\\"", "{safe_metric}", "\\": \\"", coalesce(to_string({exp["path"]}), \'\'), "\\""'
        )
    
    for exp in json_explanations:
        safe_metric = exp["metric"].replace('"', '\\"')
        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')
    
    dimensions_str = ", ".join(dimension_fields)
    
    trace_filter_parts = ["root_span_id = '{root_span_id}'", "scores IS NOT NULL"]
    metric_filter = _build_metric_filter_expr(metrics_for_filter)
    if metric_filter:
        trace_filter_parts.append(metric_filter)
    trace_explanation_checks = []
    for exp in direct_explanations:
        trace_explanation_checks.append(f"{exp['path']} IS NOT NULL")
    for exp in json_explanations:
        trace_explanation_checks.append(f"{exp['path']} IS NOT NULL")
    if trace_explanation_checks:
        trace_filter_parts.append("(" + " OR ".join(trace_explanation_checks) + ")")
    trace_filter_str = " AND ".join(trace_filter_parts)
    
    return (
        f"from: experiment('{experiment_id}') | "
        f"select: root_span_id | "
        f"dimensions: {dimensions_str} | "
        f"filter: {trace_filter_str}"
    )


@click.command()
@click.option(
    "--experiment-id",
    "-e",
    required=True,
    help="Braintrust experiment ID",
)
@click.option(
    "--metrics",
    "-m",
    help="Comma-separated list of metrics to fetch (uses all available metrics if not specified)",
)
@click.option(
    "--max-iterations",
    "-i",
    default=5,
    help="Maximum query refinement iterations (default: 5)",
)
@click.option(
    "--output",
    "-o",
    type=click.Choice(["table", "json"]),
    default="table",
    help="Output format (default: table)",
)
@click.option(
    "--verbose",
    "-v",
    is_flag=True,
    help="Show detailed refinement steps",
)
@click.option(
    "--filter",
    "-f",
    "filters",
    multiple=True,
    help="Additional filter conditions (can be used multiple times)",
)
@click.option(
    "--no-default-filters",
    is_flag=True,
    help="Disable default filters",
)
@click.option(
    "--use-llm-initial",
    is_flag=True,
    help="Use LLM to generate initial query instead of template",
)
@click.option(
    "--show-judge",
    "-j",
    is_flag=True,
    help="Show judge responses (explanations for scores)",
)
@click.option(
    "--judge-limit",
    default=5,
    help="Number of judge response samples to show (default: 5)",
)
@click.option(
    "--refresh-schema",
    is_flag=True,
    help="Force refresh schema cache (re-analyze experiment structure with LLM)",
)
@click.option(
    "--eval",
    required=True,
    help="Evaluation name for schema cache key. Different experiments with same eval share schema.",
)
@click.option(
    "--braintrust-api-key",
    help="Braintrust API key (overrides BRAINTRUST_API_KEY environment variable)",
)
@click.option(
    "--braintrust-org",
    type=click.Choice(["bt-azure", "non-pme"]),
    default=None,
    help="Braintrust organization: bt-azure or non-pme (overrides BRAINTRUST_ORG env var, default: non-pme)",
)
@click.option(
    "--template-format",
    type=click.Choice(["console", "json"]),
    default="console",
    help="Query template output format (default: console)",
)
@click.option(
    "--model-type",
    type=click.Choice(["azure", "papyrus", "copilot", "local"]),
    default=None,
    help="Model client type (overrides MODEL_CLIENT_TYPE env var, default: azure)",
)
@click.option(
    "--model-endpoint",
    default=None,
    help="Model endpoint URL. For papyrus: Papyrus endpoint. For copilot: backend URL (default: http://localhost:8237). For local: OpenAI-compatible base URL (default: http://localhost:23333/api/openai/v1)",
)
@click.option(
    "--model-name",
    default=None,
    help="Model name/ID to use. For azure: deployment name. For papyrus: model deployment. Required for copilot and local types.",
)
@click.option(
    "--cache-dir",
    default=None,
    help="Directory containing schema_cache.json. Defaults to ~/.btql-agent/ if not specified.",
)
def main(
    experiment_id: str,
    metrics: str | None,
    max_iterations: int,
    output: str,
    verbose: bool,
    filters: tuple,
    no_default_filters: bool,
    use_llm_initial: bool,
    show_judge: bool,
    judge_limit: int,
    refresh_schema: bool,
    eval: str,
    braintrust_api_key: str | None,
    braintrust_org: str | None,
    template_format: str,
    model_type: str | None,
    model_endpoint: str | None,
    model_name: str | None,
    cache_dir: str | None,
) -> None:
    """BTQL Agent - Auto-generate BTQL queries to fetch experiment metrics.

    The agent automatically discovers the data schema using LLM analysis.
    Schema is cached to the --cache-dir directory (default: ~/.btql-agent/) for faster subsequent runs.

    Example:
        btql-agent -e abc123 --eval "picasso_v1"  # Use all available metrics (Azure)
        btql-agent -e abc123 -m "completeness,groundedness" --eval "picasso_v1"
        btql-agent -e abc123 --eval "picasso_v1" --model-type papyrus --model-name gpt-51-chat
        btql-agent -e abc123 --eval "picasso_v1" --model-type copilot --model-name gpt-4o
        btql-agent -e abc123 --eval "picasso_v1" --model-type local --model-name qwen-72b

    Model Types:
        azure   - Azure OpenAI (default, uses DefaultAzureCredential)
        papyrus - Papyrus (uses DefaultAzureCredential with Papyrus scope)
        copilot - Copilot via VS Code extension proxy (requires --model-name)
        local   - Local OpenAI-compatible server (requires --model-name)
    """
    # Lazy imports - only load when actually executing (not for --help)
    from .agent import BTQLAgent

    # Get configuration with model type options
    config = get_config_from_env(
        max_iterations,
        braintrust_api_key,
        braintrust_org,
        model_type,
        model_endpoint,
        model_name,
    )
    
    # Create agent with optional cache directory
    from pathlib import Path
    cache_path = Path(cache_dir) if cache_dir else None
    agent = BTQLAgent(config, cache_dir=cache_path)
    
    # Parse metrics if provided
    metric_list = [m.strip() for m in metrics.split(",")] if metrics else None

    # Override default filters if requested
    if no_default_filters:
        config.default_filters = []

    # Use eval as cache key
    cached_schema = agent.get_cached_schema(eval)
    schema_status = "refreshing" if refresh_schema else ("cached" if cached_schema else "discovering")

    # Only show console output if not using JSON template format
    if template_format != "json":
        console.print(f"[bold]Experiment:[/bold] {experiment_id}")
        console.print(f"[bold]Eval:[/bold] {eval}")
        if metric_list:
            console.print(f"[bold]Metrics:[/bold] {', '.join(metric_list)}")
        else:
            console.print(f"[bold]Metrics:[/bold] [yellow]All available (to be discovered)[/yellow]")
        console.print(f"[bold]Schema:[/bold] {schema_status}")
        console.print()

    # Discover or get cached schema
    schema = asyncio.run(
        agent.discover_schema(
            eval, 
            force_refresh=refresh_schema, 
            experiment_id=experiment_id, 
            max_attempts=max_iterations
        )
    )
    
    if not schema.available_metrics:
        console.print("[red]Error:[/red] No metrics found in schema. Cannot proceed.")
        console.print("[yellow]Tip:[/yellow] Check the experiment has scored data.")
        sys.exit(1)
    
    if template_format != "json" and not metric_list:
        console.print(f"[yellow]Discovered metrics:[/yellow] {', '.join(schema.available_metrics)}")
        console.print()
    
    # Display final query templates
    display_final_query_template(experiment_id, schema, template_format, eval, metric_list)


if __name__ == "__main__":
    main()
