from __future__ import annotations

import asyncio
import gzip
import io
import hashlib
import json
import logging
import os
import re
import subprocess
import time
from dataclasses import dataclass
from decimal import Decimal, InvalidOperation
from heapq import heappush, heapreplace
from pathlib import Path
from typing import Any, Protocol

import aiofiles
import httpx
import random
import sys
import shutil
from prompt_auto_lab.core.auth_token import get_azure_token_provider
from autogen_ext.models.openai import AzureOpenAIChatCompletionClient
from prompt_auto_lab.core.papyrus_chat_completion_client import PapyrusChatCompletionClient
from prompt_auto_lab.config.braintrust import BRAINTRUST_ORG_ENDPOINTS

from prompt_auto_lab.core.model_clients import (
    build_azure_model_client,
    build_azure_model_thought_client,
    build_papyrus_model_client,
)
from prompt_auto_lab.core.picasso_workspace import SPARSE_CHECKOUT_PATHS

logger = logging.getLogger(__name__)


def _normalize_btql_agent_model_type(model_client_type: str | None) -> str:
    raw = (model_client_type or "azure").strip().lower()
    mapping = {
        "azure": "azure",
        "papyrus": "papyrus",
        "copilot": "copilot",
        "local": "local",
    }
    return mapping.get(raw, "azure")

DEFAULT_MAX_TEXT_CHARS = int(os.getenv("MAX_TOOL_TEXT_CHARS", "800"))
DEFAULT_MAX_DIFF_CHARS = int(os.getenv("MAX_DIFF_TEXT_CHARS", "2000"))


def _truncate_text(value: Any, max_chars: int = DEFAULT_MAX_TEXT_CHARS) -> tuple[Any, bool, int]:
    if not isinstance(value, str):
        return value, False, 0
    original_len = len(value)
    if max_chars <= 0 or original_len <= max_chars:
        return value, False, original_len
    if max_chars <= 3:
        return value[:max_chars], True, original_len
    head = max_chars // 2
    tail = max_chars - head - 3
    if tail < 1:
        truncated = value[: max_chars - 3] + "..."
    else:
        truncated = value[:head] + "..." + value[-tail:]
    return truncated, True, original_len


def _truncate_record_fields(
    record: dict[str, Any],
    fields: tuple[str, ...] | list[str],
    max_chars: int = DEFAULT_MAX_TEXT_CHARS,
) -> dict[str, Any]:
    out = dict(record)
    for field in fields:
        truncated, did_truncate, original_len = _truncate_text(out.get(field), max_chars=max_chars)
        if did_truncate:
            out[field] = truncated
            out[f"{field}_truncated"] = True
            out[f"{field}_original_len"] = original_len
    return out

async def run_subprocess_async(
    args: list[str],
    cwd: str | None = None,
    capture_output: bool = True,
    text: bool = True,
    timeout: float | None = None,
    env: dict[str, str] | None = None,
    encoding: str | None = None,
) -> subprocess.CompletedProcess:
    """Run subprocess.run in a thread pool to avoid blocking the event loop."""
    def _run():
        kwargs: dict[str, Any] = {
            "cwd": cwd,
            "capture_output": capture_output,
            "text": text,
            "timeout": timeout,
        }
        if env is not None:
            kwargs["env"] = env
        if encoding is not None:
            kwargs["encoding"] = encoding
        return subprocess.run(args, **kwargs)
    return await asyncio.to_thread(_run)


def build_model_clients():
    """Build main, think, and analysis model clients based on configuration.

    This centralizes the logic for choosing between Papyrus and Azure models
    so that callers (like ``main.py``) do not need to duplicate it.

    Returns:
        :class:`~prompt_auto_lab.core.orchestrator_builder.ModelClients` with
        ``main``, ``think``, and ``analysis`` roles populated.

        For Azure deployments the ``analysis`` client is built from dedicated
        env-vars (``AZURE_OPENAI_ANALYSIS_ENDPOINT`` /
        ``AUTOPROMPTLAB_ANALYSIS_DEPLOYMENT``) when both are set; otherwise
        it falls back to ``main`` (same object).
    """
    from prompt_auto_lab.core.orchestrator_builder import ModelClients

    token_provider: Any | None = None
    model_client: Any | None = None
    model_think_client: Any | None = None
    model_client_type = os.getenv("MODEL_CLIENT", "Azure")
    papyrus_scope = os.getenv("PAPYRUS_VERRIFY_SCOPES") or os.getenv("PAPYRUS_VERIFY_SCOPES")
    azure_scope = os.getenv("AZURE_VERIFY_SCOPES", "https://cognitiveservices.azure.com/.default")

    if model_client_type == "Papyrus":
        if not papyrus_scope:
            raise ValueError("PAPYRUS_VERIFY_SCOPES environment variable is required for Papyrus client.")
        token_provider = get_azure_token_provider(papyrus_scope)
        papyrus_endpoint = os.getenv("PAPYRUS_ENDPOINT")
        papyrus_model = os.getenv("PAPYRUS_MODEL")
        papyrus_quota_id = os.getenv("PAPYRUS_QUOTA_ID", "")
        if not papyrus_endpoint or not papyrus_model:
            raise ValueError("PAPYRUS_ENDPOINT and PAPYRUS_MODEL environment variables are required for Papyrus client.")
        model_client = build_papyrus_model_client(token_provider=token_provider, endpoint=papyrus_endpoint, model=papyrus_model, quota_id=papyrus_quota_id)
        model_think_client = build_papyrus_model_client(token_provider=token_provider, endpoint=papyrus_endpoint, model=papyrus_model, quota_id=papyrus_quota_id)
        # Papyrus has no separate analysis deployment — reuse main
        return ModelClients(main=model_client, think=model_think_client, analysis=model_client)
    elif model_client_type == "Azure":
        token_provider = get_azure_token_provider(azure_scope)
        deployment = os.environ.get("MODEL_DEPLOYMENT")
        azure_endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
        azure_api_version = os.environ.get("AZURE_API_VERSION")
        if not deployment or not azure_endpoint or not azure_api_version:
            raise ValueError("MODEL_DEPLOYMENT, AZURE_OPENAI_ENDPOINT, and AZURE_API_VERSION environment variables are required for Azure client.")
        model_client = build_azure_model_client(
            token_provider=token_provider,
            deployment=deployment,
            endpoint=azure_endpoint,
            api_version=azure_api_version,
        )
        model_think_client = build_azure_model_thought_client(
            token_provider=token_provider,
            endpoint=azure_endpoint,
            api_version=azure_api_version,
        )

        # Build dedicated analysis client when env vars are present
        analysis_endpoint = os.environ.get("AZURE_OPENAI_ANALYSIS_ENDPOINT")
        analysis_deployment = os.environ.get("AUTOPROMPTLAB_ANALYSIS_DEPLOYMENT")
        if analysis_endpoint and analysis_deployment:
            analysis_timeout_s = float(os.environ.get("AUTOPROMPTLAB_ANALYSIS_TIMEOUT_S", "120"))
            analysis_max_retries = int(os.environ.get("AUTOPROMPTLAB_ANALYSIS_MAX_RETRIES", "3"))
            analysis_client = build_azure_model_client(
                token_provider=token_provider,
                deployment=analysis_deployment,
                endpoint=analysis_endpoint,
                timeout_s=analysis_timeout_s,
                max_retries=analysis_max_retries,
            )
        else:
            # Graceful fallback: reuse main client
            analysis_client = model_client

        return ModelClients(main=model_client, think=model_think_client, analysis=analysis_client)

    else:
        raise ValueError(f"Unsupported MODEL_CLIENT type: {model_client_type!r}")


# class AipaEvalClient:
#     """Client for Aipa evaluation."""
#     def __init__(
#         self,
#         backend: Any = None,
#         eval_path: str = "",
#         conda_python_path: str = "",
#         eval_name: str = "",
#         dataset_name: str = "",
#         local_debug: bool = False,
#         max_dataset_rows: int | None = None,
#     ) -> None:
#         self.backend = backend
#         self.eval_path = eval_path
#         self.conda_python_path = conda_python_path
#         self.eval_name = eval_name
#         self.dataset_name = dataset_name
#         self.local_debug = local_debug
#         if max_dataset_rows is None:
#             self.max_dataset_rows = int(os.environ.get("MAX_DATASET_ROWS", "300"))
#         else:
#             self.max_dataset_rows = max_dataset_rows

#     async def submit_eval(self, current_branch: str, task_id: str, experiment_suffix: str) -> dict[str, Any]:
#         """
#         Submit an eval job using PowerShell command.

#         Args:
#             current_branch: The current git branch name
#             task_id: The task ID for the eval
#             experiment_suffix: The experiment suffix for the eval

#         Returns:
#             Dict containing the eval output log
#         """
#         try:
#             # Build the eval command with conda aipa environment using conda run
#             eval_command = (
#                 f"{self.conda_python_path} -m evals.run "
#                 f"--generator_config_yaml aipa_evals\\llm\\configs\\staging_picasso_cn.yaml "
#                 f"--eval_project_name user_xuyanmo_test "
#                 f"--eval {self.eval_name} "
#                 f"--dataset_name {self.dataset_name} "
#                 f"--task_id {task_id} "
#                 f"--experiment_suffix {experiment_suffix} "
#                 f"--max_workers 1 "
#                 # f"--ignore_prod_parity "
#                 f"--prompt_branch {current_branch} "
#                 f"--dataset_sampling_method first "
#                 f"--max_dataset_rows {self.max_dataset_rows}"
#             )

#             # Run the command using Popen for real-time output
#             process = subprocess.Popen(
#                 ["powershell", "-Command", eval_command],
#                 stdout=subprocess.PIPE,
#                 stderr=subprocess.STDOUT,
#                 text=True,
#                 encoding="utf-8",
#                 cwd=self.eval_path
#             )

#             output_lines = []
#             # Stream output in real-time
#             for line in process.stdout:
#                 print(line)  # Print to console 
#                 output_lines.append(line)

#             process.wait(timeout=3600)  # Wait for process to complete with 1 hour timeout
#             output_log = ''.join(output_lines)

#             return {
#                 "success": process.returncode == 0,
#                 "output_log": output_log,
#                 "return_code": process.returncode
#             }
#         except subprocess.TimeoutExpired:
#             process.kill()
#             return {"error": "Eval command timed out after 1 hour"}
#         except subprocess.CalledProcessError as e:
#             return {"error": f"Eval command failed: {e.stderr}"}

#     async def parse_experiment_id(self, output_log: str) -> dict[str, Any]:
#         """
#         Parse the result_id from the eval output log.

#         Args:
#             output_log: The output log from the eval command

#         Returns:
#             Dict containing the parsed result_id
#         """
#         # Try to find experiment_id and experiment_link pattern in the log
#         # Example: Experiment xuyanmo_hello_world_base-ef750c07 is running at https://www.braintrust.dev/app/Microsoft%20Copilot/p/leaderboard_hello_world/experiments/xuyanmo_hello_world_base-ef750c07
#         experiment_id_pattern = r"Experiment ID:\s*([a-zA-Z0-9_-]+)"
#         exp_link_pattern = r"https://www\.braintrust\.dev/app/([a-zA-Z0-9_%/\-]+)"

#         exp_id_match = re.search(experiment_id_pattern, output_log, re.IGNORECASE)
#         link_match = re.search(exp_link_pattern, output_log, re.IGNORECASE)
#         if exp_id_match and link_match:
#             real_link = r"https://www.braintrust.dev/app/" + link_match.group(1)
#             return {"result_id": exp_id_match.group(1), "link": real_link, "success": True}   

#         return {"error": "Could not parse result_id or experiment link from output log", "success": False}

#     async def submit_aipa_eval(self, payload_json: str) -> dict[str, Any]:
#         """
#         Submit an Aipa-Eval job based on the provided payload.
#         Args:
#             payload_json: JSON string containing the eval parameters, must include 'branch' and 'dataset_name'.
#         Returns:
#             Dict containing the experiment_id or error message
#         """
#         if self.local_debug:
#             suffix = random.randint(1000, 9999)
#             return {"experiment_id": f"cad3d72a-4813-4a23-8a57-775e63730a4c-{suffix}","experiment_link": "https://www.braintrust.dev/app/Microsoft%20Copilot%20PME/p/Edge_CN_CopilotContextualQuality"}
#         #
#         payload = json.loads(payload_json)
#         if not isinstance(payload, dict):
#             return {"error": "payload_json must decode to an object"}
#         if payload.get("branch") is None:
#             return {"error": "branch is required in payload for picasso eval."}
#         if payload.get("dataset_name") is not None:
#             self.dataset_name = payload.get("dataset_name")
        
#         time_str = time.strftime("%Y%m%d_%H%M%S")
#         task_id = "auto_prompt_lab_test_task_" + time_str
#         experiment_suffix = "auto_prompt_lab_experiment_" + time_str
        
#         eval_result = await self.submit_eval(payload.get("branch"), task_id, experiment_suffix)

#         if "output_log" in eval_result:
#             parse_result = await self.parse_experiment_id(eval_result["output_log"])
#             if parse_result.get("success"):
#                 return {"experiment_id": parse_result["result_id"], "experiment_link": parse_result["link"]}
#             else:
#                 return {"error": "Failed to parse experiment_id from eval output."}
#         else:
#             return {"error": "Eval submission failed."}

#     async def submit_job(self, payload: dict[str, Any]) -> str:
#         """Submit a job to the backend."""
#         job_id = str(payload.get("job_id") or "stub-job")
#         if self.backend is not None:
#             self.backend.jobs[job_id] = dict(payload)
#         return job_id
    
class GitHubActionEvalClient:
    """Client for Aipa evaluation in github actions."""

    # Default picasso repository NWO
    AIPA_EVAL_REPO = "infinity-microsoft/picasso"

    def __init__(
        self,
        backend: Any = None,
        github_cli_path: str = "",
        eval_name: str = "",
        dataset_name: str = "",
        local_debug: bool = False,
        max_dataset_rows: int | None = None,
        dataset_sampling_method: str | None = None,
        aipa_eval_branch: str = "main",
        conversation_mode: str = "Smart",
        output_dir: str | Path | None = None,
        eval_project_name: str | None = None,
        dataset_project_name: str | None = None,
        add_feature_flags: str | None = None,
        baseline_feature_flags: str | None = None,
        picasso_version: str | None = None,
        braintrust_org: str | None = None,
        environment: str | None = None,
        model: str | None = None,
        additional_cmd_args: str | None = None,
        picasso_env: str | None = None,
        task_id_prefix: str | None = None,
    ) -> None:
        self.backend = backend
        if github_cli_path:
            self.github_cli_path = github_cli_path
        else:
            self.github_cli_path = str(
                Path(__file__).resolve().parent.parent / "tools" / "trigger_picasso_eval.ps1"
            )
        self.eval_name = eval_name
        self.dataset_name = dataset_name
        self.local_debug = local_debug
        self.aipa_eval_branch = aipa_eval_branch
        self.conversation_mode = conversation_mode
        self.output_dir = output_dir
        self.picasso_version = picasso_version or os.environ.get("PICASSO_VERSION", "Baseline")
        if max_dataset_rows is None:
            self.max_dataset_rows = int(os.environ.get("MAX_DATASET_ROWS", "300"))
        else:
            self.max_dataset_rows = max_dataset_rows
        # Load eval project settings from environment or use defaults
        self.eval_project_name = eval_project_name or os.environ.get("EVAL_PROJECT_NAME", "user_jialyang_test")
        self.dataset_project_name = dataset_project_name or os.environ.get("DATASET_PROJECT_NAME", "")
        self.add_feature_flags = add_feature_flags or os.environ.get("ADD_FEATURE_FLAGS", "")
        self.baseline_feature_flags = baseline_feature_flags or os.environ.get("BASELINE_FEATURE_FLAGS", "")
        self.braintrust_org = braintrust_org or os.environ.get("BRAINTRUST_ORG", "non-pme")
        self.environment = environment
        self.model = model
        self.dataset_sampling_method = dataset_sampling_method or "first"
        self.additional_cmd_args = additional_cmd_args or ""
        self.picasso_env = picasso_env
        self.task_id_prefix = task_id_prefix
        # Track all running subprocesses for cleanup (keyed by branch for async, list for sync fallback)
        self._async_processes: dict[str, asyncio.subprocess.Process] = {}
        self._sync_processes: list[subprocess.Popen] = []
        # Track active eval branches so cancellation targets the right branch
        self._active_branches: set[str] = set()
        self._last_branch: str | None = None

    def _resolve_powershell_executable(self) -> str:
        """Resolve a PowerShell executable name available on current platform."""
        candidates = ["powershell", "pwsh"] if os.name == "nt" else ["pwsh", "powershell"]
        for candidate in candidates:
            if shutil.which(candidate):
                return candidate
        raise FileNotFoundError(
            "PowerShell executable not found. Install PowerShell and ensure 'pwsh' (or 'powershell') is in PATH."
        )

    async def submit_eval(self, picasso_current_branch: str, task_id: str | None, experiment_suffix: str | None, run_leaderboard: bool = False) -> dict[str, Any]:
        """
        Submit an eval job using PowerShell command.

        Args:
            picasso_current_branch: The current git branch name for Picasso
            task_id: The task ID for the eval
            experiment_suffix: The experiment suffix for the eval
            run_leaderboard: If True, adds EvalTags for leaderboard evaluation

        Returns:
            Dict containing the eval output log
        """
        self._active_branches.add(picasso_current_branch)
        self._last_branch = picasso_current_branch
        try:
            picasso_version = self.picasso_version
            environment = self.environment
            model = self.model

            # Common: Build feature_flags with baseline flags if applicable
            feature_flags = self.add_feature_flags
            if picasso_version == "Baseline":
                baseline_flags = self.baseline_feature_flags
                if baseline_flags and self.add_feature_flags:
                    feature_flags = f"{self.add_feature_flags},{baseline_flags}"
                elif baseline_flags:
                    feature_flags = baseline_flags

            # Build cmd_args based on mode
            is_baseline = picasso_version == "Baseline"

            if run_leaderboard:
                # Leaderboard mode: simplified cmd_args
                cmd_args = "--ignore_prod_parity" if is_baseline else ""
            else:
                # Regular eval mode: full cmd_args with eval parameters
                # Extract model version prefix from feature flags if present
                model_version_prefix = ""
                if self.add_feature_flags:
                    import re as regex_module
                    model_match = regex_module.search(r'syndication-model([a-zA-Z0-9\-]+)', self.add_feature_flags)
                    if model_match:
                        model_version_prefix = model_match.group(1) + "_"

                full_suffix = f"{model_version_prefix}{experiment_suffix}" if experiment_suffix else f"{model_version_prefix}"
                cmd_args = (
                    f"--eval {self.eval_name} "
                    f"--dataset_name '{self.dataset_name}' "
                    f"--experiment_suffix {full_suffix} "
                    f"--eval_project_name '{self.eval_project_name}' "
                )
                if self.dataset_sampling_method != "custom":
                    cmd_args += (
                        f"--dataset_sampling_method {self.dataset_sampling_method} "
                        f"--max_dataset_rows {self.max_dataset_rows} "
                    )
                if self.dataset_project_name:
                    cmd_args += f"--dataset_project_name '{self.dataset_project_name}' "
                if is_baseline:
                    cmd_args += "--ignore_prod_parity "
                if self.additional_cmd_args:
                    cmd_args += self.additional_cmd_args
            
            # Common: Build base eval_command
            eval_command = (
                f"{self.github_cli_path} -DownloadLogs "
                f'-EvalRunnerCmdArgs "{cmd_args}" '
                f"-Branch '{picasso_current_branch}' "
                f"-ConversationMode '{self.conversation_mode}' "
                f"-BraintrustOrg '{self.braintrust_org}' "
                f"-FeatureFlags '{feature_flags}' ")
            
            # Add PicassoEnv if set
            if self.picasso_env:
                eval_command += f"-PicassoEnv '{self.picasso_env}' "
            
            # Common: Add Model parameter
            if model:
                eval_command += f"-Model '{model}' "
            else:
                eval_command += "-Model default "
            
            # Configure parameters based on run mode
            if run_leaderboard:
                # Leaderboard mode: add EvalTags, no TaskId
                eval_tags = "leaderboard_china" if environment == "china" else "leaderboard"
                eval_command += f"-EvalTags '{eval_tags}' "
            elif task_id:
                # Regular eval mode: add TaskId, no EvalTags
                eval_command = eval_command.replace("-DownloadLogs ", f"-DownloadLogs -TaskId '{task_id}' ")

            ps_executable = self._resolve_powershell_executable()
            ps_args = [ps_executable]
            if os.name == "nt":
                ps_args.extend(["-ExecutionPolicy", "Bypass"])
            ps_args.extend(["-Command", eval_command])

            # Use async subprocess to avoid blocking the event loop
            process = await asyncio.create_subprocess_exec(
                *ps_args,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
                cwd=None
            )
            self._async_processes[picasso_current_branch] = process

            output_lines = []
            # Async streaming - doesn't block event loop
            async for line in process.stdout:
                decoded = line.decode("utf-8", errors="replace")
                logger.info(decoded.rstrip("\r\n"))
                output_lines.append(decoded)

            await asyncio.wait_for(process.wait(), timeout=7200)  # 2 hour timeout
            output_log = ''.join(output_lines)

            return {
                "success": process.returncode == 0,
                "output_log": output_log,
                "return_code": process.returncode
            }
        except NotImplementedError:
             return await asyncio.to_thread(self._run_sync_subprocess, eval_command)
        except asyncio.TimeoutError:
            process.kill()
            return {"error": "Eval command timed out after 2 hours"}
        except Exception as e:
            return {"error": f"Eval command failed: {str(e)}"}
        finally:
            self._async_processes.pop(picasso_current_branch, None)
            self._active_branches.discard(picasso_current_branch)
    
    def _run_sync_subprocess(self, eval_command: str) -> dict[str, Any]:
        """Synchronous version of subprocess execution for fallback on Windows SelectorEventLoop."""
        try:
            ps_executable = self._resolve_powershell_executable()
            ps_args = [ps_executable]
            if os.name == "nt":
                ps_args.extend(["-ExecutionPolicy", "Bypass"])
            ps_args.extend(["-Command", eval_command])

            process = subprocess.Popen(
                ps_args,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                cwd=None,
                text=True,
                encoding="utf-8",
                errors="replace"
            )
            self._sync_processes.append(process)
            
            output_lines = []            
            # Reads form stdout sync
            for line in process.stdout:
                logger.info(line.rstrip("\r\n"))
                output_lines.append(line)
            
            process.wait(timeout=7200)  # 120 minutes timeout
            output_log = "".join(output_lines)
            try:
                self._sync_processes.remove(process)
            except ValueError:
                pass
            
            result = {
                "success": process.returncode == 0,
                "output_log": output_log,
                "return_code": process.returncode
            }
            return result
        except subprocess.TimeoutExpired:
            process.kill()
            return {"error": "Eval command timed out after 120 minutes"}
        except Exception as e:
            return {"error": f"Eval command failed (threaded): {str(e)}"}

    def kill_subprocess(self) -> None:
        """Kill all running subprocesses. Call this during cleanup to ensure subprocesses are terminated."""
        for branch, process in list(self._async_processes.items()):
            try:
                process.kill()
                logger.info("Killed async subprocess for branch %s", branch)
            except Exception as e:
                logger.warning("Error killing async subprocess for branch %s: %s", branch, e)
        self._async_processes.clear()

        for process in list(self._sync_processes):
            try:
                process.kill()
                process.wait(timeout=5)
                logger.info("Killed sync subprocess")
            except Exception as e:
                logger.warning("Error killing sync subprocess: %s", e)
        self._sync_processes.clear()

    async def parse_experiment_id(self, output_log: str) -> dict[str, Any]:
        """
        Parse the result_id from the eval output log.

        Args:
            output_log: The output log from the eval command

        Returns:
            Dict containing the parsed result_id
        """
        # Try to find experiment_id and experiment_link pattern in the log
        # Example: Experiment xuyanmo_hello_world_base-ef750c07 is running at https://www.braintrust.dev/app/Microsoft%20Copilot/p/leaderboard_hello_world/experiments/xuyanmo_hello_world_base-ef750c07
        experiment_id_pattern = r"Experiment ID:\s*([a-zA-Z0-9_-]+)"
        exp_link_pattern = r"https://www\.braintrust\.dev/app/([a-zA-Z0-9_%/\-]+)"

        exp_id_match = re.search(experiment_id_pattern, output_log, re.IGNORECASE)
        link_match = re.search(exp_link_pattern, output_log, re.IGNORECASE)
        if exp_id_match and link_match:
            real_link = r"https://www.braintrust.dev/app/" + link_match.group(1)
            return {"result_id": exp_id_match.group(1), "link": real_link, "success": True}   

        return {"error": "Could not parse result_id or experiment link from output log", "success": False}

    async def parse_eval_dashboard(self, output_log: str) -> dict[str, Any]:
        """
        Parse the Machine Evals Dashboard URL from the eval output log.

        Args:
            output_log: The output log from the eval command

        Returns:
            Dict containing the parsed dashboard URL
        """
        # Try to find Machine Evals Dashboard URL pattern in the log
        # Example: Machine Evals Dashboard URL: https://...
        dashboard_pattern = r"Machine Evals Dashboard URL:\s*(https://[^\s]+)"

        dashboard_match = re.search(dashboard_pattern, output_log, re.IGNORECASE)
        if dashboard_match:
            dashboard_url = dashboard_match.group(1)
            return {"dashboard_url": dashboard_url, "success": True}

        return {"success": False, "error": "Machine Evals Dashboard URL not found in output log"}

    async def submit_aipa_eval(self, payload_json: str) -> dict[str, Any]:
        """
        Submit an Aipa-Eval job based on the provided payload.
        Args:
            payload_json: JSON string containing the eval parameters, must include 'branch', 'run_id', 'commit_id' and 'dataset_name'.
        Returns:
            Dict containing the experiment_id or error message
        """
        if self.local_debug:
            suffix = random.randint(1000, 9999)
            return {"experiment_id": f"cad3d72a-4813-4a23-8a57-775e63730a4c-{suffix}","experiment_link": "https://www.braintrust.dev/app/Microsoft%20Copilot%20PME/p/Edge_CN_CopilotContextualQuality"}
        #
        payload = json.loads(payload_json)
        if not isinstance(payload, dict):
            return {"error": "payload_json must decode to an object"}
        prompt_branch = payload.get("prompt_branch") or payload.get("branch")
        if prompt_branch is None:
            return {"error": "prompt_branch is required in payload for picasso eval."}
        if payload.get("run_id") is None:
            return {"error": "run_id is required in payload for picasso eval."}
        if payload.get("commit_id") is None:
            return {"error": "commit_id is required in payload for picasso eval."}
        if payload.get("dataset_name") is not None:
            self.dataset_name = payload.get("dataset_name")
        
        run_id = payload.get("run_id")        
        commit_id_short = payload.get("commit_id")[:7]
        task_id = (self.task_id_prefix + '_' if self.task_id_prefix else "auto_prompt_lab_test_task_") + run_id + '_' + commit_id_short
        experiment_suffix = "auto_prompt_lab_experiment_" + run_id + '_' + commit_id_short
        
        eval_result = await self.submit_eval(prompt_branch, task_id, experiment_suffix)

        if "output_log" in eval_result:
            parse_result = await self.parse_experiment_id(eval_result["output_log"])
            if parse_result.get("success"):
                return {"experiment_id": parse_result["result_id"], "experiment_link": parse_result["link"]}
            else:
                return {"error": "Failed to parse experiment_id from eval output.", "eval_result": eval_result}
        else:
            return {"error": "Eval submission failed.", "eval_result": eval_result}

    async def submit_job(self, payload: dict[str, Any]) -> str:
        """Submit a job to the backend."""
        job_id = str(payload.get("job_id") or "stub-job")
        if self.backend is not None:
            self.backend.jobs[job_id] = dict(payload)
        return job_id

    async def cancel_all_runs_on_branch(self, branch: str) -> dict[str, Any]:
        """
        Cancel all in-progress and queued GitHub Actions workflow runs on a specific branch.
        
        Args:
            branch: The branch name to cancel all runs for
            
        Returns:
            Dict containing success status and details of cancelled runs
        """
        cancelled_runs = []
        errors = []
        repo = self.AIPA_EVAL_REPO
        
        for status in ["in_progress", "queued"]:
            try:
                # Get all run IDs for this branch and status
                list_cmd = [
                    "gh", "run", "list",
                    "--repo", repo,
                    "--branch", branch,
                    "--status", status,
                    "--json", "databaseId",
                    "--jq", ".[].databaseId"
                ]
                
                list_result = await run_subprocess_async(
                    list_cmd,
                    timeout=60
                )
                
                if list_result.returncode != 0:
                    errors.append(f"Failed to list {status} runs: {list_result.stderr}")
                    continue
                
                # Parse run IDs from output
                run_ids = [rid.strip() for rid in list_result.stdout.strip().split('\n') if rid.strip()]
                
                # Cancel each run
                for run_id in run_ids:
                    try:
                        cancel_cmd = ["gh", "run", "cancel", run_id, "--repo", repo]
                        cancel_result = await run_subprocess_async(
                            cancel_cmd,
                            timeout=30
                        )
                        
                        if cancel_result.returncode == 0:
                            cancelled_runs.append(run_id)
                            logger.info("Cancelled run %s", run_id)
                        else:
                            errors.append(f"Failed to cancel run {run_id}: {cancel_result.stderr}")
                    except Exception as e:
                        errors.append(f"Error cancelling run {run_id}: {str(e)}")
                        
            except Exception as e:
                errors.append(f"Error listing {status} runs: {str(e)}")
        
        return {
            "success": len(errors) == 0,
            "cancelled_runs": cancelled_runs,
            "cancelled_count": len(cancelled_runs),
            "errors": errors if errors else None
        }

class BraintrustClient:
    def __init__(
        self,
        backend,
        api_key: str,
        metrics_config: list[dict] = None,  # List of metric configurations (primary)
        dimension_threshold: float = 0.5,
        bad_case_limit_per_iteration: int = 10,
        output_dir: str = r"dataset/braintrust_output",
        local_debug: bool = False,
        model_client: Any | None = None,
        analysis_model_client: Any | None = None,
        analysis_agent_client: Any | None = None,
        badcase_deep_analysis: bool = False,
        eval_name: str = "default",
        braintrust_org: str = "non-pme",
        environment: str | None = None,
        # Model config for BTQL agent subprocess
        model_client_type: str = "azure",
        model_endpoint: str | None = None,
        model_name: str | None = None,
        # Optional callback for logging to UI
        log_callback: Any | None = None,
    ) -> None:
        self.api_key = api_key
        self.backend = backend
        self.output_dir = output_dir

        # Store metrics configuration (primary source of truth)
        self.metrics_config = metrics_config or []
        # Build a quick lookup dictionary: metric_name -> config
        self.metrics_config_dict = {}
        for mc in self.metrics_config:
            if isinstance(mc, dict) and 'name' in mc:
                self.metrics_config_dict[mc['name']] = mc

        # Derive key_metrics from metrics_config (all metrics)
        cleaned_key_metrics = None
        if self.metrics_config:
            # Extract names from metrics_config
            cleaned_key_metrics = [
                mc['name'] for mc in self.metrics_config
                if isinstance(mc, dict) and 'name' in mc and isinstance(mc['name'], str) and mc['name'].strip()
            ]
            if not cleaned_key_metrics:
                cleaned_key_metrics = None

        self.key_metrics = cleaned_key_metrics

        # Derive llm_based_key_metrics from metrics_config (only LLM-judged metrics for btql-agent)
        llm_based_key_metrics = None
        if self.metrics_config:
            # Extract names from metrics_config where is_llm_judged is True
            llm_based_key_metrics = [
                mc['name'] for mc in self.metrics_config
                if isinstance(mc, dict) and 'name' in mc and isinstance(mc['name'], str) and mc['name'].strip()
                and mc.get('is_llm_judged') is True
            ]
            if not llm_based_key_metrics:
                llm_based_key_metrics = None

        self.llm_based_key_metrics = llm_based_key_metrics

        self.dimension_threshold = dimension_threshold
        self.bad_case_limit_per_iteration = bad_case_limit_per_iteration
        self.local_debug = local_debug
        # analysis_model_client takes priority for root-cause analysis;
        # falls back to model_client when not provided.
        self.model_client = analysis_model_client if analysis_model_client is not None else model_client
        self.analysis_agent_client = analysis_agent_client
        if self.analysis_agent_client is not None and hasattr(self.analysis_agent_client, "set_metrics_config"):
            try:
                self.analysis_agent_client.set_metrics_config(self.metrics_config)
            except Exception as exc:
                logger.warning("Failed to sync metrics_config into analysis_agent_client: %s", exc)
        self.badcase_deep_analysis = badcase_deep_analysis
        self.eval_name = eval_name
        self.braintrust_org = braintrust_org
        self.environment = (environment or "global").strip().lower()
        self.braintrust_base_url = BRAINTRUST_ORG_ENDPOINTS.get(braintrust_org, BRAINTRUST_ORG_ENDPOINTS["non-pme"])
        # Model config for BTQL agent
        self.model_client_type = _normalize_btql_agent_model_type(model_client_type)
        self.model_endpoint = model_endpoint
        self.model_name = model_name
        # Cache experiment-level tool registries so we parse once per experiment.
        self._experiment_tool_registry_cache: dict[str, dict[str, Any]] = {}
        # Callback for logging to UI (async function: log_callback(event_type, agent_name, message, payload))
        self._log_callback = log_callback
        self.__create_output_dir__()
    
    def __create_output_dir__(self) -> None:
        """Create output directory if it does not exist."""
        path = Path(self.output_dir)
        path.mkdir(parents=True, exist_ok=True)
    
    async def __save_output_data__(self, data: Any, path, format: str = "jsonl") -> None:
        """Save output data to file.

        Args:
            data: Content to persist.
            path: Relative output path under ``self.output_dir``.
            format: ``"jsonl"`` (default) or ``"text"``.
        """
        # check and create parent dir
        
        # Convert to absolute path to ensure consistency across async operations
        abs_output_dir = os.path.abspath(self.output_dir)
        path = os.path.normpath(os.path.join(abs_output_dir, path))
        
        # On Windows, use \\?\ prefix for long path support (>260 chars)
        if os.name == 'nt' and len(path) > 260 and not path.startswith('\\\\?\\'):
            path = '\\\\?\\' + path
        
        parent_dir = os.path.dirname(path)
        logger.info(f"Creating directory if not exists: {parent_dir}")
        try:
            os.makedirs(parent_dir, exist_ok=True)
            logger.info(f"Directory created/verified: {parent_dir}")
        except Exception as e:
            logger.error(f"Failed to create directory {parent_dir}: {e}")
            raise
        
        logger.info(f"Writing data to file: {path}")
        try:
            async with aiofiles.open(path, "w", encoding="utf-8") as f:
                if format == "text":
                    if isinstance(data, str):
                        await f.write(data)
                    else:
                        await f.write(json.dumps(data, ensure_ascii=False))
                else:
                    records: list[Any]
                    if isinstance(data, list):
                        records = data
                    elif isinstance(data, dict):
                        records = [data]
                    else:
                        raise TypeError("jsonl format requires list/dict data")
                    for record in records:
                        await f.write(json.dumps(record, ensure_ascii=False) + "\n")
            logger.info(f"Successfully wrote data to: {path}")
        except Exception as e:
            logger.error(f"Failed to write file {path}: {type(e).__name__}: {e}")
            raise
    
    async def __read_file__(self, file_path: str) -> str:
        """Read file content asynchronously."""
        async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
            content = await f.read()
        return content

    async def fetch_summary_results(self, experiment_id: str, if_fetch_details: bool = False) -> dict[str, Any]:
        logger.info(
            "[fetch_summary_results] start experiment_id=%s, if_fetch_details=%s, local_debug=%s",
            experiment_id,
            if_fetch_details,
            self.local_debug,
        )
        if self.local_debug:
            if if_fetch_details:
                logger.info("[fetch_summary_results] local_debug enabled, fetching event details for experiment_id=%s", experiment_id)
                local_debug_details = await self.fetch_event_details(experiment_id)
                if isinstance(local_debug_details, list):
                    logger.info(
                        "[fetch_summary_results] local_debug event details fetched, count=%s",
                        len(local_debug_details),
                    )
                else:
                    logger.warning(
                        "[fetch_summary_results] local_debug event details returned non-list: type=%s, value=%s",
                        type(local_debug_details).__name__,
                        local_debug_details,
                    )

            return {
                "completeness": random.uniform(0, 1),
                "clarity & readability": random.uniform(0, 1),
                "engagement": random.uniform(0, 1),
                "factual accuracy": random.uniform(0, 1),
                "groundedness": random.uniform(0, 1)
            }
        # """Fetch summarized evaluation results from Braintrust API."""
        response = await self.send_request_to_btql(
            path=f"/v1/experiment/{experiment_id}/summarize",
            method="GET",
            params={"summarize_scores": "true"},
            timeout=httpx.Timeout(60.0, connect=10.0),
        )
        scores = None
        logger.info(f"Braintrust summarize response status: {response.status_code}, text: {response.text}")
        if response.status_code != 200:
            logger.warning(
                "[fetch_summary_results] summarize API returned non-200 for experiment_id=%s, status=%s",
                experiment_id,
                response.status_code,
            )
        # # keep related key_metrics only
        if response.status_code == 200:
            try:
                respond_json = response.json()
            except json.JSONDecodeError:
                logger.exception(
                    "[fetch_summary_results] summarize response JSON decode failed for experiment_id=%s",
                    experiment_id,
                )
                respond_json = {}
            if 'scores' in respond_json:
                scores = respond_json['scores']
                if not isinstance(scores, dict):
                    logger.warning(
                        "[fetch_summary_results] summarize scores is not dict: type=%s, value=%s",
                        type(scores).__name__,
                        scores,
                    )
                    scores = None
                else:
                    logger.info(
                        "[fetch_summary_results] summarize scores loaded, metric_count=%s, metrics=%s",
                        len(scores),
                        list(scores.keys()),
                    )
                # If key_metrics is None, keep all scores; otherwise filter by key_metrics
                if self.key_metrics is not None:
                    original_metric_count = len(scores)
                    scores = {key: scores.get(key) for key in self.key_metrics if key in scores}
                    scores = {k: v for k, v in scores.items() if k in self.key_metrics}
                    logger.info(
                        "[fetch_summary_results] key_metrics filter applied, before=%s, after=%s, key_metrics=%s",
                        original_metric_count,
                        len(scores),
                        self.key_metrics,
                    )
                else:
                    # Initialize key_metrics with all score keys if not specified
                    self.key_metrics = list(scores.keys())
                    logger.info(
                        "[fetch_summary_results] key_metrics initialized from summary scores, count=%s",
                        len(self.key_metrics),
                    )
                    # intialize metrics_config_dict if empty
                    for metric_name in self.key_metrics:
                        if metric_name not in self.metrics_config_dict:
                            self.metrics_config_dict[metric_name] = {"name": metric_name, "optimization_direction": "higher_better"}

                # Normalize scores based on optimization_direction (lower_better -> invert)
                inverted_metric_count = 0
                skipped_metric_count = 0
                for metric_name, score_value in list(scores.items()):
                    # Handle both nested dict format {'metric': {'score': 0.5}} and flat format {'metric': 0.5}
                    if isinstance(score_value, dict) and 'score' in score_value:
                        actual_score = score_value['score']
                    elif isinstance(score_value, (int, float)):
                        actual_score = score_value
                    else:
                        skipped_metric_count += 1
                        continue

                    metric_config = self.metrics_config_dict.get(metric_name, {})
                    optimization_direction = metric_config.get('optimization_direction', 'higher_better')

                    # If lower is better, invert the score to make it higher_better
                    if optimization_direction == 'lower_better':
                        inverted_score = 1.0 - float(actual_score)
                        if isinstance(score_value, dict):
                            scores[metric_name]['score'] = inverted_score
                        else:
                            scores[metric_name] = inverted_score
                        inverted_metric_count += 1
                logger.info(
                    "[fetch_summary_results] summarize score normalization done, inverted=%s, skipped=%s",
                    inverted_metric_count,
                    skipped_metric_count,
                )
            else:
                logger.warning(
                    "[fetch_summary_results] summarize response missing 'scores' for experiment_id=%s, keys=%s",
                    experiment_id,
                    list(respond_json.keys()) if isinstance(respond_json, dict) else type(respond_json).__name__,
                )
        # return {"error": response.text}

        # scores = {"query": f"""from: experiment('{experiment_id}')| dimensions: root_span_id, scores| filter: scores IS NOT NULL and input.output.picasso_eval_metadata.inferences[-1].response != ""  and input.output.picasso_eval_metadata.responding.request.messages[0].content.text NOT LIKE '%pageMetadata":"ContentFiltered"%' and input.output[0].picasso_eval_metadata.safetyEvents[-1].result not contains "blocked" """}
        # async with httpx.AsyncClient() as client:
        #     # scores_resp = await client.post(url="https://mangoeastus2prodcpu-d7amaactfmcrazd5.b01.azurefd.net/btql", json=scores, headers=headers)
        #     scores_resp = await client.post(url="https://api.braintrust.dev/btql", json=scores, headers=headers)
        
        #     # Convert scores to dict, merge same root_span_id scores
        #     if scores_resp.status_code != 200:
        #         return {"error": scores_resp.text}
        
        # #aggregate scores for each dimension
        # scores_rows = scores_resp.json().get('data', [])
        logger.info("[fetch_summary_results] start fetch_event_details for experiment_id=%s", experiment_id)
        try:
            filtered_output = await self.fetch_event_details(experiment_id)
        except Exception:
            logger.exception(
                "[fetch_summary_results] fetch_event_details failed for experiment_id=%s",
                experiment_id,
            )
            raise

        if isinstance(filtered_output, list):
            logger.info(
                "[fetch_summary_results] fetch_event_details completed, returned list count=%s",
                len(filtered_output),
            )
        else:
            logger.warning(
                "[fetch_summary_results] fetch_event_details returned non-list: type=%s, value=%s",
                type(filtered_output).__name__,
                filtered_output,
            )

        experiment_output_dir = Path(self.output_dir) / f"experiments_{experiment_id}"
        event_details_path = experiment_output_dir / "event_details.jsonl"
        event_details_filtered_path = experiment_output_dir / "event_details_filtered.jsonl"
        for expected_file in (event_details_path, event_details_filtered_path):
            if expected_file.exists():
                logger.info(
                    "[fetch_summary_results] output file exists: %s (size=%s bytes)",
                    expected_file,
                    expected_file.stat().st_size,
                )
            else:
                logger.warning(
                    "[fetch_summary_results] output file missing: %s",
                    expected_file,
                )

        aggregated_scores = scores
        if aggregated_scores is None:
            # Fallback: aggregate scores from event details if API summary failed
            logger.info(
                "[fetch_summary_results] summarize scores unavailable, fallback to aggregate from event details for experiment_id=%s",
                experiment_id,
            )
            aggregated_scores = {}

            # Check if filtered_output is valid (list, not an error dict)
            if isinstance(filtered_output, list):
                for row in filtered_output:
                    score_dict = row.get('merged_scores', {})
                    for key, value in score_dict.items():
                        # If key_metrics is None, include all metrics; otherwise filter
                        if self.key_metrics is None or key in self.key_metrics:
                            if key not in aggregated_scores:
                                aggregated_scores[key] = []
                            aggregated_scores[key].append(value)

                # Compute average scores for each dimension
                for key, values in aggregated_scores.items():
                    if values:  # Avoid division by zero
                        aggregated_scores[key] = sum(values) / len(values)
                logger.info(
                    "[fetch_summary_results] fallback aggregation completed, metric_count=%s, metrics=%s",
                    len(aggregated_scores),
                    list(aggregated_scores.keys()),
                )
            else:
                logger.warning("fetch_event_details returned error: %s", filtered_output)

        logger.info(
            "[fetch_summary_results] done experiment_id=%s, final_metric_count=%s",
            experiment_id,
            len(aggregated_scores) if isinstance(aggregated_scores, dict) else 0,
        )
        return aggregated_scores
    
    async def fetch_eval_prompt(self, experiment_id: str, root_span_ids: list[str]) -> dict[str, Any]:
        """Fetch evaluation prompt from Braintrust API and persist it under experiments_<id>/eval_prompt.jsonl."""
        output_data_path = f"experiments_{experiment_id}/eval_prompt.jsonl"
        if self.local_debug:
            # Resolve path relative to project root (autopromplab), not current working directory
            # protocol.py -> core -> prompt_auto_lab -> autoPromptLab2 -> autopromplab (project root)
            project_root = Path(__file__).resolve().parents[3]
            fixture_path = project_root / "autoPromptLab2" / "eval_prompt.jsonl"
            if not fixture_path.exists():
                return {"success": False, "error": f"eval_prompt fixture not found: {fixture_path}"}

            fixture_text = fixture_path.read_text(encoding="utf-8").strip()
            if not fixture_text:
                return {"success": False, "error": f"eval_prompt fixture is empty: {fixture_path}"}

            try:
                loaded = json.loads(fixture_text)
                data_list = [loaded] if isinstance(loaded, dict) else loaded
            except json.JSONDecodeError:
                data_list = []
                for line in fixture_text.splitlines():
                    line = line.strip()
                    if not line:
                        continue
                    data_list.append(json.loads(line))

            await self.__save_output_data__(data_list, output_data_path)
            return {"success": True, "file path": str(Path(self.output_dir) / output_data_path)}

        def _build_query(*, use_root_span_filter: bool) -> str:
            root_span_filter = ""
            if use_root_span_filter:
                normalized_root_span_ids = [
                    str(span_id).strip()
                    for span_id in root_span_ids
                    if str(span_id).strip()
                ]
                if normalized_root_span_ids:
                    quoted_ids = ", ".join(f"'{span_id}'" for span_id in normalized_root_span_ids)
                    root_span_filter = f" AND root_span_id IN [{quoted_ids}]"

            if environment == "china":
                return f"""from: experiment('{experiment_id}')
| select:
metadata.request.messages[0].content[0].text AS prompt_text,
| filter: metadata.request.messages[0].content[0].text != ""{root_span_filter}
| limit: 1"""

            return f"""from: experiment('{experiment_id}')
| select:
metadata.request AS prompt_text,
| filter: metadata.request IS NOT NULL{root_span_filter}
| limit: 1"""

        def _has_prompt_payload(payload: Any) -> bool:
            if isinstance(payload, list):
                return any(_has_prompt_payload(item) for item in payload)
            if not isinstance(payload, dict):
                return False

            prompt_text = payload.get("prompt_text")
            if isinstance(prompt_text, str):
                return bool(prompt_text.strip())
            if isinstance(prompt_text, (list, dict, tuple, set)):
                return len(prompt_text) > 0
            if prompt_text is not None:
                return True

            data = payload.get("data")
            if isinstance(data, list):
                return any(_has_prompt_payload(item) for item in data)
            if isinstance(data, dict):
                return _has_prompt_payload(data)
            return False

        environment = self.environment or "global"
        filtered_query = _build_query(use_root_span_filter=True)
        unfiltered_query = _build_query(use_root_span_filter=False)

        response = await self.send_request_to_btql(filtered_query, timeout=httpx.Timeout(60.0, connect=10.0))
        response_data = response.json()

        if filtered_query != unfiltered_query and response.status_code == 200 and not _has_prompt_payload(response_data):
            logger.info(
                "[fetch_eval_prompt] no prompt found for experiment_id=%s using root_span_ids filter; retrying without root_span filter",
                experiment_id,
            )
            response = await self.send_request_to_btql(unfiltered_query, timeout=httpx.Timeout(60.0, connect=10.0))
            response_data = response.json()
        
        # Save to file (response.json() returns a dict, wrap it in a list to match original behavior)
        data_list = [response_data] if isinstance(response_data, dict) else response_data
        await self.__save_output_data__(data_list, output_data_path)
        
        if response.status_code == 200:
            return {"success": True, "file path": str(Path(self.output_dir) / output_data_path)}

        return {"success": False, "error": response.text}

    async def send_request_to_btql(
        self,
        query: str | None = None,
        *,
        path: str = "/btql",
        method: str = "POST",
        params: dict[str, Any] | None = None,
        json_body: dict[str, Any] | None = None,
        timeout: httpx.Timeout | None = None,
    ) -> httpx.Response:
        """Send a request to Braintrust API (BTQL by default)."""
        base_url = self.braintrust_base_url.rstrip("/")
        url = f"{base_url}{path}"
        headers = {"Authorization": f"Bearer {self.api_key}"}
        try_times = 3
        if timeout is None:
            timeout = httpx.Timeout(30.0, connect=10.0)
        method_upper = method.upper()
        if json_body is None and query is not None and method_upper != "GET":
            json_body = {"query": query}
        last_response: httpx.Response | None = None

        async with httpx.AsyncClient(timeout=timeout, follow_redirects=False) as client:
            for attempt in range(try_times):
                try:
                    response = await client.request(
                        method_upper,
                        url,
                        headers=headers,
                        params=params,
                        json=json_body,
                    )
                    last_response = response
                except httpx.TimeoutException as exc:
                    wait_s = 1.0 + attempt
                    logger.warning(
                        "Attempt %s, BTQL request timeout: %s. Retrying in %ss",
                        attempt,
                        exc,
                        wait_s,
                    )
                    await asyncio.sleep(wait_s)
                    continue

                if response.status_code == 200:
                    return response
                if response.status_code in {301, 302, 303, 307, 308}:
                    redirect_url = response.headers.get("Location")
                    if not redirect_url:
                        match = re.search(r"(https?://\S+)", response.text or "")
                        redirect_url = match.group(1) if match else None
                    if not redirect_url:
                        return response
                    try:
                        redirect_resp = await client.get(redirect_url)
                        if redirect_resp.status_code != 200:
                            return redirect_resp
                        content = redirect_resp.content
                        if redirect_url.endswith(".json.gz") or redirect_resp.headers.get("Content-Encoding") == "gzip":
                            if content[:2] == b"\x1f\x8b":
                                try:
                                    content = gzip.decompress(content)
                                except OSError:
                                    with gzip.GzipFile(fileobj=io.BytesIO(content)) as gz:
                                        content = gz.read()
                        return httpx.Response(
                            200,
                            content=content,
                            request=httpx.Request(method_upper, url, params=params),
                            headers={"Content-Type": "application/json"},
                        )
                    except httpx.RequestError as exc:
                        logger.warning("%s, Redirect fetch failed: %s", attempt, exc)
                        return response
                if response.status_code == 429:
                    logger.warning(
                        "Attempt %s, Rate limit exceeded when sending request to BTQL API, error: %s",
                        attempt,
                        response.text,
                    )
                    retry_after_match = re.search(r"Retry after ([\d.]+) seconds", response.text)
                    retry_after_s = float(retry_after_match.group(1)) if retry_after_match else (1.0 + attempt)
                    logger.info("Waiting for %.2f seconds before retrying...", retry_after_s)
                    await asyncio.sleep(retry_after_s + 0.5)
                    continue
                    

        return last_response if last_response is not None else httpx.Response(599, request=httpx.Request("POST", url))

    def _should_use_non_llm_fallback_queries(self) -> bool:
        """Return True when all key metrics are explicitly marked non-LLM-judged."""
        if not self.key_metrics:
            return False
        if not self.metrics_config_dict:
            return False

        metric_flags: list[bool] = []
        for metric_name in self.key_metrics:
            metric_config = self.metrics_config_dict.get(metric_name)
            if not isinstance(metric_config, dict):
                return False
            if "is_llm_judged" not in metric_config:
                return False
            metric_flags.append(metric_config.get("is_llm_judged") is True)

        return len(metric_flags) > 0 and not any(metric_flags)

    def _build_non_llm_fallback_queries(self, experiment_id: str) -> dict[str, Any]:
        """Fallback BTQL queries for experiments without LLM-judged metrics."""
        scores_query = (
            f"from: experiment('{experiment_id}') | "
            "select: root_span_id, scores | "
            "dimensions: root_span_id, scores | "
            "filter: scores IS NOT NULL"
        )
        dimensions_query = (
            f"from: experiment('{experiment_id}') | "
            "select: root_span_id, input.input AS input, \"{}\" AS judge_response | "
            "filter: scores IS NOT NULL"
        )
        return {
            "scores": scores_query,
            "dimensions": dimensions_query,
            "dimensions_grouped": [dimensions_query],
        }

    def _build_btql_agent_failure_fallback_queries(self, experiment_id: str) -> dict[str, Any]:
        """Fallback queries when btql-agent generation/execution fails."""
        scores_query = (
            f"from: experiment('{experiment_id}') | "
            "select: root_span_id, scores | "
            "dimensions: root_span_id, scores | "
            "filter: scores IS NOT NULL"
        )
        dimensions_query = (
            f"from: experiment('{experiment_id}') | "
            "select: root_span_id, input.input AS input, coalesce(to_string(metadata), '{}') AS judge_response | "
            "dimensions: root_span_id, input.input AS input, coalesce(to_string(metadata), '{}') AS judge_response | "
            "filter: scores IS NOT NULL AND metadata IS NOT NULL AND to_string(metadata) != '' AND to_string(metadata) != '{}'"
        )
        return {
            "scores": scores_query,
            "dimensions": dimensions_query,
            "dimensions_grouped": [dimensions_query],
        }

    async def generate_btql_queries_with_agent(self, experiment_id: str) -> dict[str, Any]:
        """Run btql-agent to generate BTQL queries dynamically.

        Args:
            experiment_id: The Braintrust experiment ID

        Returns:
            Dictionary with:
            - 'scores': scores query string
            - 'dimensions': legacy single dimensions query string
            - 'dimensions_grouped': list of grouped dimensions query strings
        """
        if self._should_use_non_llm_fallback_queries():
            logger.info(
                "[generate_btql_queries_with_agent] using non-LLM fallback queries, experiment_id=%s key_metrics=%s",
                experiment_id,
                self.key_metrics,
            )
            return self._build_non_llm_fallback_queries(experiment_id)

        # Check if btql-agent is installed, install if necessary
        try:
            check_result = subprocess.run(
                [sys.executable, "-m", "btql_agent", "--help"],
                capture_output=True,
                timeout=20
            )
            # returncode can be 0 even when module is missing (error only in stderr)
            stderr_text = check_result.stderr.decode("utf-8", errors="replace") if check_result.stderr else ""
            if check_result.returncode != 0 or "No module named" in stderr_text:
                raise FileNotFoundError("btql-agent not properly installed")
        except (FileNotFoundError, subprocess.TimeoutExpired):
            logger.warning("btql-agent not found or not working, attempting to install...")
            try:
                # Ensure pip is available in current interpreter
                pip_check = subprocess.run(
                    [sys.executable, "-m", "pip", "--version"],
                    capture_output=True,
                    text=True,
                    timeout=20,
                )
                pip_stderr = pip_check.stderr or ""
                if pip_check.returncode != 0 or "No module named" in pip_stderr:
                    logger.warning("pip not available, running ensurepip...")
                    ensurepip_result = subprocess.run(
                        [sys.executable, "-m", "ensurepip", "--upgrade"],
                        capture_output=True,
                        text=True,
                        timeout=60,
                    )
                    if ensurepip_result.returncode != 0:
                        raise RuntimeError(
                            f"ensurepip failed (rc={ensurepip_result.returncode}): {ensurepip_result.stderr}"
                        )
                    logger.info("ensurepip succeeded")

                # Get the tools/auto-btql directory path
                # protocol.py -> core -> prompt_auto_lab (which contains tools/)
                prompt_auto_lab_root = Path(__file__).resolve().parents[1]
                btql_path = prompt_auto_lab_root / "tools" / "auto-btql"
                
                install_result = subprocess.run(
                    [sys.executable, "-m", "pip", "install", "-e", "."],
                    cwd=str(btql_path),
                    capture_output=True,
                    text=True,
                    timeout=120
                )
                
                install_stderr = install_result.stderr or ""
                if install_result.returncode != 0 or "No module named" in install_stderr:
                    logger.error("Failed to install btql-agent (rc=%s): %s", install_result.returncode, install_stderr)
                    raise RuntimeError("Could not install btql-agent")
                
                logger.info("btql-agent installed successfully")
            except Exception as e:
                logger.error("Error installing btql-agent: %s", e)
                raise RuntimeError(f"Failed to install btql-agent: {e}")

        # Run btql-agent command using python -m to ensure it uses the current environment
        # Point --cache-dir to prompt_auto_lab/config where schema_cache.json lives
        prompt_auto_lab_root = Path(__file__).resolve().parents[1]
        cache_dir = str(prompt_auto_lab_root / "config")

        cmd = [
            sys.executable,
            "-m",
            "btql_agent",
            "-e", experiment_id,
            "--eval", self.eval_name,
            "--braintrust-api-key", self.api_key,
            "--braintrust-org", self.braintrust_org,
            "--template-format", "json",
            "--model-type", _normalize_btql_agent_model_type(self.model_client_type),
            "--cache-dir", cache_dir,
        ]

        # Add model endpoint for copilot/local
        if self.model_endpoint:
            cmd.extend(["--model-endpoint", self.model_endpoint])

        # Add model name for copilot/local (or Azure deployment override)
        if self.model_name:
            cmd.extend(["--model-name", self.model_name])
        
        # Add metrics if llm_based_key_metrics is provided (only LLM-judged metrics)
        if self.llm_based_key_metrics:
            metrics_str = ",".join(self.llm_based_key_metrics)
            cmd.extend(["-m", metrics_str])

        try:
            # Run the command and capture output
            # Set COLUMNS to prevent Rich Console from wrapping output
            env = os.environ.copy()
            env['COLUMNS'] = '10000'  # Large value to prevent line wrapping

            # Use run_subprocess_async to avoid blocking the event loop.
            # This is critical when using Copilot model client - the subprocess
            # makes HTTP requests to the backend's /api/copilot/v1 endpoint,
            # which requires the event loop to process polling requests from
            # the VS Code extension. Using sync subprocess.run() would deadlock.
            result = await run_subprocess_async(
                cmd,
                capture_output=True,
                text=True,
                encoding='utf-8',
                timeout=300,  # 5 minute timeout
                env=env
            )
            actual_result = result.returncode
            output = result.stdout

            # Log btql-agent output to UI if callback is available
            if self._log_callback is not None:
                try:
                    # Build log message with clear sections
                    log_lines = [f"experiment_id={experiment_id}"]
                    if result.stderr:
                        log_lines.append(f"--- stderr ---\n{result.stderr.strip()}")
                    log_lines.append(f"--- stdout ---\n{output.strip()}")
                    log_message = "\n".join(log_lines)

                    # Use agent_message with "btql-agent" so it's clearly identified
                    await self._log_callback(
                        "agent_message",
                        "btql-agent",
                        log_message,
                    )
                except Exception as e:
                    logger.warning("Failed to broadcast btql-agent output: %s", e)

            # Parse JSON output to extract queries
            scores_query = None
            dimensions_query = None
            grouped_dimensions_queries: list[str] = []

            try:
                # The btql-agent may output human-readable text before the JSON
                # Find the JSON portion by looking for the first '{' that starts a JSON object
                json_start = output.find('{"experiment_id"')
                if json_start == -1:
                    # Fallback: try to find any JSON object
                    json_start = output.rfind('{')
                
                if json_start != -1:
                    json_str = output[json_start:]
                    # Find the matching closing brace
                    brace_count = 0
                    json_end = -1
                    for i, char in enumerate(json_str):
                        if char == '{':
                            brace_count += 1
                        elif char == '}':
                            brace_count -= 1
                            if brace_count == 0:
                                json_end = i + 1
                                break
                    if json_end != -1:
                        json_str = json_str[:json_end]
                    json_output = json.loads(json_str)
                else:
                    # Try parsing the entire output as JSON
                    json_output = json.loads(output)

                # Extract queries from the JSON structure
                queries = json_output.get("queries", {})

                # Get the "get_scores" query
                get_scores = queries.get("get_scores", {})
                scores_query = get_scores.get("query")

                # Get the "get_judge_responses" query
                get_judge_responses = queries.get("get_judge_responses", {})
                dimensions_query = get_judge_responses.get("query")

                # Get grouped judge response queries when available
                grouped = queries.get("get_judge_responses_grouped", [])
                if isinstance(grouped, list):
                    for item in grouped:
                        if not isinstance(item, dict):
                            continue
                        q = item.get("get_judge_responses")
                        if isinstance(q, str) and q.strip():
                            grouped_dimensions_queries.append(q.strip())

            except json.JSONDecodeError as e:
                cmd_str = " ".join(cmd)
                logger.error("Failed to parse JSON output from btql-agent.")
                logger.error("Command: %s", cmd_str)
                logger.error("Output: %s", output)
                logger.error("Stderr: %s", result.stderr)
                raise RuntimeError(f"Failed to parse JSON output from btql-agent for experiment {experiment_id}: {e}\nCommand: {cmd_str}\nOutput: {output}\nStderr: {result.stderr}")

            # Validate that we successfully parsed both queries
            if not scores_query or not dimensions_query:
                logger.error("Failed to extract scores query or dimensions_query from JSON. btql-agent output:")
                logger.error(output)
                raise RuntimeError(f"Failed to extract scores query or dimensions query from btql-agent JSON output for experiment {experiment_id}")

            # Deduplicate while preserving order.
            if grouped_dimensions_queries:
                grouped_dimensions_queries = list(dict.fromkeys(grouped_dimensions_queries))
            else:
                grouped_dimensions_queries = [dimensions_query]

            logger.info(
                "[generate_btql_queries_with_agent] experiment_id=%s grouped_dimension_queries=%s",
                experiment_id,
                len(grouped_dimensions_queries),
            )

            return {
                "scores": scores_query,
                "dimensions": dimensions_query,
                "dimensions_grouped": grouped_dimensions_queries,
            }
            
        except subprocess.TimeoutExpired as e:
            logger.error("Output: %s", output)
            logger.error("Stderr: %s", result.stderr)
            raise RuntimeError(f"btql-agent command timed out after 120 seconds for experiment {experiment_id}") from e
        except FileNotFoundError as e:
            raise RuntimeError(f"btql-agent command not found. Make sure it's installed and in PATH.") from e
    
    #change input to id by hashing input content
    async def change_input_to_id(self, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
        for row in rows:
            input_content = row.get('input', '')
            if input_content:
                input_content = json.dumps(input_content, ensure_ascii=False)
                hashed_id = hashlib.md5(input_content.encode('utf-8')).hexdigest()
                row['input'] = hashed_id 
        return rows

    async def get_scores_for_root_span_id(self, root_span_ids: list[str], query_template: str) -> dict[str, Any]:
        """Requery scores for a specific root_span_id."""
        query = query_template.replace("{root_span_id}", ', '.join(f"'{id}'" for id in root_span_ids))
        response = await self.send_request_to_btql(query)
        if response.status_code == 200:
            response = response.json()
            data = response.get('data', [])
            if data:
               if_success = True
               return data
        # {"Code":"TooManyRequestsError","Message":"Too many requests. Source: checkBtqlOrgRateLimit. Rate limit: 20 requests per 60 seconds. Consumed: 80. Retry after 5.47 seconds. Please contact us at support@braintrust.dev to discuss remediation strategies. [user_email=yuqiansun@microsoft.com] [user_org=Microsoft Copilot Azure] [timestamp=1767938116.384]","InternalTraceId":"69609844000000006cefed3a153e7b31","Path":"/btql","Service":"api"}
    
        else:
            logger.warning(
                "Failed to requery scores for root_span_id: %s, error: %s",
                root_span_ids,
                response.text,
            )
            return  None
    
    async def get_miss_merged_scores_data(self, merged_data: dict[str, Any]) -> dict[str, Any]:
        missing_merged_scores_data ={}
        for root_span_id, data in merged_data.items():
            if 'merged_scores' not in data.keys() or not data['merged_scores']:
                missing_merged_scores_data[root_span_id] = data
            # if not all key_metrics are in merged_scores
            elif self.key_metrics and not all(key in data['merged_scores'] for key in self.key_metrics):
                missing_merged_scores_data[root_span_id] = data
        return missing_merged_scores_data
            
    async def update_merged_scores_with_requery(self, merged_data: dict[str, Any], scores_rows: list[dict[str, Any]], query_template: str, score_dimension_count: int) -> dict[str, Any]:

        semaphore = asyncio.Semaphore(3)  # limit concurrent requests to 3

        async def requery_scores_for_span(idx: int, root_span_ids: list[str]) -> tuple[str, list[dict[str, Any]] | None]:
            async with semaphore:
                scores = await self.get_scores_for_root_span_id(root_span_ids, query_template)
                return idx, scores
        # try moretimes to requery missing merged_scores, max 5 retries to avoid infinite loop
        max_retries = 5
        retry_count = 0
        while retry_count < max_retries:
            missing_merged_scores_data = await self.get_miss_merged_scores_data(merged_data)
            if len(missing_merged_scores_data) == 0:
                logger.info(
                    "***update_merged_scores_with_requery: All merged_scores are complete after requery***"
                )
                break

            # Print missing root_span_ids for debugging
            missing_keys = list(missing_merged_scores_data.keys())
            logger.info(
                "***update_merged_scores_with_requery: Retry %s/%s, missing root_span_ids: %s***",
                retry_count + 1,
                max_retries,
                missing_keys,
            )

            #group 5 root_span_id in missing_merged_scores_data in each batch 
            root_span_ids = list(missing_merged_scores_data.keys())
            root_span_ids_batch = [root_span_ids[i:i + 5] for i in range(0, len(root_span_ids), 5)]

            tasks = [asyncio.create_task(requery_scores_for_span(idx, batch)) for idx, batch in enumerate(root_span_ids_batch)]
            retry_scores: list[dict[str, Any]] = []
            for idx, requery_scores in await asyncio.gather(*tasks):
                if requery_scores:
                    retry_scores.extend(requery_scores)
                else:
                    logger.warning(
                        "Requery scores failed for root_span_id: %s",
                        root_span_ids[idx],
                    )

            # merge requery scores into merged_data
            await self.merge_scores(list(missing_merged_scores_data.values()), retry_scores)

            #update merged_data with missing_merged_scores_data
            for root_span_id, data in missing_merged_scores_data.items():
                merged_data[root_span_id]['merged_scores'] = data.get('merged_scores', {})

            retry_count += 1

        # Check if there are still missing scores after max retries
        if retry_count >= max_retries:
            remaining_missing_data = await self.get_miss_merged_scores_data(merged_data)
            if len(remaining_missing_data) > 0:
                remaining_missing = list(remaining_missing_data.keys())
                logger.warning(
                    "***update_merged_scores_with_requery: Max retries (%s) reached, exiting with %s missing scores: %s***",
                    max_retries,
                    len(remaining_missing),
                    remaining_missing,
                )

        return merged_data

    def _build_model_response_query(self, experiment_id: str) -> str:
        # Pull response candidates from span-level traces.
        # We return both numeric and string start timestamp and do final latest-row
        # selection client-side by max span_start only.
        return (
            f"from: experiment('{experiment_id}') spans | "
            "select: root_span_id, "
            "metrics.start AS span_start, to_string(metrics.start) AS span_start_raw, "
            "metadata.response AS test_model_response | "
            "dimensions: root_span_id, "
            "metrics.start AS span_start, to_string(metrics.start) AS span_start_raw, "
            "metadata.response AS test_model_response | "
            "filter: metrics.\"Time to first token (inference-level, sum)\" IS NOT NULL "
            "AND metadata.response IS NOT NULL AND metadata.response != '' | "
            "sort: span_start desc"
        )

    @staticmethod
    def _truncate_tail(text: str, max_chars: int) -> str:
        """Truncate *text* by keeping only the leading *max_chars* characters.

        Unlike ``_truncate_text`` which keeps head+tail with ``...`` in the
        middle, this always cuts from the end so the remaining prefix stays
        coherent — important for JSON / structured judge responses.
        """
        if not isinstance(text, str) or max_chars <= 0:
            return text
        if len(text) <= max_chars:
            return text
        suffix = "..."
        if max_chars <= len(suffix):
            return text[:max_chars]
        return text[: max_chars - len(suffix)] + suffix

    @staticmethod
    def _pre_truncate_judge_candidates(candidates: list[str], max_total_chars: int) -> list[str]:
        """Proportionally truncate each candidate so the joined result fits within *max_total_chars*.

        Each candidate is allocated a budget proportional to its share of the
        total length.  The separator is accounted for.
        Truncation is tail-only (keeps prefix) to preserve JSON readability.
        """
        if not candidates:
            return candidates
        separator = ", "
        separator_overhead = len(separator) * max(len(candidates) - 1, 0)
        available = max(max_total_chars - separator_overhead, 0)
        total_len = sum(len(c) for c in candidates)
        if total_len <= available:
            return candidates  # already fits

        truncated: list[str] = []
        for c in candidates:
            if total_len > 0:
                budget = max(int(available * len(c) / total_len), 1)
            else:
                budget = available // len(candidates) if candidates else available
            truncated.append(BraintrustClient._truncate_tail(c, budget))
        return truncated

    def _truncate_judge_response_for_analysis(self, judge_response: Any) -> str:
        """Shrink judge_response payload for bad-case analysis data."""
        max_chars = int(os.getenv("MAX_JUDGE_RESPONSE_ANALYSIS_CHARS", "2000"))
        keywords = ("rationale", "feedback", "explanation", "reason")

        def collect_text_by_keywords(value: Any, path: str = "", out: list[str] | None = None) -> list[str]:
            if out is None:
                out = []
            if isinstance(value, dict):
                for k, v in value.items():
                    next_path = f"{path}.{k}" if path else str(k)
                    collect_text_by_keywords(v, next_path, out)
                return out
            if isinstance(value, list):
                for i, item in enumerate(value):
                    collect_text_by_keywords(item, f"{path}[{i}]", out)
                return out
            if isinstance(value, str):
                path_l = path.lower()
                if any(k in path_l for k in keywords):
                    txt = value.strip()
                    if txt:
                        out.append(f"{path}: {txt}")
            return out

        raw_obj: Any = judge_response
        if isinstance(judge_response, str):
            parsed = self._safe_json_loads(judge_response)
            if parsed is not None:
                raw_obj = parsed

        candidate = ""
        if isinstance(raw_obj, (dict, list)):
            snippets = collect_text_by_keywords(raw_obj)
            if snippets:
                candidate = "\n".join(snippets)
            else:
                candidate = json.dumps(raw_obj, ensure_ascii=False)
        elif raw_obj is None:
            candidate = ""
        else:
            candidate = str(raw_obj)

        truncated, _, _ = _truncate_text(candidate, max_chars=max_chars)
        return truncated if isinstance(truncated, str) else str(truncated)

    def _apply_judge_response_truncation(self, merged_data: dict[str, Any]) -> None:
        """Apply judge_response compaction to merged rows before writing bad-case files."""
        updated = 0
        for row in merged_data.values():
            if not isinstance(row, dict):
                continue
            raw = row.get("judge_response", "")
            compact = self._truncate_judge_response_for_analysis(raw)
            if compact != raw:
                updated += 1
            row["judge_response"] = compact
        logger.info(
            "[judge_response_truncation] total_rows=%s updated_rows=%s",
            len(merged_data),
            updated,
        )

    def _to_decimal(self, value: Any, default: Decimal) -> Decimal:
        try:
            if value is None:
                return default
            if isinstance(value, Decimal):
                return value
            return Decimal(str(value))
        except (TypeError, ValueError, InvalidOperation):
            return default

    def _pick_latest_model_responses_by_root_span_id(
        self,
        response_rows: list[dict[str, Any]],
    ) -> dict[str, dict[str, Any]]:
        latest: dict[str, dict[str, Any]] = {}
        skipped_invalid_row = 0
        skipped_invalid_root = 0
        skipped_empty_response = 0
        equal_start_conflict = 0

        def rank_key(row: dict[str, Any]) -> Decimal:
            return self._to_decimal(
                row.get("span_start_raw", row.get("span_start")),
                Decimal("-1"),
            )

        for row in response_rows:
            if not isinstance(row, dict):
                skipped_invalid_row += 1
                continue
            root_span_id = row.get("root_span_id")
            if not isinstance(root_span_id, str) or not root_span_id:
                skipped_invalid_root += 1
                continue
            response_text = row.get("test_model_response")
            if response_text is None or response_text == "":
                skipped_empty_response += 1
                continue

            current = latest.get(root_span_id)
            if current is None:
                latest[root_span_id] = row
                continue

            row_key = rank_key(row)
            current_key = rank_key(current)
            if row_key > current_key:
                latest[root_span_id] = row
            elif row_key == current_key:
                equal_start_conflict += 1
        logger.info(
            "[model_response_picker] input_rows=%s selected_root_spans=%s skipped_invalid_row=%s skipped_invalid_root=%s skipped_empty_response=%s equal_start_conflict=%s",
            len(response_rows),
            len(latest),
            skipped_invalid_row,
            skipped_invalid_root,
            skipped_empty_response,
            equal_start_conflict,
        )
        if latest:
            sample_starts: list[str] = []
            for selected in list(latest.values())[:3]:
                sample_starts.append(
                    str(selected.get("span_start_raw", selected.get("span_start", "")))
                )
            logger.info(
                "[model_response_picker] sample_selected_span_start=%s",
                sample_starts,
            )
        return latest
    
        
    async def fetch_event_details(self, experiment_id: str) -> list[dict[str, Any]]:
        """Fetch evaluation event details from Braintrust API. including dimensions score and reasoning."""

        if self.local_debug:
            # Resolve path relative to project root (autopromplab), not current working directory
            # protocol.py -> core -> prompt_auto_lab -> autoPromptLab2 -> autopromplab (project root)
            project_root = Path(__file__).resolve().parents[3]
            file_path = project_root / "autoPromptLab2" / "debug" / "event_details_filtered.jsonl"
            with open(file_path, "r", encoding="utf-8") as f:
                data = [json.loads(line) for line in f if line.strip()]
            #write to output dir
            self.output_data_path = f"experiments_{experiment_id}/event_details_filtered.jsonl"
            await self.__save_output_data__(data, self.output_data_path)

            self.output_data_path = f"experiments_{experiment_id}/event_details.jsonl"
            await self.__save_output_data__(data, self.output_data_path)

            return data
        
        using_agent_queries = True
        try:
            generated_queries = await self.generate_btql_queries_with_agent(experiment_id)
            # raise()
        except Exception as exc:  # noqa: BLE001
            logger.warning(
                "[fetch_event_details] btql-agent generation failed, using metadata fallback queries. experiment_id=%s error=%s",
                experiment_id,
                exc,
            )
            generated_queries = self._build_btql_agent_failure_fallback_queries(experiment_id)
            using_agent_queries = False
        
        # Add additional filters to both queries
        # additional_filter = ' and input.output.picasso_eval_metadata.inferences[-1].response != ""  and input.output.picasso_eval_metadata.responding.request.messages[0].content.text NOT LIKE \'%pageMetadata":"ContentFiltered"%\' and input.output[0].picasso_eval_metadata.safetyEvents[-1].result not contains "blocked"'
        grouped_dimension_queries = generated_queries.get("dimensions_grouped")
        if isinstance(grouped_dimension_queries, list):
            dimension_query_list = [q for q in grouped_dimension_queries if isinstance(q, str) and q.strip()]
        else:
            dimension_query_list = []
        if not dimension_query_list:
            base_dimension_query = generated_queries.get("dimensions")
            if isinstance(base_dimension_query, str) and base_dimension_query.strip():
                dimension_query_list = [base_dimension_query]

        dimensions = {"query": generated_queries["dimensions"]}
        scores = {"query": generated_queries["scores"]}
        response = {"query": self._build_model_response_query(experiment_id)}
        logger.info(
            "[fetch_event_details] experiment_id=%s query_ready: dimension_queries=%s scores/model_response",
            experiment_id,
            len(dimension_query_list),
        )

        # prompt_download=await self.fetch_eval_prompt(experiment_id)
        query_tasks = [self.send_request_to_btql(q) for q in dimension_query_list]
        query_tasks.extend(
            [
                self.send_request_to_btql(scores["query"]),
                self.send_request_to_btql(response["query"]),
            ]
        )
        query_results = await asyncio.gather(*query_tasks)
        dimension_responses = query_results[: len(dimension_query_list)]
        scores_resp = query_results[len(dimension_query_list)]
        response_resp = query_results[len(dimension_query_list) + 1]
        first_dimension_status = dimension_responses[0].status_code if dimension_responses else "NA"
        successful_dimension_count = sum(1 for resp in dimension_responses if resp.status_code == 200)
        logger.info(
            "[fetch_event_details] query_status experiment_id=%s dimension_first=%s dimension_success=%s/%s scores=%s model_response=%s",
            experiment_id,
            first_dimension_status,
            successful_dimension_count,
            len(dimension_responses),
            scores_resp.status_code,
            response_resp.status_code,
        )

        if successful_dimension_count == 0 or scores_resp.status_code != 200:
            if using_agent_queries:
                logger.warning(
                    "[fetch_event_details] btql-agent query execution failed, retrying with metadata fallback. experiment_id=%s dimension_success=%s/%s scores_status=%s",
                    experiment_id,
                    successful_dimension_count,
                    len(dimension_responses),
                    scores_resp.status_code,
                )
                fallback_queries = self._build_btql_agent_failure_fallback_queries(experiment_id)
                dimensions["query"] = fallback_queries["dimensions"]
                scores["query"] = fallback_queries["scores"]
                using_agent_queries = False
                dimension_query_list = fallback_queries.get("dimensions_grouped", [dimensions["query"]])
                query_tasks = [self.send_request_to_btql(q) for q in dimension_query_list]
                query_tasks.extend(
                    [
                        self.send_request_to_btql(scores["query"]),
                        self.send_request_to_btql(response["query"]),
                    ]
                )
                query_results = await asyncio.gather(*query_tasks)
                dimension_responses = query_results[: len(dimension_query_list)]
                scores_resp = query_results[len(dimension_query_list)]
                response_resp = query_results[len(dimension_query_list) + 1]
                successful_dimension_count = sum(1 for resp in dimension_responses if resp.status_code == 200)

        if successful_dimension_count == 0 or scores_resp.status_code != 200:
            sample_dimension_error = ""
            for resp in dimension_responses:
                if resp.status_code != 200:
                    sample_dimension_error = (resp.text or "")[:600]
                    break
            logger.error(
                "[fetch_event_details] query failed: dimension_success=%s/%s scores_status=%s dimension_error=%s scores_error=%s",
                successful_dimension_count,
                len(dimension_responses),
                scores_resp.status_code,
                sample_dimension_error,
                (scores_resp.text or "")[:600],
            )
            return []

        dimension_rows: list[dict[str, Any]] = []
        for idx, dim_resp in enumerate(dimension_responses):
            if dim_resp.status_code != 200:
                logger.warning(
                    "[fetch_event_details] grouped dimension query failed idx=%s status=%s error=%s",
                    idx,
                    dim_resp.status_code,
                    (dim_resp.text or "")[:300],
                )
                continue
            try:
                dimension_payload = dim_resp.json()
                rows = dimension_payload.get("data", [])
                if isinstance(rows, list):
                    dimension_rows.extend(rows)
            except Exception as exc:  # noqa: BLE001
                logger.warning(
                    "[fetch_event_details] failed to parse grouped dimension response json idx=%s: %s",
                    idx,
                    exc,
                )

        if not dimension_rows:
            logger.error("[fetch_event_details] no dimension rows collected from grouped queries")
            return []

        try:
            scores_payload = scores_resp.json()
            scores_rows = scores_payload.get("data", [])
        except Exception as exc:  # noqa: BLE001
            logger.error("[fetch_event_details] failed to parse scores response json: %s", exc)
            return []

        response_rows: list[dict[str, Any]] = []
        if response_resp.status_code == 200:
            try:
                response_payload = response_resp.json()
                response_rows = response_payload.get("data", [])
            except Exception as exc:  # noqa: BLE001
                logger.warning("[fetch_event_details] failed to parse response query json: %s", exc)
        else:
            logger.warning(
                "[fetch_event_details] model response query failed, status=%s, error=%s",
                response_resp.status_code,
                (response_resp.text or "")[:600],
            )
        logger.info(
            "[fetch_event_details] rows experiment_id=%s dimensions=%s scores=%s model_response=%s",
            experiment_id,
            len(dimension_rows),
            len(scores_rows),
            len(response_rows),
        )

        # --- MERGING LOGIC ---

        # Step 0: change input to id by hashing input content
        dimension_rows = await self.change_input_to_id(dimension_rows)

        # Step 1: Create a dictionary of dimension data keyed by root_span_id
        merged_data, score_dimensions = await self.merge_scores(
            dimension_rows,
            scores_rows,
            response_rows=response_rows,
        )
        logger.info(
            "[fetch_event_details] merged experiment_id=%s merged_rows=%s score_dimensions=%s",
            experiment_id,
            len(merged_data),
            score_dimensions,
        )

        # Step 1.1: requery scores for case which merged_scores is empty
        score_query_template = scores['query'] + " AND root_span_id IN [{root_span_id}]"
        merged_data = await self.update_merged_scores_with_requery(merged_data, scores_rows, score_query_template, score_dimensions)

        # Step 2: filter merged_scores according to key_metrics
        await self.filter_key_metrics(merged_data)

        # Step 2.1: compact judge_response before writing event_details*.jsonl
        self._apply_judge_response_truncation(merged_data)
        
        # Step 3: calculate average score from 'merged_scores'
        all_average_scores = await self.get_average_scores(merged_data)

        # Step 4: filter low average score distribution of scores
        filtered_output = await self.filter_low_average_score_cases(merged_data, all_average_scores)

        # Step 5: filter low score which has at least one dimension score below self.dimension_threshold
        filtered_output.extend(await self.filter_low_dimension_cases(merged_data, self.dimension_threshold))


        # Step 6: deduplicate filtered output
        filtered_output = list({row['root_span_id']: row for row in filtered_output}.values())

        # Convert the dictionary of results back to a list of records
        final_output_list = list(merged_data.values())

        #save all output data to jsonl file
        self.output_data_path = f"experiments_{experiment_id}/event_details.jsonl"
        await self.__save_output_data__(final_output_list, self.output_data_path)
        
        # rank low score cases by average_score ascending
        final_output_list.sort(key=lambda x: x.get('average_score', 0.0))
        final_output_list_filtered = final_output_list[:self.bad_case_limit_per_iteration]

        #save filtered output data to jsonl file
        self.output_data_path = f"experiments_{experiment_id}/event_details_filtered.jsonl"
        await self.__save_output_data__(final_output_list_filtered, self.output_data_path)


        prompt_root_span_ids = [
            str(row.get("root_span_id", "")).strip()
            for row in final_output_list_filtered
            if str(row.get("root_span_id", "")).strip()
        ]
        prompt_download = await self.fetch_eval_prompt(experiment_id, prompt_root_span_ids)

        # # Step 7: post-process group by key metrics to only keep cases where all key metrics are below 1
        # filtered_output = await self.post_process_filtered_cases(filtered_output)

        # # Step 8: save filtered_output to jsonl file
        # self.output_data_path = f"experiments_{experiment_id}/event_filtered_group.json"
        # self.__save_output_data__(filtered_output, self.output_data_path)


        # # merge prompt file path
        # if prompt_download.get("success"):
        #     for row in filtered_output:
        #         row['current_prompt_file_path'] = prompt_download.get("file path")

        return final_output_list

    def _event_details_filtered_path(self, experiment_id: str) -> Path:
        return Path(os.path.join(self.output_dir, f"experiments_{experiment_id}", "event_details_filtered.jsonl"))

    async def _read_jsonl_records(self, path: Path) -> list[dict[str, Any]]:
        records: list[dict[str, Any]] = []
        try:
            async with aiofiles.open(path, "r", encoding="utf-8") as f:
                async for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        records.append(json.loads(line))
                    except json.JSONDecodeError:
                        continue
        except FileNotFoundError:
            return []
        return records

    def _extract_root_span_ids(self, records: list[dict[str, Any]]) -> list[str]:
        root_span_ids: list[str] = []
        seen = set()
        for row in records:
            root_span_id = row.get("root_span_id")
            if isinstance(root_span_id, str) and root_span_id and root_span_id not in seen:
                root_span_ids.append(root_span_id)
                seen.add(root_span_id)
        return root_span_ids

    def _get_conversation_compact_limits(self) -> tuple[int, int]:
        """Return (max_messages, max_segment_chars) for conversation compaction."""
        max_messages = int(os.getenv("AUTOPROMPTLAB_CONVERSATION_MAX_MESSAGES", "30"))
        max_segment_chars = int(os.getenv("AUTOPROMPTLAB_CONVERSATION_SEGMENT_MAX_CHARS", "1200"))
        return max(1, max_messages), max(128, max_segment_chars)

    def _truncate_text(self, text: str, max_chars: int) -> str:
        if len(text) <= max_chars:
            return text
        omitted = len(text) - max_chars
        return f"{text[:max_chars]} ... [truncated {omitted} chars]"

    def _truncate_nested_strings(self, value: Any, max_chars: int) -> Any:
        """Recursively truncate long strings to control context growth."""
        if isinstance(value, str):
            return self._truncate_text(value, max_chars)
        if isinstance(value, list):
            return [self._truncate_nested_strings(v, max_chars) for v in value]
        if isinstance(value, tuple):
            return tuple(self._truncate_nested_strings(v, max_chars) for v in value)
        if isinstance(value, dict):
            return {k: self._truncate_nested_strings(v, max_chars) for k, v in value.items()}
        return value

    def _parse_im_serialized_messages(self, payload_text: str) -> list[dict[str, Any]] | None:
        """Parse '<|im_start|>role<|im_sep|>content<|im_end|>' serialized text."""
        pattern = re.compile(r"<\|im_start\|>(.*?)<\|im_sep\|>(.*?)<\|im_end\|>", re.DOTALL)
        matches = pattern.findall(payload_text)
        if not matches:
            return None

        messages: list[dict[str, Any]] = []
        for role_raw, content_raw in matches:
            role = str(role_raw).strip() or "user"
            content_text = str(content_raw).strip()
            messages.append(
                {
                    "role": role,
                    "content": [{"type": "text", "text": content_text}],
                }
            )
        return messages

    def _normalize_conversation_payload(self, payload: Any) -> Any:
        """Normalize payload into a message list when possible."""
        if payload is None:
            return None
        if isinstance(payload, list):
            return payload
        if isinstance(payload, dict):
            messages = payload.get("messages")
            if messages is not None:
                return messages
            return payload
        if isinstance(payload, str):
            text = payload.strip()
            if not text:
                return None

            if text[0] in "[{":
                try:
                    parsed = json.loads(text)
                    if isinstance(parsed, dict):
                        parsed_messages = parsed.get("messages")
                        if parsed_messages is not None:
                            return parsed_messages
                    if isinstance(parsed, list):
                        return parsed
                except json.JSONDecodeError:
                    pass

            im_messages = self._parse_im_serialized_messages(text)
            if im_messages is not None:
                return im_messages

            return [{"role": "system", "content": [{"type": "text", "text": text}]}]
        return payload

    def _extract_messages_payload(self, input_payload: Any, response_payload: Any) -> Any:
        """Prefer input_payload; otherwise derive messages from response payload."""
        normalized_input = self._normalize_conversation_payload(input_payload)
        if normalized_input is not None:
            return normalized_input
        return self._normalize_conversation_payload(response_payload)

    def _conversation_payload_length(self, payload: Any) -> int:
        """Estimate conversation payload length for tie-breaking between BTQL rows."""
        if payload is None:
            return 0
        if isinstance(payload, str):
            return len(payload)
        if isinstance(payload, list):
            total = 0
            for item in payload:
                total += self._conversation_payload_length(item)
            return total
        if isinstance(payload, dict):
            if "messages" in payload:
                return self._conversation_payload_length(payload.get("messages"))
            total = 0
            for value in payload.values():
                total += self._conversation_payload_length(value)
            if total > 0:
                return total
            try:
                return len(json.dumps(payload, ensure_ascii=False, default=str))
            except TypeError:
                return len(str(payload))
        return len(str(payload))

    def _select_longest_conversation_record(
        self,
        data: list[Any],
    ) -> tuple[dict[str, Any] | None, int]:
        """Select the row whose extracted conversation payload is longest."""
        best_record: dict[str, Any] | None = None
        best_len = -1
        for item in data:
            if not isinstance(item, dict):
                continue
            extracted = self._extract_messages_payload(
                item.get("input_payload"),
                item.get("response"),
            )
            current_len = self._conversation_payload_length(extracted)
            if current_len > best_len:
                best_len = current_len
                best_record = item

        if best_record is None and data:
            first = data[0]
            if isinstance(first, dict):
                best_record = first
                extracted = self._extract_messages_payload(
                    first.get("input_payload"),
                    first.get("response"),
                )
                best_len = self._conversation_payload_length(extracted)
            else:
                best_len = 0
        return best_record, max(0, best_len)

    def _compact_conversation_state(
        self,
        conversation_state: Any,
        *,
        max_messages: int | None = None,
        max_segment_chars: int | None = None,
    ) -> Any:
        """Normalize + compact conversation payload while keeping shape stable."""
        default_max_messages, default_max_segment_chars = self._get_conversation_compact_limits()
        effective_max_messages = max_messages if isinstance(max_messages, int) and max_messages > 0 else default_max_messages
        effective_max_segment_chars = (
            max_segment_chars if isinstance(max_segment_chars, int) and max_segment_chars > 0 else default_max_segment_chars
        )
        if not isinstance(conversation_state, list):
            return self._truncate_nested_strings(conversation_state, effective_max_segment_chars)

        original_count = len(conversation_state)
        messages = conversation_state
        if original_count > effective_max_messages:
            head = effective_max_messages // 2
            tail = effective_max_messages - head
            omitted = original_count - effective_max_messages
            sentinel = {
                "role": "system",
                "content": [
                    {
                        "type": "text",
                        "text": f"[omitted {omitted} middle messages for context budget]",
                    }
                ],
            }
            messages = messages[:head] + [sentinel] + messages[-tail:]

        return [self._truncate_nested_strings(msg, effective_max_segment_chars) for msg in messages]

    def _conversation_to_text(self, message: Any) -> str:
        if isinstance(message, str):
            return message
        if isinstance(message, list):
            parts: list[str] = []
            for item in message:
                if isinstance(item, dict):
                    text = item.get("text")
                    if isinstance(text, str) and text.strip():
                        parts.append(text)
                elif isinstance(item, str):
                    parts.append(item)
            return "\n".join(parts)
        if isinstance(message, dict):
            text = message.get("text")
            if isinstance(text, str):
                return text
        return ""

    def _normalize_tool_name(self, raw_name: Any) -> str:
        text = str(raw_name or "").strip().strip("`\"'")
        if not text:
            return ""
        text = re.split(r"[\s(,]", text, maxsplit=1)[0].strip()
        if not re.match(r"^[A-Za-z_][A-Za-z0-9_.-]{1,99}$", text):
            return ""
        return text

    def _shorten_text(self, text: str, limit: int = 140) -> str:
        compact = " ".join(str(text or "").split())
        if len(compact) <= limit:
            return compact
        return compact[: max(1, limit - 3)].rstrip() + "..."

    def _extract_prompt_text_candidates(self, payload: Any, depth: int = 0) -> list[str]:
        """Recursively extract prompt text snippets from eval_prompt payloads."""
        if depth > 6 or payload is None:
            return []
        if isinstance(payload, str):
            text = payload.strip()
            if not text:
                return []
            # IM-serialized payload may contain system/developer/user/assistant turns.
            # For tool-registry extraction we only keep system/developer instructions.
            if "<|im_start|>" in text and "<|im_end|>" in text:
                parsed_messages = self._parse_im_serialized_messages(text)
                if isinstance(parsed_messages, list) and parsed_messages:
                    prompt_parts: list[str] = []
                    for msg in parsed_messages:
                        if not isinstance(msg, dict):
                            continue
                        role = str(msg.get("role", "")).strip().lower()
                        if role not in {"system", "developer"}:
                            continue
                        part = self._conversation_to_text(msg.get("content")).strip()
                        if part:
                            prompt_parts.append(part)
                    if prompt_parts:
                        return ["\n\n".join(prompt_parts)]
                    return []
            return [text]
        if isinstance(payload, list):
            collected: list[str] = []
            for item in payload[:30]:
                collected.extend(self._extract_prompt_text_candidates(item, depth + 1))
            return collected
        if not isinstance(payload, dict):
            return []

        collected: list[str] = []
        if "prompt_text" in payload:
            collected.extend(self._extract_prompt_text_candidates(payload.get("prompt_text"), depth + 1))

        messages = payload.get("messages")
        if isinstance(messages, list):
            prompt_parts: list[str] = []
            for msg in messages:
                if not isinstance(msg, dict):
                    continue
                role = str(msg.get("role", "")).strip().lower()
                if role not in {"system", "developer"}:
                    continue
                text = self._conversation_to_text(msg.get("content")).strip()
                if text:
                    prompt_parts.append(text)
            if prompt_parts:
                collected.append("\n\n".join(prompt_parts))

        for key in ("data", "request", "metadata", "input_payload", "response", "output"):
            if key in payload:
                collected.extend(self._extract_prompt_text_candidates(payload.get(key), depth + 1))

        return collected

    def _collect_prompt_text_variants_from_eval_prompt_records(
        self,
        records: list[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        """Collect prompt variants with hash/count from eval_prompt records."""
        stats: dict[str, dict[str, Any]] = {}
        for record in records:
            if not isinstance(record, dict):
                continue
            candidates = self._extract_prompt_text_candidates(record)
            for candidate in candidates:
                text = candidate.strip()
                if not text:
                    continue
                prompt_hash = hashlib.sha256(text.encode("utf-8")).hexdigest()[:16]
                entry = stats.get(prompt_hash)
                if entry is None:
                    stats[prompt_hash] = {
                        "prompt_hash": prompt_hash,
                        "prompt_text": text,
                        "count": 1,
                        "char_count": len(text),
                    }
                else:
                    entry["count"] = int(entry.get("count", 0)) + 1
                    if len(text) > int(entry.get("char_count", 0)):
                        entry["prompt_text"] = text
                        entry["char_count"] = len(text)
        variants = list(stats.values())
        variants.sort(key=lambda item: (-int(item.get("count", 0)), -int(item.get("char_count", 0))))
        return variants

    def _sanitize_tool_registry_item(self, item: Any) -> dict[str, Any] | None:
        if not isinstance(item, dict):
            return None
        name = self._normalize_tool_name(item.get("tool_name") or item.get("name"))
        if not name:
            return None
        description = self._shorten_text(
            str(item.get("short_description") or item.get("description") or "").strip(),
            limit=160,
        )
        when_to_use = self._shorten_text(str(item.get("when_to_use") or "").strip(), limit=240)
        evidence_snippet = self._shorten_text(
            str(item.get("evidence_snippet") or item.get("evidence") or "").strip(),
            limit=240,
        )
        confidence_raw = item.get("confidence", 0.0)
        try:
            confidence = float(confidence_raw)
            if confidence > 1.0:
                confidence = confidence / 100.0
            confidence = min(1.0, max(0.0, confidence))
        except (TypeError, ValueError):
            confidence = 0.0
        return {
            "name": name,
            "short_description": description,
            "when_to_use": when_to_use,
            "evidence_snippet": evidence_snippet,
            "confidence": confidence,
        }

    def _parse_tool_registry_response(self, content: str) -> list[dict[str, Any]]:
        if not isinstance(content, str):
            return []
        text = content.strip()
        if not text:
            return []

        if "```" in text:
            lines = text.splitlines()
            text = "\n".join(line for line in lines if not line.strip().startswith("```")).strip()

        payload: Any = None
        try:
            payload = json.loads(text)
        except json.JSONDecodeError:
            match = re.search(r"\{[\s\S]*\}|\[[\s\S]*\]", text)
            if match:
                try:
                    payload = json.loads(match.group(0))
                except json.JSONDecodeError:
                    payload = None
        if payload is None:
            return []

        candidates: list[Any]
        if isinstance(payload, dict):
            tools_value = payload.get("tools")
            if isinstance(tools_value, list):
                candidates = tools_value
            else:
                return []
        elif isinstance(payload, list):
            candidates = payload
        else:
            return []

        sanitized: list[dict[str, Any]] = []
        seen: set[str] = set()
        for item in candidates:
            normalized = self._sanitize_tool_registry_item(item)
            if not normalized:
                continue
            name = normalized["name"]
            if name in seen:
                continue
            seen.add(name)
            sanitized.append(normalized)
        return sanitized

    async def _extract_tools_from_prompt_with_llm(self, prompt_text: str) -> list[dict[str, Any]]:
        if self.model_client is None:
            return []
        from autogen_core.models import SystemMessage, UserMessage

        trimmed_prompt = prompt_text.strip()
        # if len(trimmed_prompt) > 24000:
        #     trimmed_prompt = (
        #         trimmed_prompt[:18000]
        #         + "\n...[truncated middle content]...\n"
        #         + trimmed_prompt[-4000:]
        #     )

        system_prompt = (
            "You extract callable tools from a system/developer prompt.\n"
            "Return strict JSON only using this schema:\n"
            "{\n"
            '  "tools": [\n'
            "    {\n"
            '      "name": "<tool_name>",\n'
            '      "short_description": "<one sentence>",\n'
            '      "when_to_use": "<trigger condition>",\n'
            '      "evidence_snippet": "<short quote/paraphrase from prompt>",\n'
            '      "confidence": 0.0\n'
            "    }\n"
            "  ]\n"
            "}\n"
            "Rules:\n"
            "- Include only explicitly callable tools.\n"
            "- Extract tools from explicit declarations and section patterns, including:\n"
            "  1) markdown headers like '### <tool_name>'\n"
            "  2) phrases like 'Decision boundary for <tool_name>'\n"
            "  3) tags like '<situations_where_I_always_use_<tool_name>>' or '<situations_where_I_never_use_<tool_name>>'\n"
            "  4) sections that describe 'when and how I call and use tools'\n"
            "- Treat these pattern matches as valid tool declarations even if no function signature is shown.\n"
            "- Do not include argument keys or generic fields.\n"
            "- Never extract tool names from user messages, assistant outputs, or real tool-call traces.\n"
            "- Keep all tools.\n"
            "- Keep names normalized (snake_case), unique, and confidence between 0 and 1.\n"
        )
        user_prompt = f"Prompt text:\n{trimmed_prompt}"
        result = await self.model_client.create(
            messages=[
                SystemMessage(content=system_prompt),
                UserMessage(content=user_prompt, source="tool_extractor"),
            ]
        )
        raw = result.content.strip() if hasattr(result, "content") and isinstance(result.content, str) else str(result)
        return self._parse_tool_registry_response(raw)

    async def _build_experiment_tool_registry(
        self,
        experiment_id: str,
        root_span_ids: list[str],
    ) -> dict[str, Any]:
        cached = self._experiment_tool_registry_cache.get(experiment_id)
        if isinstance(cached, dict):
            return cached

        eval_prompt_path = Path(self.output_dir) / f"experiments_{experiment_id}" / "eval_prompt.jsonl"
        if not eval_prompt_path.exists() and root_span_ids:
            await self.fetch_eval_prompt(experiment_id, root_span_ids)

        records = await self._read_jsonl_records(eval_prompt_path)
        if not records:
            result = {
                "ok": False,
                "experiment_id": experiment_id,
                "source": "eval_prompt_llm",
                "tools": [],
                "error": "eval_prompt_records_empty",
            }
            self._experiment_tool_registry_cache[experiment_id] = result
            return result

        prompt_candidates = self._extract_prompt_text_candidates(records[0])
        prompt_text = prompt_candidates[0] if prompt_candidates else ""
        if not isinstance(prompt_text, str) or not prompt_text.strip():
            result = {
                "ok": False,
                "experiment_id": experiment_id,
                "source": "eval_prompt_llm",
                "tools": [],
                "error": "prompt_text_missing_in_eval_prompt",
            }
            self._experiment_tool_registry_cache[experiment_id] = result
            return result

        tools = await self._extract_tools_from_prompt_with_llm(prompt_text)
        result = {
            "ok": True,
            "experiment_id": experiment_id,
            "source": "eval_prompt_llm",
            "tools": tools,
        }
        if not tools:
            if self.model_client is None:
                result["warning"] = "model_client_not_available_for_tool_registry"
            else:
                result["warning"] = "tool_registry_empty_after_llm_extraction"
        self._experiment_tool_registry_cache[experiment_id] = result
        return result

    def _summarize_tool_trace(self, conversation_state: Any) -> dict[str, Any]:
        """Summarize actual tool calls from conversation messages."""
        if not isinstance(conversation_state, list):
            return {
                "called_tools": [],
                "called_tool_count": 0,
                "has_tool_calls": False,
            }

        called_counts: dict[str, int] = {}
        json_name_pattern = re.compile(r'"(?:name|tool_name|function_name)"\s*:\s*"([A-Za-z_][A-Za-z0-9_.-]{1,80})"')
        xml_name_pattern = re.compile(r"<tool_name>\s*([A-Za-z_][A-Za-z0-9_.-]{1,80})\s*</tool_name>", re.IGNORECASE)

        def _add_name(raw_name: Any) -> None:
            name = self._normalize_tool_name(raw_name)
            if not name:
                return
            called_counts[name] = called_counts.get(name, 0) + 1

        for msg in conversation_state:
            if not isinstance(msg, dict):
                continue
            role = str(msg.get("role", "")).strip().lower()
            text = self._conversation_to_text(msg.get("content"))

            if role == "tool":
                _add_name(msg.get("name"))
                _add_name(msg.get("tool_name"))
                _add_name(msg.get("tool"))

            function_call = msg.get("function_call")
            if isinstance(function_call, dict):
                _add_name(function_call.get("name"))

            tool_calls = msg.get("tool_calls")
            if isinstance(tool_calls, list):
                for call in tool_calls:
                    if not isinstance(call, dict):
                        continue
                    _add_name(call.get("name"))
                    function_payload = call.get("function")
                    if isinstance(function_payload, dict):
                        _add_name(function_payload.get("name"))

            for candidate in json_name_pattern.findall(text):
                _add_name(candidate)
            for candidate in xml_name_pattern.findall(text):
                _add_name(candidate)

        called_tools = [
            {"name": name, "count": count}
            for name, count in sorted(called_counts.items(), key=lambda kv: (-kv[1], kv[0]))
        ]
        return {
            "called_tools": called_tools,
            "called_tool_count": len(called_tools),
            "has_tool_calls": bool(called_tools),
        }

    def _build_tool_gap_assessment(
        self,
        available_tools: list[dict[str, str]],
        tool_trace_summary: dict[str, Any],
    ) -> dict[str, Any]:
        """Compare available tools from prompt vs actual tool calls."""
        available_names = []
        available_desc_map: dict[str, str] = {}
        for item in available_tools:
            if not isinstance(item, dict):
                continue
            name = self._normalize_tool_name(item.get("name"))
            if not name:
                continue
            if name not in available_desc_map:
                available_names.append(name)
                available_desc_map[name] = str(item.get("description", "")).strip()

        called_names = []
        called_raw = tool_trace_summary.get("called_tools") if isinstance(tool_trace_summary, dict) else []
        if isinstance(called_raw, list):
            for item in called_raw:
                if not isinstance(item, dict):
                    continue
                name = self._normalize_tool_name(item.get("name"))
                if name and name not in called_names:
                    called_names.append(name)

        available_set = set(available_names)
        called_set = set(called_names)
        unused_available = sorted(available_set - called_set)
        unknown_called = sorted(called_set - available_set)
        matched_called = sorted(called_set & available_set)

        return {
            "available_tool_count": len(available_names),
            "called_tool_count": len(called_names),
            "unused_available_tools": [
                {"name": name, "description": available_desc_map.get(name, "")}
                for name in unused_available
            ],
            "unknown_called_tools": unknown_called,
            "matched_called_tools": matched_called,
            "no_tool_called_despite_available_tools": bool(available_names) and not bool(called_names),
        }

    def _project_conversation_for_strategy(
        self,
        conversation_state: Any,
        *,
        strategy: str,
        max_messages: int,
    ) -> Any:
        if not isinstance(conversation_state, list):
            return conversation_state
        if not conversation_state:
            return conversation_state

        s = (strategy or "full").strip().lower()
        if s == "recent":
            return conversation_state[-max_messages:]

        if s == "tool_focus":
            tool_keywords = ("tool", "function_call", "tool_call", "action", "plugin")
            selected: list[Any] = []
            for msg in conversation_state:
                if not isinstance(msg, dict):
                    continue
                role = str(msg.get("role", "")).lower()
                content = self._conversation_to_text(msg.get("content"))
                if role == "tool" or any(k in content.lower() for k in tool_keywords):
                    selected.append(msg)
            if not selected:
                return conversation_state[-max_messages:]
            if len(selected) > max_messages:
                return selected[-max_messages:]
            return selected

        return conversation_state

    def _build_conversation_query(self, experiment_id: str, root_span_id: str) -> str:
        environment = (self.environment or "global").strip().lower()
        if environment == "china":
            return f"""from: experiment('{experiment_id}')
| select: root_span_id,
    metadata.request.messages AS input_payload,
    metadata.request AS response
| filter: metadata.request.messages != "" AND root_span_id = '{root_span_id}'
| limit: 10"""

        return f"""from: experiment('{experiment_id}')
| select: root_span_id,
    metadata.request AS input_payload

| filter: root_span_id = '{root_span_id}' AND metadata.request != ""
| limit: 10"""

    async def fetch_conversation_state_by_root_span_id(
        self,
        experiment_id: str,
        root_span_id: str,
        *,
        compact: bool = True,
        max_messages: int | None = None,
        max_segment_chars: int | None = None,
        strategy: str = "full",
    ) -> dict[str, Any]:
        """
        Fetch full conversation state for a specific root_span_id.
        """
        if self.local_debug:
            return {
                "ok": True,
                "root_span_id": root_span_id,
                "conversation_state": None,
                "judge_response": None,
                "test_model_response": None,
                "raw": None,
            }

        query = self._build_conversation_query(experiment_id, root_span_id)
        response = await self.send_request_to_btql(query)
        if response.status_code != 200:
            return {"ok": False, "root_span_id": root_span_id, "error": response.text}

        payload = response.json()
        data = payload.get("data", [])
        if not data:
            return {"ok": False, "root_span_id": root_span_id, "error": "No data found for root_span_id"}

        record, selected_payload_length = self._select_longest_conversation_record(data)
        if not isinstance(record, dict):
            return {
                "ok": False,
                "root_span_id": root_span_id,
                "error": "No valid record found for root_span_id",
            }
        raw_input_payload = record.get("input_payload")
        raw_response_payload = record.get("response")
        conversation_state_full = self._extract_messages_payload(raw_input_payload, raw_response_payload)
        projected_state = self._project_conversation_for_strategy(
            conversation_state_full,
            strategy=strategy,
            max_messages=max_messages if isinstance(max_messages, int) and max_messages > 0 else 80,
        )
        if compact:
            conversation_state = self._compact_conversation_state(
                projected_state,
                max_messages=max_messages,
                max_segment_chars=max_segment_chars,
            )
        else:
            conversation_state = projected_state

        return {
            "ok": True,
            "root_span_id": root_span_id,
            "conversation_state": conversation_state,
            "conversation_state_full": conversation_state_full,
            "strategy": strategy,
            "candidate_count": len(data),
            "selected_payload_length": selected_payload_length,
            # These two fields are sourced from event_details records in
            # analyze_root_cases_from_event_details, not this conversation query.
            "judge_response": None,
            "test_model_response": None,
            "raw": record,
        }

    def _event_record_richness_score(self, record: dict[str, Any]) -> int:
        """Estimate how informative an event_details row is for merge preference."""
        return (
            self._conversation_payload_length(record.get("judge_response"))
            + self._conversation_payload_length(record.get("test_model_response"))
            + self._conversation_payload_length(record.get("merged_scores"))
            + self._conversation_payload_length(record.get("average_score"))
            + self._conversation_payload_length(record.get("input"))
        )

    def _build_event_details_by_root_span_id(
        self,
        records: list[dict[str, Any]],
    ) -> dict[str, dict[str, Any]]:
        """
        Build root_span_id -> event row map.

        If duplicated root_span_id exists, keep the richer row so downstream
        analysis gets the most complete judge/model/score payload.
        """
        event_map: dict[str, dict[str, Any]] = {}
        for record in records:
            if not isinstance(record, dict):
                continue
            root_span_id = record.get("root_span_id")
            if not isinstance(root_span_id, str) or not root_span_id:
                continue
            existing = event_map.get(root_span_id)
            if existing is None:
                event_map[root_span_id] = record
                continue
            if self._event_record_richness_score(record) > self._event_record_richness_score(existing):
                event_map[root_span_id] = record
        return event_map

    async def _fetch_conversation_states(
        self,
        experiment_id: str,
        root_span_ids: list[str],
        max_concurrency: int = 3,
        event_details_by_root_span_id: dict[str, dict[str, Any]] | None = None,
    ) -> dict[str, dict[str, Any]]:
        semaphore = asyncio.Semaphore(max_concurrency)

        async def fetch_one(root_span_id: str) -> tuple[str, dict[str, Any]]:
            async with semaphore:
                data = await self.fetch_conversation_state_by_root_span_id(experiment_id, root_span_id)
                if isinstance(data, dict):
                    conversation_full = data.get("conversation_state_full")
                    if conversation_full is None:
                        conversation_full = data.get("conversation_state")
                    tool_trace_summary = self._summarize_tool_trace(conversation_full)
                    data["tool_trace_summary"] = tool_trace_summary
                    event_row = (
                        event_details_by_root_span_id.get(root_span_id)
                        if event_details_by_root_span_id
                        else None
                    )
                    if isinstance(event_row, dict):
                        for key in ("judge_response", "test_model_response", "merged_scores", "average_score", "input"):
                            if key in event_row and event_row.get(key) is not None:
                                data[key] = event_row.get(key)
                return root_span_id, data

        tasks = [asyncio.create_task(fetch_one(root_span_id)) for root_span_id in root_span_ids]
        results: dict[str, dict[str, Any]] = {}
        for root_span_id, data in await asyncio.gather(*tasks):
            results[root_span_id] = data
        return results

    def _safe_json_loads(self, content: str) -> dict[str, Any] | None:
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            return None

    def _build_root_case_prompt(self, root_span_id: str, conversation_state: Any, judge_response: Any, test_model_response: Any, event_row: dict[str, Any]) -> str:
        payload = {
            "root_span_id": root_span_id,
            "conversation_state": conversation_state,
            "judge_response": judge_response,
            "test_model_response": test_model_response,
            "event_row": {
                "input": event_row.get("input"),
                "merged_scores": event_row.get("merged_scores", {}),
                "average_score": event_row.get("average_score"),
            },
        }
        return json.dumps(payload, ensure_ascii=False, indent=2)

    async def _analyze_root_case_with_llm(
        self,
        *,
        root_span_id: str,
        conversation_state: Any,
        judge_response: Any,
        test_model_response: Any,
        event_row: dict[str, Any],
    ) -> dict[str, Any]:
        if self.model_client is None:
            return {"ok": False, "root_span_id": root_span_id, "error": "model_client not set in BraintrustClient"}

        from autogen_core.models import SystemMessage, UserMessage

        system_prompt = (
            """You are a Failure Reason Analyzer.
Role:
Analyze a SINGLE failed evaluation case by deeply inspecting the full
conversation, tool usage, page context, and model output.
Your task is to explain WHY the model failed in this case,
not how to fix it.

Your output will be consumed by a higher-level BadCaseAnalyzer
that performs clustering and decision-making.

---

INPUT YOU WILL RECEIVE:
- Full conversation transcript (system, developer, user, assistant, tool calls)
- Page context (if any), including metadata
- Final model response
- Evaluation score and evaluator notes (if available)

---

YOUR TASKS (STRICT ORDER):

1. Reconstruct the execution reality:
   - Did the model actually receive page content?
   - Was the page empty, login-gated, error page, or content-rich?
   - Were there conflicting system instructions?
   - Did tools already return sufficient information?

2. Identify the PRIMARY failure reason for THIS case.
   Choose ONE dominant reason from the taxonomy below.
   (You may list secondary reasons, but one must be primary.)

3. Extract concrete evidence from the conversation
   that supports this failure diagnosis.

4. Describe the failure in neutral, factual terms
   WITHOUT proposing fixes or improvements.

---

FAILURE REASON TAXONOMY (use these labels verbatim):

- "context_inaccessible":
    Page or required content was not actually accessible
    (login-gated, permission denied, empty page, error page).

- "context_misdetected":
    Page content was provided, but the model incorrectly believed
    it was missing or inaccessible.

- "context_underutilized":
    Page content was available and relevant, but the model
    failed to extract or use it.

- "policy_overtrigger":
    Model incorrectly applied safety or policy constraints
    to otherwise allowed content.

- "hallucinated_external_knowledge":
    Model introduced facts not present in the page or conversation
    and not requested by the user.

- "tool_misuse":
    Model invoked tools unnecessarily, redundantly,
    or instead of answering with available information.

- "instruction_conflict":
    Conflicting system/developer instructions led to incorrect behavior.

- "sparse_context_overreach":
    Page content was minimal, but the model extrapolated
    unsupported details.

- "evaluation_misalignment":
    Model response may be reasonable, but failed evaluator expectations
    (format, style, scope mismatch).

- "unknown_or_mixed":
    Failure cannot be confidently attributed to a single cause.

---

OUTPUT FORMAT (STRICT JSON):

{
  "case_id": "<case_id>",
  "primary_failure_reason": "<one taxonomy label>",
  "secondary_factors": ["<optional labels>"],
  "failure_explanation": "<concise explanation of why the model failed in this case>",
  "evidence": {
    "conversation_snippet": "<verbatim excerpt showing the failure>",
    "page_context_summary": "<what the page actually contained>",
    "model_assumption": "<what the model appeared to believe>",
    "reality_mismatch": "<how that belief differed from reality>"
  },
  "actionability_hint": {
    "is_prompt_fixable": true | false | "uncertain",
    "notes": "<purely descriptive, no optimization suggestions>"
  }
}

---

RULES:
- Do NOT suggest prompt changes.
- Do NOT cluster cases.
- Do NOT judge whether this case should be ignored.
- Do NOT speculate beyond the provided data.
- Be precise, evidence-backed, and neutral.
"""
            # "You are analyzing a root case failure.\n"
            # "Use conversation_state + judge_response to identify the root cause.\n"
            # "Return a single JSON object with keys:\n"
            # "root_span_id, issue_type, root_cause, evidence, fix_suggestion, requires_tool_call.\n"
            # "evidence must include conversation_snippet, judge_response_snippet, model_response_snippet.\n"
            # "Do not wrap output in code fences."
        )

        user_prompt = self._build_root_case_prompt(
            root_span_id,
            conversation_state,
            judge_response,
            test_model_response,
            event_row,
        )

        result = await self.model_client.create(
            messages=[
                SystemMessage(content=system_prompt),
                UserMessage(content=user_prompt, source="user"),
            ]
        )

        content = result.content.strip() if hasattr(result, "content") else ""
        # parsed = self._safe_json_loads(content) if isinstance(content, str) else None
        parsed = {'raw_llm_output': content}
        # if not parsed:
        #     return {"ok": False, "root_span_id": root_span_id, "error": "LLM output is not valid JSON", "raw": content}

        parsed["root_span_id"] = root_span_id
        return parsed

    async def _analyze_root_cases_batch(
        self,
        *,
        records: list[dict[str, Any]],
        conversation_map: dict[str, dict[str, Any]],
        max_concurrency: int = 2,
    ) -> list[dict[str, Any]]:
        semaphore = asyncio.Semaphore(max_concurrency)

        async def analyze_one(record: dict[str, Any]) -> dict[str, Any]:
            root_span_id = record.get("root_span_id")
            if not isinstance(root_span_id, str):
                return {"ok": False, "error": "missing root_span_id"}
            conversation_data = conversation_map.get(root_span_id, {})
            async with semaphore:
                return await self._analyze_root_case_with_llm(
                    root_span_id=root_span_id,
                    conversation_state=conversation_data.get("conversation_state"),
                    judge_response=conversation_data.get("judge_response") or record.get("judge_response"),
                    test_model_response=conversation_data.get("test_model_response") or record.get("test_model_response"),
                    event_row=record,
                )

        tasks = [asyncio.create_task(analyze_one(record)) for record in records]
        return list(await asyncio.gather(*tasks))

    async def analyze_root_cases_from_event_details(
        self,
        experiment_id: str,
        fetch_concurrency: int = 3,
        analyze_concurrency: int = 2,
    ) -> dict[str, Any]:
        """
        Analyze root cases using conversation state + judge response for each root_span_id
        in event_details_filtered.jsonl of the given experiment.
        """
        
        logger.info(f"[analyze_root_cases_from_event_details] self.output_dir: {self.output_dir}")
        event_path = self._event_details_filtered_path(experiment_id)
        logger.info(f"[analyze_root_cases_from_event_details] event_path: {event_path}")
        
        records = await self._read_jsonl_records(event_path)
        if not records:
            return {"ok": False, "error": f"No records found in {event_path}"}

        # Deduplicate by root_span_id before analysis.
        # If duplicates exist, keep the richer row (see _build_event_details_by_root_span_id).
        dedup_event_map = self._build_event_details_by_root_span_id(records)
        dedup_root_span_ids = self._extract_root_span_ids(records)
        deduped_records = [
            dedup_event_map[root_span_id]
            for root_span_id in dedup_root_span_ids
            if root_span_id in dedup_event_map
        ]
        if deduped_records and len(deduped_records) != len(records):
            logger.info(
                "[analyze_root_cases_from_event_details] deduplicated records by root_span_id: %d -> %d",
                len(records),
                len(deduped_records),
            )
            records = deduped_records

        if self.bad_case_limit_per_iteration is not None:
            records = records[:self.bad_case_limit_per_iteration]

        root_span_ids = self._extract_root_span_ids(records)
        event_details_by_root_span_id = self._build_event_details_by_root_span_id(records)
        conversation_map = await self._fetch_conversation_states(
            experiment_id,
            root_span_ids,
            max_concurrency=fetch_concurrency,
            event_details_by_root_span_id=event_details_by_root_span_id,
        )

        if self.analysis_agent_client is not None and self.model_client is not None:
            # Skill-guided per-case analysis:
            # 1) first pass uses compact conversation_state.
            # 2) analysis agent may request richer context via conversation_fetcher.
            tool_registry = await self._build_experiment_tool_registry(experiment_id, root_span_ids)
            available_tools_for_experiment = tool_registry.get("tools", []) if isinstance(tool_registry, dict) else []
            if not isinstance(available_tools_for_experiment, list):
                available_tools_for_experiment = []
            tool_registry_meta = {
                "source": str(tool_registry.get("source", "eval_prompt_llm")) if isinstance(tool_registry, dict) else "eval_prompt_llm",
                "prompt_hash": tool_registry.get("prompt_hash") if isinstance(tool_registry, dict) else None,
                "prompt_variant_count": int(tool_registry.get("prompt_variant_count", 0)) if isinstance(tool_registry, dict) else 0,
                "tool_count": len(available_tools_for_experiment),
                "warning": tool_registry.get("warning") if isinstance(tool_registry, dict) else None,
            }
            enriched_records = []
            for record in records:
                root_span_id = record.get("root_span_id")
                conv = conversation_map.get(root_span_id, {}) if root_span_id else {}
                available_tools_snapshot = [
                    dict(item) for item in available_tools_for_experiment if isinstance(item, dict)
                ]
                tool_trace_summary = conv.get("tool_trace_summary", {})
                if not isinstance(tool_trace_summary, dict):
                    tool_trace_summary = {}
                enriched_records.append(
                    {
                        **record,
                        "experiment_id": experiment_id,
                        "conversation_state": conv.get("conversation_state"),
                        "available_tools_from_system_prompt": available_tools_snapshot,
                        "tool_trace_summary": tool_trace_summary,
                        "tool_gap_assessment": self._build_tool_gap_assessment(
                            available_tools_snapshot,
                            tool_trace_summary,
                        ),
                        "tool_registry_meta": tool_registry_meta,
                        "judge_response": conv.get("judge_response") or record.get("judge_response"),
                        "test_model_response": conv.get("test_model_response")
                        or record.get("test_model_response"),
                    }
                )

            async def fetch_full_conversation_context(
                *,
                experiment_id: str,
                root_span_id: str,
                strategy: str = "full",
                max_messages: int = 80,
                max_segment_chars: int = 1800,
            ) -> dict[str, Any]:
                cached = conversation_map.get(root_span_id, {})
                full_state = cached.get("conversation_state_full")
                if full_state is None:
                    fetched = await self.fetch_conversation_state_by_root_span_id(
                        experiment_id,
                        root_span_id,
                        compact=False,
                        strategy=strategy,
                        max_messages=max_messages,
                        max_segment_chars=max_segment_chars,
                    )
                    if not fetched.get("ok"):
                        return fetched
                    full_state = fetched.get("conversation_state")
                projected = self._project_conversation_for_strategy(
                    full_state,
                    strategy=strategy,
                    max_messages=max_messages,
                )
                # For full strategy, return expanded context directly to preserve
                # complete evidence for root-cause verification.
                if (strategy or "full").strip().lower() == "full":
                    conversation_state = projected
                else:
                    conversation_state = self._compact_conversation_state(
                        projected,
                        max_messages=max_messages,
                        max_segment_chars=max_segment_chars,
                    )
                return {
                    "ok": True,
                    "root_span_id": root_span_id,
                    "strategy": strategy,
                    "conversation_state": conversation_state,
                }

            agent_result = await self.analysis_agent_client.analyze_experiment_failures(
                records=enriched_records,
                model_client=self.model_client,
                experiment_id=experiment_id,
                max_concurrency=analyze_concurrency,
                conversation_fetcher=fetch_full_conversation_context,
                save_output=False,
            )
            if not agent_result.get("ok"):
                return agent_result

            # Re-save to output_dir so callers read from a consistent location
            analysis_results = agent_result.get("results", [])
        else:
            analysis_results = await self._analyze_root_cases_batch(
                records=records,
                conversation_map=conversation_map,
                max_concurrency=analyze_concurrency,
            )

        output_path = f"experiments_{experiment_id}/root_case_analysis.jsonl"
        await self.__save_output_data__(analysis_results, output_path)

        return {
            "ok": True,
            "count": len(analysis_results),
            "file_path": str(Path(self.output_dir) / output_path),
        }
    
    async def get_low_case_details(
        self,
        experiment_id: str,
        include_analysis: bool | None = None,
        include_conversation: bool = False,
        root_span_id: str | None = None,
    ) -> Any:
        """
        Get evaluation event details from local event_details_filtered.jsonl.

        - include_analysis: when True, run skill-guided root cause analysis and return results.
          Defaults to self.badcase_deep_analysis (configurable at client construction time).
        - include_conversation=True: return conversation state for a specific root_span_id.
        """
        if include_analysis is None:
            include_analysis = self.badcase_deep_analysis
        if include_analysis:
            analysis_result = await self.analyze_root_cases_from_event_details(experiment_id)
            if not analysis_result.get("ok"):
                return analysis_result
            file_path = analysis_result.get("file_path")
            if not file_path:
                return {"ok": False, "error": "analysis file_path missing"}
            return await self.__read_file__(file_path)

        if include_conversation:
            if not root_span_id:
                return {"ok": False, "error": "root_span_id is required when include_conversation=True"}
            return await self.fetch_conversation_state_by_root_span_id(experiment_id, root_span_id)

        logger.info(f"[get_low_case_details] self.output_dir: {self.output_dir}")
        file_path = os.path.join(self.output_dir, f"experiments_{experiment_id}", "event_details_filtered.jsonl")
        logger.info(f"[get_low_case_details] file_path: {file_path}")
        return await self.__read_file__(file_path)
    

    def _normalize_score_value(self, metric_name: str, score_value: Any) -> Any:
        normalized_value = score_value
        if isinstance(score_value, str):
            try:
                normalized_value = float(score_value)
            except ValueError:
                return normalized_value

        if isinstance(normalized_value, (int, float)):
            metric_config = self.metrics_config_dict.get(metric_name, {})
            optimization_direction = metric_config.get("optimization_direction", "higher_better")
            if optimization_direction == "lower_better":
                return 1.0 - float(normalized_value)
        return normalized_value

    def _normalize_scores(self, scores: dict[str, Any]) -> dict[str, Any]:
        return {
            metric_name: self._normalize_score_value(metric_name, score_value)
            for metric_name, score_value in scores.items()
        }

    def _merge_dict_recursive(self, base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]:
        def _is_empty_value(value: Any) -> bool:
            if value is None:
                return True
            if isinstance(value, str):
                return value.strip() == ""
            if isinstance(value, (dict, list, tuple, set)):
                return len(value) == 0
            return False

        for k, v in extra.items():
            if k in base and isinstance(base[k], dict) and isinstance(v, dict):
                base[k] = self._merge_dict_recursive(base[k], v)
            elif k in base and (not _is_empty_value(base[k])) and _is_empty_value(v):
                # Do not let empty incoming values overwrite non-empty existing values.
                continue
            else:
                base[k] = v
        return base

    def _merge_dimension_rows_by_root_span_id(self, dimension_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Merge duplicate dimension rows for same root_span_id into one stable row."""
        def _is_trivial_judge_payload(value: Any) -> bool:
            if value is None:
                return True
            if isinstance(value, str):
                s = value.strip()
                if not s or s in {"{}", "[]", "null", '""'}:
                    return True
                parsed = self._safe_json_loads(s)
                if isinstance(parsed, (dict, list)) and len(parsed) == 0:
                    return True
                return False
            if isinstance(value, (dict, list, tuple, set)):
                return len(value) == 0
            return False

        grouped: dict[str, list[dict[str, Any]]] = {}
        for row in dimension_rows:
            if not isinstance(row, dict):
                continue
            root_id = row.get("root_span_id")
            if not isinstance(root_id, str) or not root_id:
                continue
            grouped.setdefault(root_id, []).append(row)

        merged_rows: list[dict[str, Any]] = []
        duplicate_root_count = 0

        for root_id, rows in grouped.items():
            if len(rows) > 1:
                duplicate_root_count += 1

            merged_row: dict[str, Any] = {"root_span_id": root_id}
            judge_candidates: list[str] = []
            judge_dict_candidates: list[dict[str, Any]] = []
            for row in rows:
                for k, v in row.items():
                    if k == "root_span_id":
                        continue
                    if k == "judge_response":
                        if isinstance(v, dict):
                            if v:
                                judge_dict_candidates.append(v)
                                judge_candidates.append(json.dumps(v, ensure_ascii=False))
                        elif isinstance(v, str) and v.strip():
                            if not _is_trivial_judge_payload(v):
                                judge_candidates.append(v)
                            parsed = self._safe_json_loads(v)
                            if isinstance(parsed, dict) and parsed:
                                judge_dict_candidates.append(parsed)
                        elif v is not None and not _is_trivial_judge_payload(v):
                            judge_candidates.append(str(v))
                        continue

                    if k == "input":
                        current = merged_row.get("input")
                        if (current is None or current == "") and v not in (None, ""):
                            merged_row["input"] = v
                        continue

                    current = merged_row.get(k)
                    if current in (None, "") and v not in (None, ""):
                        merged_row[k] = v

            # Pre-truncate judge candidates proportionally so the merged
            # result fits within MAX_JUDGE_RESPONSE_ANALYSIS_CHARS.
            max_judge_chars = int(os.getenv("MAX_JUDGE_RESPONSE_ANALYSIS_CHARS", "2000"))

            if judge_dict_candidates and len(judge_dict_candidates) == len(judge_candidates):
                merged_judge: dict[str, Any] = {}
                for d in judge_dict_candidates:
                    merged_judge = self._merge_dict_recursive(merged_judge, d)
                if merged_judge:
                    merged_text = json.dumps(merged_judge, ensure_ascii=False)
                    merged_row["judge_response"] = self._truncate_tail(merged_text, max_judge_chars)
                elif judge_candidates:
                    judge_candidates = self._pre_truncate_judge_candidates(judge_candidates, max_judge_chars)
                    merged_row["judge_response"] = ", ".join(judge_candidates)
            elif judge_candidates:
                # Fall back to the most informative textual value.
                judge_candidates = self._pre_truncate_judge_candidates(judge_candidates, max_judge_chars)
                merged_row["judge_response"] = ", ".join(judge_candidates)

            merged_rows.append(merged_row)

        if duplicate_root_count > 0:
            logger.info(
                "[merge_dimension_rows] merged duplicate root_span_id rows: before=%s after=%s duplicate_roots=%s",
                len(dimension_rows),
                len(merged_rows),
                duplicate_root_count,
            )
        empty_judge_count = sum(1 for row in merged_rows if not str(row.get("judge_response", "")).strip())
        if empty_judge_count > 0:
            logger.info(
                "[merge_dimension_rows] empty_judge_response_rows=%s/%s",
                empty_judge_count,
                len(merged_rows),
            )
        return merged_rows

    async def merge_scores(
        self,
        dimension_rows: list[dict[str, Any]],
        scores_rows: list[dict[str, Any]],
        response_rows: list[dict[str, Any]] | None = None,
    ) -> tuple[dict[str, Any], int]:
        dimension_rows = self._merge_dimension_rows_by_root_span_id(dimension_rows)
        merged_data = {
            row["root_span_id"]: row
            for row in dimension_rows
            if isinstance(row, dict) and row.get("root_span_id")
        }
        max_dimensions = 0

        for row in scores_rows:
            root_id = row.get("root_span_id")
            if root_id not in merged_data:
                continue

            scores = row.get("scores")
            if not isinstance(scores, dict) or not scores:
                continue

            target_row = merged_data[root_id]
            target_scores = target_row.setdefault("merged_scores", {})
            target_scores.update(self._normalize_scores(scores))

            max_dimensions = max(max_dimensions, len(target_scores))
            target_row.pop("scores", None)

        for target_row in merged_data.values():
            target_row.setdefault("input", "")
            target_row.setdefault("judge_response", "")

        # Merge separately queried model responses.
        if response_rows:
            latest_responses = self._pick_latest_model_responses_by_root_span_id(response_rows)
            attached_count = 0
            missing_count = 0
            for root_id, target_row in merged_data.items():
                latest_row = latest_responses.get(root_id)
                if latest_row is None:
                    target_row.setdefault("test_model_response", "")
                    missing_count += 1
                    continue
                model_response = latest_row.get("test_model_response")
                if model_response is None:
                    target_row.setdefault("test_model_response", "")
                    missing_count += 1
                else:
                    target_row["test_model_response"] = model_response
                    attached_count += 1
            logger.info(
                "[merge_scores] model_response_merge dimension_rows=%s response_rows=%s selected_root_spans=%s attached=%s missing=%s",
                len(merged_data),
                len(response_rows),
                len(latest_responses),
                attached_count,
                missing_count,
            )
        else:
            for target_row in merged_data.values():
                target_row.setdefault("test_model_response", "")
            logger.info(
                "[merge_scores] model_response_merge skipped (response_rows empty), dimension_rows=%s",
                len(merged_data),
            )
        return merged_data, max_dimensions
    
    async def get_average_scores(self, merged_data: dict[str, Any]) -> list[float]:
        all_average_scores = []
        for row in merged_data.values():
            merged_scores = row.get('merged_scores', {})
            if merged_scores:
                score_values = [float(v) for k, v in merged_scores.items() if isinstance(v, (int, float))]
                if score_values:
                    average_score = sum(score_values) / len(score_values)
                    row['average_score'] = average_score
                    all_average_scores.append(row['average_score'])
                else:
                    row['average_score'] = 1000
            else:
                row['average_score'] = 1000
            

        return all_average_scores

    async def filter_low_average_score_cases(self, merged_data: dict[str, Any], all_average_scores: list[float] = None) -> list[dict[str, Any]]:
        all_average_scores.sort()
        score_threshold_index = int(len(all_average_scores) * 0.1)  # bottom 10%
        score_threshold = all_average_scores[score_threshold_index] if score_threshold_index < len(all_average_scores) else 1.0
        score_threshold = 1.0 if score_threshold > 1.0 else score_threshold
        filtered_output = [row for row in merged_data.values() if row.get('average_score', 0.0) < score_threshold] 

        return filtered_output

    async def filter_key_metrics(self, merged_data: dict[str, Any]) -> dict[str, Any]:
        if not self.key_metrics:
            return merged_data

        key_metrics = set(self.key_metrics)
        # llm_judged_metrics = {
        #     metric_name
        #     for metric_name in key_metrics
        #     if self.metrics_config_dict.get(metric_name, {}).get("is_llm_judged", True)
        # }

        for row in merged_data.values():
            merged_scores = row.get("merged_scores") or {}
            row["merged_scores"] = {k: v for k, v in merged_scores.items() if k in key_metrics}

            # raw_judge_response = row.get("judge_response")
            # if isinstance(raw_judge_response, dict):
            #     judge_response = raw_judge_response
            # else:
            #     judge_response = self._safe_json_loads(raw_judge_response or "{}")

            # if not isinstance(judge_response, dict):
            #     judge_response = {}

            # row["judge_response"] = json.dumps(
            #     {k: v for k, v in judge_response.items() if k in llm_judged_metrics},
            #     ensure_ascii=False,
            # )

        return merged_data

    async def filter_low_dimension_cases(self, merged_data: dict[str, Any], threshold: float) -> dict[str, Any]:
        filtered_output = []
        for row in merged_data.values():
            row['merged_scores'] = row.get('merged_scores', {})
            low_score_dimensions = [k for k, v in row['merged_scores'].items() if isinstance(v, (int, float)) and float(v) <= threshold]
            if low_score_dimensions:
                filtered_output.append(row)

        return filtered_output

    async def post_process_filtered_cases(self, filtered_output: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Collect cases whose configured key metrics are all below 1."""
        if not self.key_metrics:
            return filtered_output
        key_metric_group = {}
        for key_metric in self.key_metrics:
            if key_metric not in key_metric_group.keys():
                key_metric_group[f"low_{key_metric}_score_cases"] = []
            for row in filtered_output:
                merged_scores = row.get('merged_scores', {})
                if key_metric in merged_scores and isinstance(merged_scores[key_metric], (int, float)) and float(merged_scores[key_metric]) < 1.0:
                    # only keep relevant key metric score in the case
                    row['merged_scores'] = {key_metric: merged_scores[key_metric]}

                    judge_response = row.get('judge_response', '{}')
                    # re to extract json part from  key_metric judge_response
                    pattern = re.compile(rf'"{key_metric}":\s*(\{{.*?\}})', re.DOTALL)
                    match = pattern.search(judge_response)
                    row['judge_response']   = match.group(1) if match else judge_response
                    key_metric_group[f"low_{key_metric}_score_cases"].append(row)

        return [key_metric_group]
    
    def read_analysis_txt(self, experiment_id: str) -> dict[str, Any]:
        """
        Read analysis artifact for a given experiment_id.

        Selection rule:
        - if badcase_deep_analysis is enabled: read root_case_analysis.jsonl
        - otherwise: read event_details_filtered.jsonl

        Args:
            experiment_id: Experiment ID.
        
        Returns:
            {"ok": True/False, "data": [...], "file_path": str, "error": "..."}
            - data is a list for downstream tool compatibility
        """
        try:
            file_name = "root_case_analysis.jsonl" if self.badcase_deep_analysis else "event_details_filtered.jsonl"
            file_path = os.path.join(self.output_dir, f"experiments_{experiment_id}", file_name)
            path = Path(file_path)
            raw_text = path.read_text(encoding="utf-8")

            data: list[Any] = []
            # Both branches currently target JSONL artifacts.
            for line in raw_text.splitlines():
                line = line.strip()
                if not line:
                    continue
                parsed = self._safe_json_loads(line)
                data.append(parsed if parsed is not None else line)

            return {
                "ok": True,
                "data": data,
                "source_file": file_name,
                "file_path": str(path)
            }
        except FileNotFoundError:
            return {"ok": False, "error": f"File not found: {file_path}"}
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    def read_memory_summaries(self) -> dict[str, Any]:
        """
        Read memory summaries from memory.jsonl file.
        Returns all memory summaries from previous iterations, sorted by iteration_id.
        The file is located at output_dir/memory.jsonl (same level as version.jsonl).

        Returns:
            {
                "ok": True/False,
                "summaries": list[dict] (sorted by iteration_id),
                "count": int,
                "file_path": str,
                "error": str | None
            }
        """
        try:
            # Build path: output_dir/memory.jsonl
            memory_path = Path(self.output_dir) / "memory.jsonl"

            if not memory_path.exists():
                return {
                    "ok": True,
                    "summaries": [],
                    "count": 0,
                    "file_path": str(memory_path),
                    "message": "memory.jsonl not found, no previous summaries"
                }

            # Read all lines from memory.jsonl
            summaries = []
            with open(memory_path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        try:
                            record = json.loads(line)
                            if isinstance(record, dict):
                                git_diff = record.get("git_diff")
                                if isinstance(git_diff, dict) and "diff" in git_diff:
                                    diff_display, diff_truncated, diff_original_len = _truncate_text(
                                        git_diff.get("diff"),
                                        max_chars=DEFAULT_MAX_DIFF_CHARS,
                                    )
                                    if diff_truncated:
                                        git_diff = dict(git_diff)
                                        git_diff["diff"] = diff_display
                                        git_diff["diff_truncated"] = True
                                        git_diff["diff_original_len"] = diff_original_len
                                        record = dict(record)
                                        record["git_diff"] = git_diff
                            summaries.append(record)
                        except json.JSONDecodeError:
                            continue

            # Sort by iteration_id (root < 1 < 2 < ...)
            def sort_key(record):
                iter_id = record.get("iteration_id", "root")
                if iter_id == "root":
                    return (0, "root")
                try:
                    return (1, int(iter_id))
                except (ValueError, TypeError):
                    return (2, str(iter_id))

            summaries.sort(key=sort_key)

            return {
                "ok": True,
                "summaries": summaries,
                "count": len(summaries),
                "file_path": str(memory_path)
            }
        except FileNotFoundError:
            return {"ok": False, "error": f"File not found: {memory_path}"}

    def read_priority_templates(self, priority_templates_path: str) -> dict[str, Any]:
        """
        Read priority templates markdown file generated by LiquidStructureAgent.
        This file contains 2-3 priority templates identified from bad case analysis
        that the optimizer should focus on modifying.

        Args:
            priority_templates_path: Path to the priority_templates.md file

        Returns:
            {
                "ok": True/False,
                "content": str (full markdown content),
                "priority_templates": list[dict] (parsed template info),
                "file_path": str,
                "error": str | None
            }
        """
        try:
            path = Path(priority_templates_path)

            if not path.exists():
                return {
                    "ok": False,
                    "error": f"Priority templates file not found: {priority_templates_path}"
                }

            content = path.read_text(encoding="utf-8")

            # Parse the markdown to extract template information
            priority_templates = []
            current_template = None

            for line in content.split('\n'):
                line = line.strip()

                # Match template header (e.g., "## 1. detailed_format.liquid")
                if line.startswith('## ') and '. ' in line:
                    if current_template:
                        priority_templates.append(current_template)
                    # Extract template name
                    template_name = line.split('. ', 1)[1].strip()
                    current_template = {"template_name": template_name, "path": "", "reason": "", "suggested_changes": []}

                elif current_template:
                    # Parse template details
                    if line.startswith('- **Path:**'):
                        current_template["path"] = line.replace('- **Path:**', '').strip()
                    elif line.startswith('- **Purpose:**'):
                        current_template["purpose"] = line.replace('- **Purpose:**', '').strip()
                    elif line.startswith('- **DO modify here:**'):
                        current_template["do_modify"] = line.replace('- **DO modify here:**', '').strip()
                    elif line.startswith("- **DON'T modify here:**"):
                        current_template["dont_modify"] = line.replace("- **DON'T modify here:**", '').strip()
                    elif line.startswith('- **Why prioritized:**'):
                        current_template["reason"] = line.replace('- **Why prioritized:**', '').strip()
                    elif line.startswith('  - ') and 'suggested_changes' in current_template:
                        current_template["suggested_changes"].append(line.replace('  - ', '').strip())

            # Add last template
            if current_template:
                priority_templates.append(current_template)

            return {
                "ok": True,
                "content": content,
                "priority_templates": priority_templates,
                "count": len(priority_templates),
                "file_path": str(path)
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}

    async def write_analysis_to_file(self, analysis_text: Any, experiment_id: str) -> dict[str, Any]:
        """
        Write analysis results to file
        
        Args:
            analysis_text: Text or structured object to persist as-is.
            experiment_id: Experiment ID. File will be saved to experiments_{experiment_id}/analysis.txt
        
        Returns:
            {"ok": True/False, "file_path": "...", "message": "..."}
        """
        try:
            if isinstance(analysis_text, str):
                content = analysis_text
            elif isinstance(analysis_text, (dict, list)):
                content = json.dumps(analysis_text, ensure_ascii=False)
            else:
                content = str(analysis_text)

            self.output_data_path = f"experiments_{experiment_id}/analysis.txt"
            await self.__save_output_data__(content, self.output_data_path, format="text")
            output_abs_path = Path(self.output_dir) / self.output_data_path

            return {
                "ok": True,
                "file_path": str(output_abs_path),
                "message": f"Successfully wrote analysis to {output_abs_path}"
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}

    async def get_bad_case_scores(self, experiment_id: str, file_name: str = "event_details_filtered.jsonl") -> dict[str, Any]:
        """
        Get bad case scores for a given experiment.
        Args:
            experiment_id: The ID of the experiment.
        Returns:
            A dictionary containing bad case scores.
        """
        data_path = Path(self.output_dir) / f"experiments_{experiment_id}" / file_name
        bad_case_scores = {}
        try:
            async with aiofiles.open(data_path, "r", encoding="utf-8") as f:
                async for line in f:
                    event = json.loads(line)
                    id = event.get("input", "")
                    scores = event.get("merged_scores", {})
                    judge_response = event.get("judge_response", "{}")
                    average_score = event.get("average_score", 0.0)
                    model_response = event.get("test_model_response", "")
                    bad_case_scores[id] = {
                        "scores": scores,
                        "average_score": average_score,
                        "judge_response": judge_response,
                        "model_response": model_response
                    }
        except Exception as e:
            return {"ok": False, "error": str(e)}
        return {"ok": True, "bad_case_scores": bad_case_scores}

    async def analyze_scores_for_each_dimension(self, previous_scores: dict[str, Any], current_scores: dict[str, Any]) -> dict[str, Any]:
        """
        Analyze scores for each dimension between two experiments.
        Args:
            previous_scores: The scores from the previous experiment.
            current_scores: The scores from the current experiment.
        Returns:
            A dictionary containing the analysis results.
        """
        analysis_results = {}
        metric_keys = self.key_metrics
        if not metric_keys:
            prev_keys = set(previous_scores.keys()) if isinstance(previous_scores, dict) else set()
            curr_keys = set(current_scores.keys()) if isinstance(current_scores, dict) else set()
            metric_keys = list(prev_keys | curr_keys)
        for dimension in metric_keys:
            prev_score = previous_scores.get(dimension, 0.0)
            curr_score = current_scores.get(dimension, 0.0)
            score_diff = curr_score - prev_score
            analysis_results[dimension] = {
                "previous_score": prev_score,
                "current_score": curr_score,
                "score_improve": score_diff,
                "improved": score_diff > 0
            }
        return {"ok": True, "analysis_results": analysis_results}

    async def compare_bad_cases_scores(
        self,
        previous_experiment_id: str,
        current_experiment_id: str,
        reverse: bool = False,
        key_map: dict[str, str] | None = None,
        case_detail_level: str = "summary",
    ) -> dict[str, Any]:
        """
        Compare bad cases scores between two experiments.
        Args:
            previous_experiment_id: The ID of the previous experiment.
            current_experiment_id: The ID of the current experiment.
            reverse: Whether to flip the comparison view.
            key_map: Mapping from canonical keys to output keys.
            case_detail_level: "summary" drops large case text fields; "full" keeps them (truncated).
        Returns:
            A dictionary containing the comparison results.
        """
        default_key_map = {
            "overall_scores_improve": "bad_case_overall_scores_improve",
            "dimension_scores_improve": "bad_case_dimension_scores_improve",
            "average_score_improve": "bad_case_average_score_improve",
            "previous_average_score": "bad_case_previous_average_score",
            "current_average_score": "bad_case_current_average_score",
            "dimension_compare": "bad_case_dimension_compare",
            "previous_judge_response": "previous_judge_response",
            "current_judge_response": "current_judge_response",
            "previous_model_response": "previous_model_response",
            "current_model_response": "current_model_response",
            "improved_cases": "improved_cases",
            "worse_cases": "worse_cases",
        }
        if key_map is None:
            key_map = default_key_map
        else:
            merged_key_map = dict(default_key_map)
            merged_key_map.update(key_map)
            key_map = merged_key_map
        detail_level = (case_detail_level or "summary").lower()
        if detail_level not in {"summary", "full"}:
            detail_level = "summary"
        case_text_fields = (
            "previous_judge_response",
            "current_judge_response",
            "previous_model_response",
            "current_model_response",
        )

        def _reverse_dimension_compare(dim_compare: dict[str, Any]) -> dict[str, Any]:
            analysis_results = dim_compare.get("analysis_results", {})
            if not isinstance(analysis_results, dict):
                return dim_compare
            reversed_results = {}
            for dim, dim_result in analysis_results.items():
                if isinstance(dim_result, dict):
                    updated = dict(dim_result)
                    if "score_improve" in updated:
                        updated["score_improve"] = -updated["score_improve"]
                    if "previous_score" in updated or "current_score" in updated:
                        prev_score = updated.get("previous_score")
                        curr_score = updated.get("current_score")
                        updated["previous_score"] = curr_score
                        updated["current_score"] = prev_score
                    if "improved" in updated:
                        if "score_improve" in updated:
                            updated["improved"] = updated["score_improve"] > 0
                        else:
                            updated["improved"] = not updated["improved"]
                    reversed_results[dim] = updated
                else:
                    reversed_results[dim] = dim_result
            updated_compare = dict(dim_compare)
            updated_compare["analysis_results"] = reversed_results
            return updated_compare

        def _format_case(case: dict[str, Any], *, reverse: bool) -> dict[str, Any]:
            case_data = dict(case)
            if reverse:
                case_data["average_score_improve"] = -case_data["average_score_improve"]
                case_data["previous_average_score"], case_data["current_average_score"] = (
                    case_data["current_average_score"],
                    case_data["previous_average_score"],
                )
                case_data["previous_judge_response"], case_data["current_judge_response"] = (
                    case_data["current_judge_response"],
                    case_data["previous_judge_response"],
                )
                case_data["previous_model_response"], case_data["current_model_response"] = (
                    case_data["current_model_response"],
                    case_data["previous_model_response"],
                )
                dim_compare = case_data.get("dimension_compare")
                if isinstance(dim_compare, dict):
                    case_data["dimension_compare"] = _reverse_dimension_compare(dim_compare)
            if detail_level == "full":
                case_data = _truncate_record_fields(case_data, case_text_fields)
            else:
                for field in case_text_fields:
                    case_data.pop(field, None)
            return {key_map.get(k, k): v for k, v in case_data.items()}

        def _format_result(result: dict[str, Any], *, reverse: bool) -> dict[str, Any]:
            improved = [_format_case(case, reverse=reverse) for case in result.get("improved_cases", [])]
            worse = [_format_case(case, reverse=reverse) for case in result.get("worse_cases", [])]
            if reverse:
                improved, worse = worse, improved
            overall_scores_improve = result.get("overall_scores_improve", 0.0)
            dimension_scores_improve = result.get("dimension_scores_improve", {})
            if reverse:
                overall_scores_improve = -overall_scores_improve
                if isinstance(dimension_scores_improve, dict):
                    dimension_scores_improve = {
                        dim: (-val if isinstance(val, (int, float)) else val)
                        for dim, val in dimension_scores_improve.items()
                    }
            return {
                key_map.get("overall_scores_improve", "overall_scores_improve"): overall_scores_improve,
                key_map.get("dimension_scores_improve", "dimension_scores_improve"): dimension_scores_improve,
                key_map.get("improved_cases", "improved_cases"): improved,
                key_map.get("worse_cases", "worse_cases"): worse,
                "case_detail_level": detail_level,
            }

        previous_scores = await self.get_bad_case_scores(previous_experiment_id, file_name="event_details_filtered.jsonl")
        current_scores = await self.get_bad_case_scores(current_experiment_id, file_name="event_details.jsonl")
        if not previous_scores.get("ok"):
            return {"ok": False, "error": f"Failed to get previous bad case scores: {previous_scores.get('error')}"}
        if not current_scores.get("ok"):
            return {"ok": False, "error": f"Failed to get current bad case scores: {current_scores.get('error')}"}
        
        # find all bad cases from previous experiment that are improved in current experiment
        improved_cases_heap = []
        worse_cases_heap = []
        total_score_change = 0.0
        total_score_count = 0
        metric_keys = list(self.key_metrics) if self.key_metrics else []
        dimension_changes_sum = {dim: 0.0 for dim in metric_keys}
        dimension_changes_count = {dim: 0 for dim in metric_keys}
        remain_item_count = self.bad_case_limit_per_iteration // 2
        heap_counter = 0
        for case_id, previous_score in previous_scores.get("bad_case_scores", {}).items():
            current_score = current_scores.get("bad_case_scores", {}).get(case_id)
            if not current_score:
                continue
            else:   
                current_average_score = current_score.get("average_score", 0) if current_score else 0
                previous_average_score = previous_score.get("average_score", 0)
                scores_compare = await self.analyze_scores_for_each_dimension(previous_score.get("scores", {}), current_score.get("scores", {}))
                score_delta = current_average_score - previous_average_score
                compare_result = {
                        "id": case_id,
                        "average_score_improve": score_delta,
                        "previous_average_score": previous_average_score,
                        "current_average_score": current_average_score,
                        "dimension_compare": scores_compare,
                        "previous_judge_response": previous_score.get("judge_response", ""),
                        "current_judge_response": current_score.get("judge_response", ""),
                        "previous_model_response": previous_score.get("model_response", ""),
                        "current_model_response": current_score.get("model_response", ""),
                    }
                total_score_change += score_delta
                total_score_count += 1
                analysis_results = scores_compare.get("analysis_results", {}) if isinstance(scores_compare, dict) else {}
                dims_to_update = metric_keys or list(analysis_results.keys())
                for dim in dims_to_update:
                    dim_change = analysis_results.get(dim, {}).get("score_improve", 0.0)
                    if dim not in dimension_changes_sum:
                        dimension_changes_sum[dim] = 0.0
                        dimension_changes_count[dim] = 0
                    dimension_changes_sum[dim] += dim_change
                    dimension_changes_count[dim] += 1

                if remain_item_count > 0:
                    heap_counter += 1
                    if current_score and current_average_score > previous_average_score:
                        if len(improved_cases_heap) < remain_item_count:
                            heappush(improved_cases_heap, (score_delta, heap_counter, compare_result))
                        elif score_delta > improved_cases_heap[0][0]:
                            heapreplace(improved_cases_heap, (score_delta, heap_counter, compare_result))
                    else:
                        neg_score = -score_delta
                        if len(worse_cases_heap) < remain_item_count:
                            heappush(worse_cases_heap, (neg_score, heap_counter, compare_result))
                        elif neg_score > worse_cases_heap[0][0]:
                            heapreplace(worse_cases_heap, (neg_score, heap_counter, compare_result))
        improved_cases = [item[2] for item in sorted(improved_cases_heap, key=lambda x: x[0], reverse=True)]
        worse_cases = [item[2] for item in sorted(worse_cases_heap, key=lambda x: x[0], reverse=True)]
        # Save comparison result to file
        raw_result = {
            "overall_scores_improve": (total_score_change / total_score_count) if total_score_count else 0.0,
            "dimension_scores_improve": {dim: (dimension_changes_sum[dim] / dimension_changes_count[dim]) if dimension_changes_count[dim] else 0.0 for dim in dimension_changes_sum},
            "improved_cases": improved_cases,
            "worse_cases": worse_cases,
        }
        comparison_result = _format_result(raw_result, reverse=reverse)
        logger.info(f"Comparison result between experiment {previous_experiment_id} and {current_experiment_id}: {comparison_result}")
        output_data_path = f"experiments_{current_experiment_id}/bad_case_comparison_with_{previous_experiment_id}.jsonl"
        logger.info(f"Saving bad case comparison result to {output_data_path}")
        await self.__save_output_data__([comparison_result], output_data_path)

        # Implement analysis logic here
        return {"ok": True, "message": f"Analyzed bad cases for experiment {previous_experiment_id} and {current_experiment_id}", "comparison_result": comparison_result}

    async def score_change_all_candidates(self, shared_state: Any, enable_compare_bad_cases_scores: bool = True) -> dict[str, Any]:
        """
        Compare overall_score and five metrics between all candidates in current node and parent node from shared_state.
        This reads directly from shared_state which is more up-to-date than version.jsonl.
        
        Args:
            shared_state: SharedState instance to read scores from
        
        Returns:
            Dict with 'ok', 'current_node', 'parent_node', 'candidates_score_changes', 'error'
        """
        state = await shared_state.read(None)
        
        # Get current iteration index
        iteration_idx = state.get("iteration_idx", "root")
        iterations = state.get("iteration", {})
        
        # Get current node info
        current_node_info = iterations.get(str(iteration_idx), {})
        current_node = current_node_info.get("current_node", iteration_idx)
        candidates = current_node_info.get("candidates", {})

        # Get analysis_source node info
        base_node = current_node_info.get("analysis_source_node")
        base_candidate = current_node_info.get("analysis_candidate")
        base_node_info = iterations.get(str(base_node), {}) if base_node is not None else {}
        if base_node is None or not base_node_info:
            return {
                "ok": True,
                "current_node": str(current_node),
                "analysis_source_node": None,
                "candidates_score_changes": {},
                "message": "This is the root node, no analysis_source_node to compare with"
            }
        base_candidate_snapshot = self._extract_candidate_snapshot(base_node_info, base_candidate)
        if not base_candidate_snapshot:
            return {
                "ok": False,
                "error": f"no analysis candidate info found for candidate {base_candidate} in node {base_node}"
            }
       
        final_results = []

        for candidate_id, candidate_info in candidates.items():

            res = await self.score_change_candidate_info(
                candidate_info,
                base_candidate_snapshot,
                enable_compare_bad_cases_scores,
            )
            res['current_candidate_id'] = candidate_id
            res['base_candidate_id'] = base_candidate
            final_results.append(res)

        return {
            "ok": True,
            "current_node": str(current_node),
            "base_node": str(base_node),
            "candidates_score_changes": final_results,
        }


    @staticmethod
    def _extract_candidate_snapshot(node_info: dict[str, Any], candidate_id: str | None) -> dict[str, Any]:
        """
        Extract candidate snapshot from node_info for the given candidate_id.
        If candidate_id is None, use selected_candidates from node_info.
        """
        candidates = node_info.get("candidates", {})
        return candidates.get(candidate_id, {}) if isinstance(candidate_id, str) else {}

    async def compare_changes_between_base_given_candidate(self, shared_state: any, candidate: str, enable_compare_bad_cases_scores: bool = False) -> dict[str, Any]:
        """
        Compare changes between the current snapshot and the base snapshot for a given candidate.
        """
        state = await shared_state.read(None)
        # Get current iteration index
        iteration_idx = state.get("iteration_idx", "root")
        iterations = state.get("iteration", {})
        # Get current node info
        current_node_info = iterations.get(str(iteration_idx), {})
        candidates = current_node_info.get("candidates", {})
        current_candidate_snapshot = candidates.get(candidate, {})
        if not current_candidate_snapshot:
            return {
                "ok": False,
                "error": f"no candidate info found for candidate {candidate} in current node {iteration_idx}"
            }
        base_node = current_node_info.get("analysis_source_node")
        if base_node is None:
            return {
                "ok": True,
                "current_node": str(iteration_idx),
                "analysis_source_node": None,
                "message": "This is the root node, no analysis_source_node to compare with"
            }
        base_candidate = current_node_info.get("analysis_candidate")
        base_node_info = iterations.get(str(base_node), {}) if base_node is not None else {}
        base_candidate_snapshot = self._extract_candidate_snapshot(base_node_info, base_candidate)  
        if not base_candidate_snapshot:
            return {
                "ok": False,
                "error": f"no candidate info found for candidate {base_candidate} in base node {base_node}"
            }   
        res = await self.score_change_candidate_info(
            current_candidate_snapshot, base_candidate_snapshot, enable_compare_bad_cases_scores
        )
        res['current_candidate_id'] = candidate
        res['base_candidate_id'] = base_candidate
        if 'ok' not in res:
            res['ok'] = True
        
        return res

    def compare_scores_between_candidates(
        self,
        current_snapshot: dict[str, Any],
        baseline_snapshot: dict[str, Any]
    ) -> tuple[float | None, float | None, float | None, dict[str, float]]:
        """
        Compare overall_score and key metrics between current_snapshot and baseline_snapshot.
        
        Args:
            current_snapshot: Current candidate snapshot
            baseline_snapshot: Baseline candidate snapshot
        
        Returns:
            overall_score_change: float | None
            metrics_changes: dict[str, float]
        """
        current_overall_score = current_snapshot.get("overall_score")
        base_overall_score = baseline_snapshot.get("overall_score") if baseline_snapshot else None

        current_metrics = current_snapshot.get("dimension_metrics", {})
        base_metrics = baseline_snapshot.get("dimension_metrics", {}) if baseline_snapshot else {}
        # Calculate overall_score change
        overall_score_change = None
        if current_overall_score is not None and base_overall_score is not None:
            overall_score_change = current_overall_score - base_overall_score

         # Calculate metrics changes
        metrics_changes = {}
        metric_keys = self.key_metrics if self.key_metrics else (set(current_metrics.keys()) | set(base_metrics.keys()))

        for metric_key in metric_keys:
            current_value = current_metrics.get(metric_key)
            base_value = base_metrics.get(metric_key)

            if current_value is not None and base_value is not None:
                change = current_value - base_value
                metrics_changes[metric_key] = {
                    "current": current_value,
                    "base": base_value,
                    "change": change,
                    "change_percent": (change / base_value * 100) if base_value != 0 else 0
                }
            elif current_value is not None:
                metrics_changes[metric_key] = {
                    "current": current_value,
                    "base": None,
                    "change": None,
                    "change_percent": None
                }
            elif base_value is not None:
                metrics_changes[metric_key] = {
                    "current": None,
                    "base": base_value,
                    "change": None,
                    "change_percent": None
                }
        return current_overall_score, base_overall_score, overall_score_change, metrics_changes

    async def score_change_candidate_info(self, candidate_info: dict[str, Any], base_candidate_info: dict[str, Any], enable_compare_bad_cases_scores: bool) -> dict[str, Any]:
        """
        Compare overall_score and five metrics between current node and parent node from shared_state.
        This reads directly from shared_state which is more up-to-date than version.jsonl.
        
        Args:
            shared_state: SharedState instance to read scores from
        
        Returns:
            Dict with 'ok', 'current_node', 'parent_node', 'overall_score_change', 'metrics_changes', 'error'
        """
        try:
            result = {}
            # Extract current and baseline snapshots
            current_overall_score, base_overall_score, overall_score_change, metrics_changes = self.compare_scores_between_candidates(
                candidate_info, base_candidate_info
            )
            current_experiment_id = candidate_info.get("experiment_id")
            analysis_experiment_id = base_candidate_info.get("experiment_id")
            result['overall_score_change'] = {
                    "current": current_overall_score,
                    "baseline": base_overall_score,
                    "change": overall_score_change,
                    "change_percent": (overall_score_change / base_overall_score * 100) if base_overall_score and base_overall_score != 0 else None
                } if overall_score_change is not None else None
            result ['overall_dimensions_metrics_changes'] = metrics_changes

            # Compare bad cases scores if enabled
            bad_cases_compares,current_bad_cases_compares = None, None
            if enable_compare_bad_cases_scores:
                if current_experiment_id and analysis_experiment_id:
                    bad_cases_compares, current_bad_cases_compares = await asyncio.gather(
                        self.compare_bad_cases_scores(analysis_experiment_id, current_experiment_id),
                        self.compare_bad_cases_scores(current_experiment_id, analysis_experiment_id, reverse=True),
                    )
                result['bad_cases_comparison'] = bad_cases_compares
                result['current_bad_cases_comparison'] = current_bad_cases_compares
            
            result['ok'] = True
            return result
        except Exception as e:
            return {"ok": False, "error": str(e)}
        
class PicassoClient:
    """Client for Picasso prompt file operations."""
    
    def __init__(
        self,
        repo_path: str | Path | None = None,
    ) -> None:
        """
        Initialize PicassoClient

        Args:
            repo_path: Picasso repository path (required).
        """
        if repo_path is None:
            raise ValueError("repo_path is required for PicassoClient")
        self.repo_path = Path(repo_path)

    async def read_prompt_template(self, file_paths_json: str) -> dict[str, Any]:
        """
        Read prompt files from Picasso repository.

        Args:
            file_paths_json: JSON string containing list of file paths to read.
                            Paths should be relative to repo (e.g., "Service.Chat/Inference/Prompts/Templates/...")
                            Format: ["path/to/file1.liquid", "path/to/file2.liquid"]

        Returns:
            {"ok": True/False, "files": {filename: content, ...}, "errors": [...]}
        """
        import json

        try:
            file_paths = json.loads(file_paths_json)
            if not isinstance(file_paths, list):
                return {"ok": False, "error": "file_paths_json must be a JSON array of file paths"}
        except json.JSONDecodeError as e:
            return {"ok": False, "error": f"Invalid JSON: {e}"}

        result: dict[str, Any] = {"ok": True, "files": {}}
        errors = []

        for rel_path in file_paths:
            # Get filename from path for the key
            filename = Path(rel_path).name

            # Check for filtered sidecar first (transparent condition-aware filtering)
            sidecar_path = self.repo_path / ".filtered" / rel_path
            abs_path = self.repo_path / rel_path

            try:
                if sidecar_path.exists():
                    with open(sidecar_path, "r", encoding="utf-8") as f:
                        content = f.read()
                else:
                    with open(abs_path, "r", encoding="utf-8") as f:
                        content = f.read()
                result["files"][filename] = {
                    "content": content,
                    "path": rel_path,
                }
            except FileNotFoundError:
                errors.append(f"{filename}: File not found - {abs_path}")
            except Exception as e:
                errors.append(f"{filename}: {str(e)}")

        if errors:
            result["ok"] = False
            result["errors"] = errors

        return result
    
    async def liquid_check(self, files_json: str) -> dict[str, Any]:
        """
        Validate syntax of multiple Liquid templates
        
        Args:
            files_json: JSON string containing mapping from filename to content
                        Example: '{"clarifying": "content...", "idiolect": "content..."}'
        
        Returns:
            {"ok": True/False, "results": {filename: {"ok": bool, "errors": [...]}}}
        """
        files = json.loads(files_json)
        if not isinstance(files, dict):
            return {"ok": False, "error": "files_json must decode to an object"}
        
        results = {}
        all_ok = True
        
        for name, content in files.items():
            errors = []
            
            # 1. Check basic Liquid tag pairing {% %}
            open_tags = content.count("{%")
            close_tags = content.count("%}")
            if open_tags != close_tags:
                errors.append(f"Mismatched Liquid tags: {open_tags} '{{%' vs {close_tags} '%}}'")
            
            # 2. Check variable tag pairing {{ }}
            open_vars = content.count("{{")
            close_vars = content.count("}}")
            if open_vars != close_vars:
                errors.append(f"Mismatched variable tags: {open_vars} '{{{{' vs {close_vars} '}}}}'")
            
            # 3. Check if common Liquid blocks are paired
            blocks = ["if", "for", "unless", "case", "capture", "comment", "raw"]
            for block in blocks:
                opens = len(re.findall(rf'\{{% *{block}\b', content))
                closes = len(re.findall(rf'\{{% *end{block} *%\}}', content))
                if opens != closes:
                    errors.append(f"Mismatched {{% {block} %}}/{{%end{block}%}}: {opens} opens vs {closes} closes")
            
            # 4. Check if elsif/else are within if block
            elsif_count = len(re.findall(r'\{% *elsif\b', content))
            else_count = len(re.findall(r'\{% *else *%\}', content))
            if_count = len(re.findall(r'\{% *if\b', content))
            if elsif_count > 0 or else_count > 0:
                if if_count == 0:
                    errors.append("Found elsif/else without any if block")
            
            results[name] = {
                "ok": len(errors) == 0,
                "errors": errors
            }
            if errors:
                all_ok = False
        
        return {"ok": all_ok, "results": results}
    
    async def make_diff(self, old_prompt: str, new_prompt: str) -> dict[str, Any]:
        """Generate unified diff"""
        from prompt_auto_lab.core.data_models import unified_diff
        return {"diff": unified_diff(old_prompt, new_prompt)}

    @staticmethod
    def _apply_filtered_diff_to_source(
        source: str,
        filtered_before: str,
        filtered_after: str,
        line_map: list[int],
    ) -> str:
        """Apply changes made to filtered content back onto the original source.

        Uses line_map to translate line positions from the filtered content
        to original source line numbers, handling nested conditionals correctly.

        Args:
            source: The original (unfiltered) file content.
            filtered_before: Filtered content before modifications.
            filtered_after: Filtered content after modifications.
            line_map: Mapping from filtered line index to 1-based original line number.

        Returns:
            Modified source with changes applied at the correct original positions.
        """
        import difflib

        before_lines = filtered_before.split("\n")
        after_lines = filtered_after.split("\n")
        source_lines = source.split("\n")

        # Use SequenceMatcher to find operations needed
        sm = difflib.SequenceMatcher(None, before_lines, after_lines)

        # Build list of (action, original_line_pos, new_content_lines) patches
        # Process in reverse order so line number offsets don't interfere
        patches: list[tuple[str, int, int, list[str]]] = []

        for op, i1, i2, j1, j2 in sm.get_opcodes():
            if op == "equal":
                continue
            elif op == "replace":
                # Map filtered line range [i1, i2) to original line numbers
                orig_start = line_map[i1] - 1 if i1 < len(line_map) else len(source_lines)
                orig_end = line_map[i2 - 1] if i2 - 1 < len(line_map) else len(source_lines)
                new_lines = after_lines[j1:j2]
                patches.append(("replace", orig_start, orig_end, new_lines))
            elif op == "insert":
                # Insert after the preceding filtered line's original position
                if i1 > 0 and i1 - 1 < len(line_map):
                    insert_after = line_map[i1 - 1]  # 1-based
                else:
                    # Inserting at the very beginning — use position before first mapped line
                    insert_after = line_map[0] - 1 if line_map else 0
                new_lines = after_lines[j1:j2]
                patches.append(("insert", insert_after, insert_after, new_lines))
            elif op == "delete":
                orig_start = line_map[i1] - 1 if i1 < len(line_map) else len(source_lines)
                orig_end = line_map[i2 - 1] if i2 - 1 < len(line_map) else len(source_lines)
                patches.append(("delete", orig_start, orig_end, []))

        # Apply patches in reverse order to preserve line numbers
        for action, start, end, new_lines in reversed(patches):
            if action == "replace":
                source_lines[start:end] = new_lines
            elif action == "insert":
                source_lines[start:start] = new_lines
            elif action == "delete":
                del source_lines[start:end]

        return "\n".join(source_lines)

    async def write_prompt_files(self, modifications_json: str) -> dict[str, Any]:
        """
        Apply modifications to prompt files based on modification instructions

        Args:
            modifications_json: JSON string containing modification instructions
                        Format (standard list):
                        {
                            "Service.Chat/Inference/Prompts/Templates/responding-common/idiolect.liquid": [
                                {"operation": "add_after", "search": "existing text", "content": "new content"},
                                {"operation": "replace", "search": "old text", "content": "new text"},
                                {"operation": "add_before", "search": "existing text", "content": "new content"}
                            ],
                            ...
                        }

                        Format (scoped dict — restricts modifications to a line range):
                        {
                            "Service.Chat/.../search_web_instruction.liquid": {
                                "line_range": [175, 249],
                                "modifications": [
                                    {"operation": "replace", "search": "old text", "content": "new text"}
                                ]
                            }
                        }

                        File paths should be relative to repo root (e.g., "Service.Chat/Inference/Prompts/Templates/...")

                        Operations:
                        - "add_after": Add content after the first occurrence of search text
                        - "add_before": Add content before the first occurrence of search text
                        - "replace": Replace the first occurrence of search text with new content

        Returns:
            {"ok": True/False, "written": [...], "errors": [...]}
        """
        try:
            # Try to parse the JSON string
            modifications = json.loads(modifications_json)
        except json.JSONDecodeError as e:
            # Try to fix common JSON issues
            cleaned_json = modifications_json.strip()

            # 1. Remove markdown code blocks if present
            if cleaned_json.startswith("```"):
                lines = cleaned_json.split("\n")
                if lines[0].startswith("```"):
                    lines = lines[1:]
                if lines and lines[-1].strip() == "```":
                    lines = lines[:-1]
                cleaned_json = "\n".join(lines).strip()

            # 2. Try to fix trailing extra characters (common issue: extra quotes/brackets)
            # If error is "Extra data", try to find the first valid JSON object
            if "Extra data" in str(e):
                # Try removing trailing characters one by one until we get valid JSON
                for i in range(1, min(10, len(cleaned_json))):  # Try removing up to 10 trailing chars
                    try:
                        test_json = cleaned_json[:-i]
                        modifications = json.loads(test_json)
                        # Success! Use this cleaned version
                        cleaned_json = test_json
                        break
                    except json.JSONDecodeError:
                        continue
                else:
                    # If removing trailing chars didn't work, try finding the first complete JSON object
                    # by looking for balanced braces
                    brace_count = 0
                    last_valid_pos = -1
                    for i, char in enumerate(cleaned_json):
                        if char == '{':
                            brace_count += 1
                        elif char == '}':
                            brace_count -= 1
                            if brace_count == 0:
                                last_valid_pos = i + 1
                                break

                    if last_valid_pos > 0:
                        try:
                            modifications = json.loads(cleaned_json[:last_valid_pos])
                            cleaned_json = cleaned_json[:last_valid_pos]
                        except json.JSONDecodeError:
                            pass

            # 3. Try parsing again with cleaned JSON
            try:
                if 'modifications' not in locals():
                    modifications = json.loads(cleaned_json)
            except json.JSONDecodeError as e2:
                # If still fails, provide detailed error information
                error_msg = f"Invalid JSON in modifications_json: {str(e2)}"
                # Show a preview of the problematic JSON
                error_pos = e2.pos if hasattr(e2, 'pos') else len(cleaned_json) // 2
                preview_start = max(0, error_pos - 100)
                preview_end = min(len(cleaned_json), error_pos + 100)
                preview = cleaned_json[preview_start:preview_end]
                error_msg += f"\nJSON preview around error position {error_pos}: ...{preview}..."
                error_msg += f"\nFull JSON length: {len(cleaned_json)} characters"
                error_msg += f"\nOriginal JSON length: {len(modifications_json)} characters"
                return {"ok": False, "error": error_msg}

        if not isinstance(modifications, dict):
            return {"ok": False, "error": "modifications_json must decode to an object"}

        result: dict[str, Any] = {"ok": True, "written": [], "errors": []}

        # Load filtered manifest for transparent write-back
        manifest_path = self.repo_path / ".filtered" / "_manifest.json"
        manifest: dict[str, Any] = {}
        if manifest_path.exists():
            try:
                manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
            except Exception:
                pass  # If manifest is broken, fall through to normal write

        for rel_path, mod_value in modifications.items():
            # Resolve relative path to absolute path
            file_path = self.repo_path / rel_path

            if not file_path.exists():
                result["errors"].append(f"{rel_path}: File not found at {file_path}")
                continue

            # Support both list format (legacy) and dict format (scoped)
            if isinstance(mod_value, dict):
                line_range = mod_value.get("line_range")
                mod_list = mod_value.get("modifications", [])
            elif isinstance(mod_value, list):
                line_range = None
                mod_list = mod_value
            else:
                result["errors"].append(f"{rel_path}: Modifications must be a list or dict")
                continue

            try:
                # Read current file content
                with open(file_path, "r", encoding="utf-8") as f:
                    current_content = f.read()

                if not isinstance(mod_list, list):
                    result["errors"].append(f"{rel_path}: Modifications must be a list")
                    continue

                # Transparent write-back: if file is in the filtered manifest
                # and no explicit line_range, apply modifications to the filtered
                # sidecar content, diff, and map changes back to the real file.
                manifest_entry = manifest.get(rel_path) if not line_range else None
                sidecar_path = self.repo_path / ".filtered" / rel_path
                if manifest_entry and sidecar_path.exists():
                    line_map = manifest_entry.get("line_map", [])
                    if not line_map:
                        result["errors"].append(f"{rel_path}: Filtered manifest has no line_map, cannot write back")
                        continue

                    filtered_before = sidecar_path.read_text(encoding="utf-8")

                    # Apply modifications to the filtered content
                    filtered_after = filtered_before
                    mod_errors = False
                    for mod in mod_list:
                        if not isinstance(mod, dict):
                            result["errors"].append(f"{rel_path}: Each modification must be a dict")
                            mod_errors = True
                            continue
                        operation = mod.get("operation")
                        search = mod.get("search")
                        content_val = mod.get("content")
                        if not operation or not search:
                            result["errors"].append(f"{rel_path}: Modification must have 'operation' and 'search' fields")
                            mod_errors = True
                            continue
                        if search not in filtered_after:
                            result["errors"].append(f"{rel_path}: Search text not found in filtered content: {search[:50]}...")
                            mod_errors = True
                            continue
                        if operation == "add_after":
                            actual = content_val
                            if content_val and content_val.startswith(search):
                                actual = content_val[len(search):]
                            if actual and "{% include " in actual:
                                result["errors"].append(f"{rel_path}: Cannot add {{% include %}} statements via add_after.")
                                mod_errors = True
                                continue
                            filtered_after = filtered_after.replace(search, search + actual, 1)
                        elif operation == "add_before":
                            actual = content_val
                            if content_val and content_val.endswith(search):
                                actual = content_val[:-len(search)]
                            if actual and "{% include " in actual:
                                result["errors"].append(f"{rel_path}: Cannot add {{% include %}} statements via add_before.")
                                mod_errors = True
                                continue
                            filtered_after = filtered_after.replace(search, actual + search, 1)
                        elif operation == "replace":
                            s_inc = "{% include " in search
                            c_inc = content_val and "{% include " in content_val
                            if s_inc and not c_inc:
                                result["errors"].append(f"{rel_path}: Cannot remove {{% include %}} statements via replace.")
                                mod_errors = True
                                continue
                            if not s_inc and c_inc:
                                result["errors"].append(f"{rel_path}: Cannot add {{% include %}} statements via replace.")
                                mod_errors = True
                                continue
                            filtered_after = filtered_after.replace(search, content_val, 1)
                        else:
                            result["errors"].append(f"{rel_path}: Unknown operation: {operation}")
                            mod_errors = True
                            continue

                    if mod_errors:
                        continue

                    # Diff filtered before/after and map changes back to real file
                    modified_content = self._apply_filtered_diff_to_source(
                        current_content, filtered_before, filtered_after, line_map
                    )

                    # # Also update the sidecar so subsequent reads see the new content
                    # sidecar_path.write_text(filtered_after, encoding="utf-8")
                else:
                    # Legacy: apply to full content (existing behavior)
                    modified_content = current_content

                    for mod in mod_list:
                        if not isinstance(mod, dict):
                            result["errors"].append(f"{rel_path}: Each modification must be a dict")
                            continue

                        operation = mod.get("operation")
                        search = mod.get("search")
                        content = mod.get("content")

                        if not operation or not search:
                            result["errors"].append(f"{rel_path}: Modification must have 'operation' and 'search' fields")
                            continue

                        # Find the search text in current content
                        if search not in modified_content:
                            result["errors"].append(f"{rel_path}: Search text not found: {search[:50]}...")
                            continue

                        # Apply operation
                        if operation == "add_after":
                            # Fix duplication: if content starts with search text, remove it
                            actual_content = content
                            if content and content.startswith(search):
                                actual_content = content[len(search):]
                            # Protect {% include %} statements: content cannot add include directives
                            if actual_content and "{% include " in actual_content:
                                result["errors"].append(f"{rel_path}: Cannot add {{%}} include {{%}} statements via add_after. Content contains include directive.")
                                continue
                            modified_content = modified_content.replace(search, search + actual_content, 1)
                        elif operation == "add_before":
                            # Fix duplication: if content ends with search text, remove it
                            actual_content = content
                            if content and content.endswith(search):
                                actual_content = content[:-len(search)]
                            # Protect {% include %} statements: content cannot add include directives
                            if actual_content and "{% include " in actual_content:
                                result["errors"].append(f"{rel_path}: Cannot add {{%}} include {{%}} statements via add_before. Content contains include directive.")
                                continue
                            modified_content = modified_content.replace(search, actual_content + search, 1)
                        elif operation == "replace":
                            # Protect {% include %} statements: cannot remove existing includes or add new ones
                            search_has_include = "{% include " in search
                            content_has_include = content and "{% include " in content
                            if search_has_include and not content_has_include:
                                result["errors"].append(f"{rel_path}: Cannot remove {{%}} include {{%}} statements via replace.")
                                continue
                            if not search_has_include and content_has_include:
                                result["errors"].append(f"{rel_path}: Cannot add {{%}} include {{%}} statements via replace.")
                                continue
                            modified_content = modified_content.replace(search, content, 1)
                        else:
                            result["errors"].append(f"{rel_path}: Unknown operation: {operation}. Valid: add_after, add_before, replace")
                            continue

                # Write modified content
                with open(file_path, "w", encoding="utf-8") as f:
                    f.write(modified_content)

                result["written"].append({
                    "path": rel_path,
                    "abs_path": str(file_path),
                    "size": len(modified_content),
                    "modifications_applied": len(mod_list)
                })
            except FileNotFoundError:
                result["errors"].append(f"{rel_path}: File not found - {file_path}")
            except Exception as e:
                result["errors"].append(f"{rel_path}: {str(e)}")

        if result["errors"]:
            result["ok"] = False

        return result

class CommenClients:
    """Common clients for various operations."""
    async def ping(self) -> dict[str, Any]:
        """Simple ping method to check connectivity."""
        return {"ok": True, "message": "pong"}

class GitClient:
    """Client for Git operations on Picasso repository."""

    def __init__(
        self,
        repo_path: str | Path,
        store_dir: str | Path | None = None,
    ) -> None:
        """
        Initialize GitClient

        Args:
            repo_path: Git repository path
            store_dir: Directory where version.jsonl is stored
        """
        self.repo_path = Path(repo_path)
        self.store_dir = Path(store_dir) if store_dir else None
    
    async def git_get_current_commit(self) -> dict[str, Any]:
        """
        Get current commit hash of the repository
        
        Returns:
            {"ok": True/False, "commit": str, "error": str}
        """
        
        try:
            # Get current commit hash
            result = await run_subprocess_async(
                ["git", "rev-parse", "HEAD"],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git rev-parse failed: {result.stderr}"}
            current_commit = result.stdout.strip()
            
            return {
                "ok": True,
                "commit": current_commit
            }
        
        except Exception as e:
            return {"ok": False, "error": str(e)}

    async def git_reset(self, commit: str, branch: str) -> dict[str, Any]:
        """
        Reset repository to specified commit and force push to remote
        
        Args:
            commit: Commit hash to reset to
            branch: Target branch name to reset on
        
        Returns:
            {"ok": True/False, "commit": str, "branch": str, "error": str}
        """
        
        try:
            # Validate branch is an autopromptlab managed branch
            if not branch.startswith("autopromptlab/ga/") and not branch.startswith("autopromptlab/session/"):
                return {"ok": False, "error": f"git reset refused: branch '{branch}' is not an autopromptlab managed branch."}

            # Get current branch
            result = await run_subprocess_async(
                ["git", "rev-parse", "--abbrev-ref", "HEAD"],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"Failed to get current branch: {result.stderr}"}

            current_branch = result.stdout.strip()

            # Check if current branch is the target branch, if not, checkout
            if current_branch != branch:
                result = await run_subprocess_async(
                    ["git", "checkout", branch],
                    cwd=str(self.repo_path),
                )
                if result.returncode != 0:
                    return {"ok": False, "error": f"git checkout to {branch} failed: {result.stderr}"}

            # Git reset to specified commit
            result = await run_subprocess_async(
                ["git", "reset", commit, "--hard"],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git reset failed: {result.stderr}"}

            # git push --force-with-lease to remote
            result = await run_subprocess_async(
                ["git", "push", "--force-with-lease"],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git push failed: {result.stderr}"}

            return {
                "ok": True,
                "commit": commit,
                "branch": branch,
                "message": f"Git reset successful on branch {branch}."
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    async def git_commit(self, branch: str, commit_message: str) -> dict[str, Any]:
        """
        Execute git checkout, add, commit in repository
        
        Args:
            branch: Target branch name
            commit_message: Commit message
        
        Returns:
            {"ok": True/False, "branch": str, "commit": str, "error": str}
        """
        
        try:
            # 1. Checkout to designated branch
            result = await run_subprocess_async(
                ["git", "checkout", "-B", branch],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git checkout failed: {result.stderr}"}

            # 2. Git add all changes in prompt directories
            # Use SPARSE_CHECKOUT_PATHS as single source of truth for staged paths
            result = await run_subprocess_async(
                ["git", "add", "-A"] + SPARSE_CHECKOUT_PATHS,
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git add failed: {result.stderr}"}

            # 3. Git commit
            result = await run_subprocess_async(
                ["git", "commit", "-m", commit_message],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                # If there are no changes, consider it successful
                if "nothing to commit" in result.stdout or "nothing to commit" in result.stderr:
                    return {"ok": True, "branch": branch, "commit": "no-changes", "message": "No changes to commit."}
                return {"ok": False, "error": f"git commit failed: {result.stderr}"}
            
            # Extract commit hash
            commit_hash = "unknown"
            if result.stdout:
                # Git commit format: [branch hash] message
                match = re.search(r'\[[\w/\-]+ ([a-f0-9]+)\]', result.stdout)
                if match:
                    commit_hash = match.group(1)
            
            return {
                "ok": True,
                "branch": branch,
                "commit": commit_hash,
                "message": "Git commit successful."
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    async def git_new_branch(self, branch: str, iteration: int) -> dict[str, Any]:
        """
        Create a new branch from the specified branch with auto-generated name.
        Branch name format: autopromptlab/ga/{commit_hash_short}_{iteration}
        Note: First checkout to the specified branch, then create a new branch from its HEAD.
        The iteration number allows multiple parallel branches from the same commit.
        
        Args:
            branch: The base branch to checkout before creating the new branch
            iteration: The iteration number (1-based) to include in branch name
        
        Returns:
            {"ok": True/False, "branch": str, "commit": str, "error": str}
        """
        
        try:
            # First checkout to the specified branch
            result = subprocess.run(
                ["git", "checkout", branch],
                cwd=str(self.repo_path),
                capture_output=True,
                text=True
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"Failed to checkout branch {branch}: {result.stderr}"}
            
            # Get current commit hash
            result = subprocess.run(
                ["git", "rev-parse", "HEAD"],
                cwd=str(self.repo_path),
                capture_output=True,
                text=True
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"Failed to get current commit: {result.stderr}"}
            
            current_commit = result.stdout.strip()
            commit_hash_short = current_commit[:8]  # Use first 8 characters of commit hash
            # Generate a random 6-character alphanumeric hash to prevent duplicates
            random_hash = ''.join(random.choices('0123456789abcdefghijklmnopqrstuvwxyz', k=6))
            branch_name = f"autopromptlab/ga/{commit_hash_short}_{iteration}_{random_hash}"
            
            # Create new branch from current HEAD (always create new, don't check if exists)
            result = subprocess.run(
                ["git", "checkout", "-b", branch_name],
                cwd=str(self.repo_path),
                capture_output=True,
                text=True
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git checkout -b failed: {result.stderr}"}
            
            # Push the new branch to remote with upstream tracking
            result = subprocess.run(
                ["git", "push", "-u", "origin", branch_name],
                cwd=str(self.repo_path),
                capture_output=True,
                text=True
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git push failed: {result.stderr}"}
            
            return {
                "ok": True,
                "branch": branch_name,
                "commit": current_commit,
                "message": f"Successfully created new branch {branch_name} from branch {branch} (commit {current_commit}) and pushed to remote."
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    async def git_delete_branch(self, branch_name: str) -> dict[str, Any]:
        """
        Delete a branch from both local and remote repository.
        
        Args:
            branch_name: Name of the branch to delete
        
        Returns:
            {"ok": True/False, "branch": str, "local_deleted": bool, "remote_deleted": bool, "error": str}
        """
        
        try:
            import re
            ga_pattern = r'^autopromptlab/ga/[^_]+_\d_[^_]+$'
            session_pattern = r'^autopromptlab/session/[a-zA-Z0-9_-]{1,20}$'
            if not re.match(ga_pattern, branch_name) and not re.match(session_pattern, branch_name):
                return {
                    "ok": True,
                    "branch": branch_name,
                    "local_deleted": False,
                    "remote_deleted": False,
                    "skipped": True,
                    "message": f"Skipped deletion of {branch_name} because it does not match expected formats (autopromptlab/ga/ or autopromptlab/session/)"
                }
            
            # Get current branch
            result = subprocess.run(
                ["git", "rev-parse", "--abbrev-ref", "HEAD"],
                cwd=str(self.repo_path),
                capture_output=True,
                text=True
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"Failed to get current branch: {result.stderr}"}
            
            current_branch = result.stdout.strip()
            
            # If currently on the branch to delete, skip deletion
            # The current branch should be the accepted child generation or the rollback parent generation, which should not be deleted
            if current_branch == branch_name:
                return {
                    "ok": True,
                    "branch": branch_name,
                    "local_deleted": False,
                    "remote_deleted": False,
                    "skipped": True,
                    "message": f"Skipped deletion of {branch_name} because it is the current branch (accepted child or rollback parent)"
                }
            
            local_deleted = False
            remote_deleted = False
            errors = []
            
            # Delete local branch
            result = subprocess.run(
                ["git", "branch", "-D", branch_name],
                cwd=str(self.repo_path),
                capture_output=True,
                text=True
            )
            if result.returncode == 0:
                local_deleted = True
            else:
                # Branch might not exist locally, check if it exists
                result_check = subprocess.run(
                    ["git", "show-ref", "--verify", "--quiet", f"refs/heads/{branch_name}"],
                    cwd=str(self.repo_path),
                    capture_output=True,
                    text=True
                )
                if result_check.returncode != 0:
                    # Branch doesn't exist locally, that's fine
                    pass
                else:
                    errors.append(f"Failed to delete local branch: {result.stderr}")
            
            # Delete remote branch
            result = subprocess.run(
                ["git", "push", "origin", "--delete", branch_name],
                cwd=str(self.repo_path),
                capture_output=True,
                text=True
            )
            if result.returncode == 0:
                remote_deleted = True
            else:
                # Branch might not exist on remote, check if it exists
                result_check = subprocess.run(
                    ["git", "ls-remote", "--heads", "origin", branch_name],
                    cwd=str(self.repo_path),
                    capture_output=True,
                    text=True
                )
                if result_check.returncode != 0 or not result_check.stdout.strip():
                    # Branch doesn't exist on remote, that's fine
                    pass
                else:
                    errors.append(f"Failed to delete remote branch: {result.stderr}")
            
            if errors:
                return {
                    "ok": False,
                    "branch": branch_name,
                    "local_deleted": local_deleted,
                    "remote_deleted": remote_deleted,
                    "error": "; ".join(errors)
                }
            
            return {
                "ok": True,
                "branch": branch_name,
                "local_deleted": local_deleted,
                "remote_deleted": remote_deleted,
                "message": f"Successfully deleted branch {branch_name} (local: {local_deleted}, remote: {remote_deleted})"
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    async def git_fetch(self, branch: str | None = None) -> dict[str, Any]:
        """
        Fetch from remote repository.
        
        Args:
            branch: Optional branch name to fetch. If None, fetches all.
        
        Returns:
            {"ok": True/False, "message": str, "error": str}
        """
        try:
            if branch:
                cmd = ["git", "fetch", "origin", branch]
            else:
                cmd = ["git", "fetch", "origin"]
            
            result = await run_subprocess_async(
                cmd,
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git fetch failed: {result.stderr}"}
            
            return {
                "ok": True,
                "message": "Successfully fetched from origin" + (f"/{branch}" if branch else "")
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    async def git_checkout(self, branch: str) -> dict[str, Any]:
        """
        Checkout to specified branch.
        
        Args:
            branch: Branch name to checkout
        
        Returns:
            {"ok": True/False, "branch": str, "error": str}
        """
        try:
            result = await run_subprocess_async(
                ["git", "checkout", branch],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git checkout failed: {result.stderr}"}
            
            return {
                "ok": True,
                "branch": branch,
                "message": f"Successfully checked out to {branch}"
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    async def git_pull(self, branch: str | None = None) -> dict[str, Any]:
        """
        Pull latest changes from remote repository.
        
        Args:
            branch: Optional branch name to pull. If None, pulls current branch.
        
        Returns:
            {"ok": True/False, "message": str, "error": str}
        """
        try:
            if branch:
                cmd = ["git", "pull", "origin", branch]
            else:
                cmd = ["git", "pull"]
            
            result = await run_subprocess_async(
                cmd,
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git pull failed: {result.stderr}"}
            
            return {
                "ok": True,
                "message": "Successfully pulled from origin" + (f"/{branch}" if branch else "")
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    async def git_create_and_push_branch(self, new_branch_name: str) -> dict[str, Any]:
        """
        Create a new branch from current HEAD and push to remote.
        
        Args:
            new_branch_name: Name for the new branch
        
        Returns:
            {"ok": True/False, "branch": str, "error": str}
        """
        try:
            # Create and checkout new branch
            result = await run_subprocess_async(
                ["git", "checkout", "-b", new_branch_name],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git checkout -b failed: {result.stderr}"}
            
            # Push new branch to remote with upstream tracking
            result = await run_subprocess_async(
                ["git", "push", "-u", "origin", new_branch_name],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git push failed: {result.stderr}"}
            
            return {
                "ok": True,
                "branch": new_branch_name,
                "message": f"Successfully created and pushed branch {new_branch_name}"
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    async def git_delete_remote_and_local_branch(self, branch_name: str, checkout_branch: str | None = None) -> dict[str, Any]:
        """
        Delete a branch from both local and remote repository.
        Optionally checkout to a different branch first.
        
        Args:
            branch_name: Name of the branch to delete
            checkout_branch: Optional branch to checkout before deletion
        
        Returns:
            {"ok": True/False, "branch": str, "local_deleted": bool, "remote_deleted": bool, "error": str}
        """
        try:
            # Checkout to another branch first if specified
            if checkout_branch:
                result = await run_subprocess_async(
                    ["git", "checkout", checkout_branch],
                    cwd=str(self.repo_path),
                )
                if result.returncode != 0:
                    return {"ok": False, "error": f"git checkout to {checkout_branch} failed: {result.stderr}"}
            
            local_deleted = False
            remote_deleted = False
            errors = []
            
            # Delete remote branch
            result = await run_subprocess_async(
                ["git", "push", "origin", "--delete", branch_name],
                cwd=str(self.repo_path),
            )
            if result.returncode == 0:
                remote_deleted = True
            else:
                errors.append(f"Failed to delete remote branch: {result.stderr}")
            
            # Delete local branch
            result = await run_subprocess_async(
                ["git", "branch", "-D", branch_name],
                cwd=str(self.repo_path),
            )
            if result.returncode == 0:
                local_deleted = True
            else:
                errors.append(f"Failed to delete local branch: {result.stderr}")
            
            if errors and not (local_deleted or remote_deleted):
                return {"ok": False, "error": "; ".join(errors)}
            
            return {
                "ok": True,
                "branch": branch_name,
                "local_deleted": local_deleted,
                "remote_deleted": remote_deleted,
                "message": f"Deleted branch {branch_name} (local: {local_deleted}, remote: {remote_deleted})"
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}

    @staticmethod
    def _normalize_branch_name(branch: str) -> str:
        """Collapse duplicate slashes and trim whitespace in branch names."""
        if not isinstance(branch, str):
            return branch
        return re.sub(r"/{2,}", "/", branch.strip())

    async def git_push(self, branch: str) -> dict[str, Any]:
        """
        Push specified branch to remote repository
        
        Args:
            branch: Branch name to push
        
        Returns:
            {"ok": True/False, "branch": str, "error": str}
        """
        
        try:
            normalized_branch = self._normalize_branch_name(branch)
            if not normalized_branch:
                return {"ok": False, "error": "git push failed: empty branch name"}
            if not normalized_branch.startswith("autopromptlab/ga/") and not normalized_branch.startswith("autopromptlab/session/"):
                return {"ok": False, "error": f"git push refused: branch '{normalized_branch}' is not an autopromptlab managed branch. Push is only allowed on autopromptlab/ga/ or autopromptlab/session/ branches."}
            if normalized_branch != branch:
                logger.warning(
                    "Normalizing branch name from %s to %s before git push",
                    branch,
                    normalized_branch,
                )
            # Git push
            result = await run_subprocess_async(
                ["git", "push", "-u", "origin", normalized_branch],
                cwd=str(self.repo_path),
            )
            if result.returncode != 0:
                return {"ok": False, "error": f"git push failed: {result.stderr}"}
            
            return {
                "ok": True,
                "branch": normalized_branch,
                "message": f"Successfully pushed to origin/{normalized_branch}"
            }
        except Exception as e:
            return {"ok": False, "error": str(e)}
    
    async def get_commit_ids(self, shared_state: Any = None) -> dict[str, Any]:
        """
        Get current node and parent node commit IDs from shared_state.
        - Reads directly from shared_state which is more up-to-date than version.jsonl.
        - Falls back to version.jsonl only if shared_state is not provided.
        
        Args:
            shared_state: SharedState instance to read commit IDs from (recommended)
        
        Returns:
            {
                "ok": True/False,
                "current_commit_id": str | None,
                "parent_commit_id": str | None,
                "current_node": str | None,
                "parent_node": str | None,
                "error": str | None
            }
        """
    
        # Prefer shared_state (most up-to-date)
        if shared_state is not None:
            try:
                state = await shared_state.read(None)
                
                # Get current iteration index
                iteration_idx = state.get("iteration_idx", "root")
                iterations = state.get("iteration", {})
                # Get analysis node info
                analysis_node = current_node_info.get("analysis_source_node")
                analysis_commit_id = current_node_info.get("analysis_commit_id")
                analysis_candidate_id = current_node_info.get("analysis_candidate_id")
                if analysis_commit_id is None:
                    return {"ok": False, "error": "analysis_commit_id not found in shared_state"}

                # Get current node info
                current_node_info = iterations.get(str(iteration_idx), {})
                current_candidates = current_node_info.get("candidates")
                res = [] 
                for candidate_id, candidate_info in current_candidates.items():
                    current_commit_id = candidate_info.get("commit_id")
                    if current_commit_id is None or analysis_commit_id is None:
                        continue
                    current_node = current_node_info.get("current_node", iteration_idx)
                    res.append({
                        "current_commit_id": current_commit_id,
                        "current_node": current_node,
                        "analysis_commit_id": analysis_commit_id,
                        "analysis_node": analysis_node,
                        "current_candidate": candidate_id,
                        "analysis_candidate": analysis_candidate_id,
                    })
                return res
            except Exception as e:
                # Fall back to version.jsonl if shared_state read fails
                return {"ok": False, "error": f"{str(e)}"}
        
    
    async def git_diff(self, commit_id_old: str, commit_id_new: str, max_chars: int | None = None) -> dict[str, Any]:
        """
        Get git diff between two commits.

        Args:
            commit_id_old: Previous commit ID
            commit_id_new: Current commit ID
            max_chars: Maximum characters for diff output. None uses DEFAULT_MAX_DIFF_CHARS.
                       Set to 0 or negative to disable truncation.

        Returns:
            {
                "ok": True/False,
                "diff": str (unified diff string),
                "files_changed": list[str],
                "error": str | None
            }
        """
        try:
            # Get diff between commits
            result = await run_subprocess_async(
                ["git", "diff", commit_id_old, commit_id_new],
                cwd=str(self.repo_path),
            )

            if result.returncode != 0:
                return {"ok": False, "error": f"git diff failed: {result.stderr}"}
            
            diff_output = result.stdout
            
            # Extract changed files from diff
            files_changed = []
            for line in diff_output.split("\n"):
                if line.startswith("diff --git"):
                    # Format: diff --git a/path/to/file b/path/to/file
                    parts = line.split()
                    if len(parts) >= 4:
                        file_path = parts[2].lstrip("a/")
                        if file_path not in files_changed:
                            files_changed.append(file_path)

            #save to experiment_id diff.txt
            # diff_path = self.store_dir / f"diff_{commit_id_old[:7]}_{commit_id_new[:7]}.txt"
            # with open(diff_path, "w", encoding="utf-8") as f:
            #     f.write(diff_output)

            diff_limit = max_chars if max_chars is not None else DEFAULT_MAX_DIFF_CHARS
            diff_display, diff_truncated, diff_original_len = _truncate_text(
                diff_output,
                max_chars=diff_limit,
            )
            response = {
                "ok": True,
                "diff": diff_display,
                "files_changed": files_changed,
            }
            if diff_truncated:
                response["diff_truncated"] = True
                response["diff_original_len"] = diff_original_len
            return response
        except Exception as e:
            return {"ok": False, "error": str(e)}

    async def git_diff_by_candidate(self, shared_state: Any, candidate_id: str) -> dict[str, Any]:
        """
        Get git diff between analysis commit and the commit for a given candidate in current node.
        """
        try:
            state = await shared_state.read(None)
            iteration_idx = state.get("iteration_idx", "root")
            iterations = state.get("iteration", {})
            current_node_info = iterations.get(str(iteration_idx), {})

            candidates = current_node_info.get("candidates", {})
            current_candidate_info = candidates.get(candidate_id, {})
            current_commit_id = current_candidate_info.get("commit_id")
            if current_commit_id is None:
                return {"ok": False, "error": f"commit_id not found for candidate {candidate_id} in node {iteration_idx}"}

            analysis_commit_id = current_node_info.get("analysis_commit_id")
            if analysis_commit_id is None:
                base_node = current_node_info.get("analysis_source_node")
                base_candidate = current_node_info.get("analysis_candidate")
                base_node_info = iterations.get(str(base_node), {}) if base_node is not None else {}
                base_candidate_info = base_node_info.get("candidates", {}).get(base_candidate, {})
                analysis_commit_id = base_candidate_info.get("commit_id")

            if analysis_commit_id is None:
                return {"ok": False, "error": "analysis_commit_id not found for base node/candidate"}
            if analysis_commit_id == current_commit_id:
                return {"ok": True, "error": "current commit equals analysis commit; nothing to diff"}

            diff_result = await self.git_diff(analysis_commit_id, current_commit_id)
            diff_result["analysis_commit_id"] = analysis_commit_id
            diff_result["current_commit_id"] = current_commit_id
            diff_result["current_candidate_id"] = candidate_id
            return diff_result
        except Exception as e:
            return {"ok": False, "error": str(e)}

    async def write_memory_summary(self, summary_json: str, shared_state: Any, git_diff: dict[str, Any]) -> dict[str, Any]:
        """
        Write memory summary to memory.jsonl file and store file path in shared_state.
        The memory.jsonl file is located at store_dir/memory.jsonl (same level as version.jsonl).
        Only the file path is stored in shared_state (like analysis_results_path), not the full content.
        
        Args:
            summary_json: JSON string containing memory summary with:
                - effectiveness: 
                - change_summary: str (what changed)
                - impact_analysis: str (how it affected scores)
                - lessons_learned: list[str] (key takeaways)
                - recommendations: list[str] (suggestions for future)
            shared_state: SharedState instance to read current iteration and write file path
        
        Returns:
            Dict with 'success', 'iteration_id', 'file_path', and 'summary'
        """
        try:
            
            summary = json.loads(summary_json) if isinstance(summary_json, str) else summary_json
            
            # Validate summary structure
            required_fields = ["effectiveness", "change_summary", "impact_analysis"]
            for field in required_fields:
                if field not in summary:
                    return {"success": False, "error": f"Missing required field: {field}"}
            
            # Get current state to determine iteration
            state = await shared_state.read(None)
            iteration_id = state.get("iteration_idx", "root")
            
            # Add metadata to summary
            memory_record = {
                "iteration_id": iteration_id,
                "git_diff": git_diff,
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                **summary
            }
            
            # Append to memory.jsonl file first
            if self.store_dir is None:
                return {
                    "success": False,
                    "error": "store_dir not set in GitClient, cannot write to memory.jsonl"
                }
            
            memory_path = self.store_dir / "memory.jsonl"
            memory_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Append to file (one JSON object per line)
            line = json.dumps(memory_record, ensure_ascii=False) + "\n"
            async with aiofiles.open(memory_path, "a", encoding="utf-8") as f:
                await f.write(line)
            
            # Store only file path in shared_state (like analysis_results_path)
            # patch_data = {
            #     "iteration": {
            #         str(iteration_id): {
            #             "memory_summary_path": str(memory_path)
            #         }
            #     }
            # }
            # await shared_state.patch(patch_data)
            
            return {
                "success": True,
                "iteration_id": iteration_id,
                "file_path": str(memory_path),
                "summary": summary
            }
        except json.JSONDecodeError as e:
            return {"success": False, "error": f"Invalid JSON: {str(e)}"}
        except Exception as e:
            return {"success": False, "error": str(e)}


class GitAgent(Protocol):
    async def push_prompt(self, *, branch: str, prompt_text: str) -> str: ...


class CodeExecutor(Protocol):
    async def liquid_check(self, prompt_text: str) -> dict[str, Any]: ...


class PythonCodeExecutor(Protocol):
    async def run_python(self, *, code: str, input: dict[str, Any]) -> dict[str, Any]: ...




def _load_dataset(dataset_path: str) -> list[dict[str, Any]]:
    path = Path(dataset_path)
    if not path.is_absolute():
        path = (Path.cwd() / path).resolve()
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except FileNotFoundError:
        return []
    if isinstance(data, list):
        return data
    if isinstance(data, dict) and isinstance(data.get("samples"), list):
        return list(data["samples"])
    raise ValueError(f"Unsupported dataset format: {type(data).__name__}")


def _extract_message(record: dict[str, Any]) -> str:
    if isinstance(record.get("message"), str) and record["message"]:
        return record["message"]
    try:
        turns = record["input"]["turns"]
        if isinstance(turns, list) and turns and isinstance(turns[0].get("message"), str):
            return turns[0]["message"]
    except Exception:
        pass
    return ""


def _extract_response(record: dict[str, Any]) -> str:
    if isinstance(record.get("picasso_ai_response"), str) and record["picasso_ai_response"]:
        return record["picasso_ai_response"]
    out = record.get("output")
    if isinstance(out, list) and out:
        first = out[0]
        if isinstance(first, dict) and isinstance(first.get("picasso_ai_response"), str):
            return first["picasso_ai_response"]
    return ""


def _compute_overall(samples: list[dict[str, Any]]) -> float:
    scores: list[float] = []
    for s in samples:
        v = s.get("score")
        if isinstance(v, (int, float)):
            scores.append(float(v))
        elif isinstance(s.get("scores"), dict):
            vals = [float(x) for x in s["scores"].values() if isinstance(x, (int, float))]
            if vals:
                scores.append(sum(vals) / len(vals))
    return (sum(scores) / len(scores)) if scores else 0.0


@dataclass
class InMemoryEvalBackend:
    jobs: dict[str, dict[str, Any]] | None = None
    results: dict[str, dict[str, Any]] | None = None
    pushed_prompts: list[dict[str, Any]] | None = None
    _dataset_cache: dict[str, list[dict[str, Any]]] | None = None

    def __post_init__(self) -> None:
        self.jobs = {} if self.jobs is None else self.jobs
        self.results = {} if self.results is None else self.results
        self.pushed_prompts = [] if self.pushed_prompts is None else self.pushed_prompts
        self._dataset_cache = {} if self._dataset_cache is None else self._dataset_cache

    def load_dataset(self, dataset_path: str) -> list[dict[str, Any]]:
        key = str(Path(dataset_path).resolve()) if dataset_path else ""
        if key in self._dataset_cache:  # type: ignore[operator]
            return self._dataset_cache[key]  # type: ignore[index]
        records = _load_dataset(dataset_path)
        self._dataset_cache[key] = records  # type: ignore[index]
        return records

    def build_metrics(self, *, dataset_path: str) -> dict[str, Any]:
        data = self.load_dataset(dataset_path)
        samples: list[dict[str, Any]] = []
        for i, r in enumerate(data):
            sample_id = r.get("id") if isinstance(r.get("id"), str) else f"sample{i+1}"
            score_dict = r.get("scores", {})
            if isinstance(score_dict, dict):
                vals = [float(x) for x in score_dict.values() if isinstance(x, (int, float))]
                score = (sum(vals) / len(vals)) if vals else 0.0
            else:
                score = float(r.get("score", 0.0) or 0.0)
                score_dict = {}
            samples.append(
                {
                    "id": sample_id,
                    "message": _extract_message(r),
                    "picasso_ai_response": _extract_response(r),
                    "scores": score_dict,
                    "score": float(score),
                }
            )
        overall = _compute_overall(samples)
        return {"overall_score": float(overall), "samples": samples}


@dataclass
class StubBraintrustClient:
    backend: InMemoryEvalBackend | None = None

    async def fetch_eval_results(self, result_id: str) -> dict[str, Any]:
        if self.backend is None:
            return {"result_id": result_id, "overall_score": 0.0, "samples": []}
        return self.backend.results.get(
            result_id, {"result_id": result_id, "overall_score": 0.0, "samples": []}
        )  # type: ignore[union-attr]


@dataclass
class StubGitAgent:
    backend: InMemoryEvalBackend | None = None

    async def push_prompt(self, *, branch: str, prompt_text: str) -> str:
        commit = hashlib.sha256((branch + "\n" + prompt_text).encode("utf-8")).hexdigest()[:12]
        if self.backend is not None:
            self.backend.pushed_prompts.append({"branch": branch, "prompt": prompt_text, "commit": commit})  # type: ignore[union-attr]
        return commit


@dataclass
class SimpleLiquidCodeExecutor:
    async def liquid_check(self, prompt_text: str) -> dict[str, Any]:
        opens = prompt_text.count("{{") + prompt_text.count("{%-") + prompt_text.count("{%")
        closes = prompt_text.count("}}") + prompt_text.count("-%}") + prompt_text.count("%}")
        ok = opens == closes
        return {"ok": ok, "details": "balanced" if ok else f"unbalanced: opens={opens} closes={closes}"}


@dataclass
class LocalPythonCodeExecutor:
    async def run_python(self, *, code: str, input: dict[str, Any]) -> dict[str, Any]:
        if not isinstance(code, str) or not code.strip():
            raise ValueError("code must be a non-empty string")
        if not isinstance(input, dict):
            raise ValueError("input must be a dict")

        safe_builtins: dict[str, Any] = {
            "len": len,
            "sum": sum,
            "min": min,
            "max": max,
            "sorted": sorted,
            "float": float,
            "int": int,
            "str": str,
            "dict": dict,
            "list": list,
            "enumerate": enumerate,
            "range": range,
        }
        glb: dict[str, Any] = {"__builtins__": safe_builtins}
        loc: dict[str, Any] = {"input": input, "output": None}

        exec(code, glb, loc)

        out = loc.get("output")
        if out is None and callable(loc.get("compute")):
            out = loc["compute"](input)

        if not isinstance(out, dict):
            raise ValueError("Python executor must return a dict via `output` or `compute()`")
        return out

async def test():
    from dotenv import load_dotenv
    from prompt_auto_lab.core.auth_token import get_azure_token_provider
    from prompt_auto_lab.core.model_clients import build_azure_model_client
    from prompt_auto_lab.clients.picasso.analysis_agent import AnalysisAgentClient

    load_dotenv()
    braintrust_org = os.getenv("BRAINTRUST_ORG", "non-pme")
    braintrust_org = "bt-azure"
    api_key = os.getenv("BRAINTRUST_AZURE_API_KEY", "") if braintrust_org == "bt-azure" else os.getenv("BRAINTRUST_API_KEY", "")
    root_span_id = os.getenv("ROOT_SPAN_ID")
    # pme_key = "sk-wRdeMEPX4VcX1bIQye13vEAbnkzdfY0lchXZC911rSkVXVIa"
    token_provider = get_azure_token_provider("https://cognitiveservices.azure.com/.default")
    model_client = build_azure_model_client(
        token_provider=token_provider,
        deployment="gpt-5.1-chat",
        endpoint="https://cnmoderation-evaluation.openai.azure.com/",
        api_version="2024-08-01-preview",
    )

    analysis_agent_client = AnalysisAgentClient()
    metrics_config =  [{"name": "10gram_overlap" , "is_llm_judged": False, "optimization_direction": "lower_better"}, {"name": "3gram_overlap" , "is_llm_judged": False, "optimization_direction": "lower_better"},]
    
    # experiment_ids = ["25111ddf-5ae2-4562-a18d-856507d112a3","444f6a08-d293-48af-be8b-24b0c02fba83"]
    experiment_ids = {
        # "443da735-07d5-4f84-8eaf-634e91c72f67": {
        #     "eval": "edge_side_pane_comprehensiveness_freechat",
        #     "metrics": []
        # },
        # "1c12bdd3-45e4-4b2d-950d-464b13129489": {
        #     "eval": "groundedness_v2",
        #     "metrics": []
        # },
        # "df70ef47-4bd0-4ce6-acf8-80e9d167e590": {
        #     "eval": "ipcopyright_v2",
        #     "metrics": []
        # },
        # "cb30d3b9-935b-447e-b178-a0994e80e1e5": {
        #     "eval": "citations",
        #     "metrics": []
        # },
        # "7ce617e8-97bc-43bb-bafe-c2bcf85d4f7b": {
        #     "eval": "instruction_following",
        #     "metrics": []
        # },
        "b3c1c43e-1742-4f84-8f4c-9d28979064d8": {
            "eval": "windows_filesagent_planner_positive",
            "metrics": [
                {
                    "name": "windows_actions_positive_prompt_reasoning_score",
                    "is_llm_judged": True,
                    "description": "",
                    "measurement_purpose": "",
                    "optimization_direction": "higher_better",
                },
                {
                    "name": "windows_aicommands_prompt_reasoning_score",
                    "is_llm_judged": True,
                    "description": "",
                    "measurement_purpose": "",
                    "optimization_direction": "higher_better",
                },
            ],
        },
        "a94fc307-c778-4ba2-9135-1a8d7d56a3df": {"eval": "factuality", "metrics": []},
        "020660f2-73a5-4d6f-86f5-8507b3c564e6": {
            "eval": "personality_evals",
            "metrics": [],
        },
        "b8b8e60e-2ddb-4704-a479-68386cfe6f30": {
            "eval": "model_hallucination",
            "metrics": [],
        },
        "b83fbc73-035d-4682-a1d2-60aaff504016": {
            "eval": "completeness_v2",
            "metrics": [],
        },
        "fad03857-698c-4011-97ce-032ca9efde42": {"eval": "refusal_v2", "metrics": []},
        "9c06d887-91e0-4576-877e-7c989ba0498c": {"eval": "backstory", "metrics": []},
        "a4cb164f-1b3c-4445-9daa-f22efe3cdf45": {"eval": "freshness", "metrics": []},
        "5b7cebbe-9332-4c77-b1fe-1025248826cc": {"eval": "verbosity", "metrics": []},
        "69d02ec8-037e-484e-99d3-fb767e00f723": {
            "eval": "overpromising",
            "metrics": [],
        },
        "fafd9812-49aa-4671-8298-9225a3e8a091": {
            "eval": "formatting_v2",
            "metrics": [],
        },
        "908c70d0-64b4-40c9-bb3d-7616a41b7fc3": {
            "eval": "td_response_quality",
            "metrics": [],
        },
    }


    # experiment_id = "155da1bc-5fba-42d9-b4e0-a1c84e277836"
    for experiment_id,v in experiment_ids.items():
        eval_name =v['eval']
        exp_metrics = v.get('metrics', [])
        braintrustclient = BraintrustClient(
            backend=InMemoryEvalBackend(),
            api_key=api_key,
            local_debug=False,
            model_client=model_client,
            braintrust_org=braintrust_org,
            eval_name=eval_name,
            metrics_config=exp_metrics,
            analysis_agent_client=analysis_agent_client,
        )
        logger.info("BraintrustClient initialized: %s", braintrustclient)
        logger.info("Fetching eval results for experiment_id: %s", experiment_id)
        logger.info("processing experiment id %s", experiment_id)
        res = await braintrustclient.fetch_summary_results(experiment_id)
        # res = await braintrustclient.fetch_event_details(experiment_id)
        # # print("Eval results:", res)
        # # Process the evaluation results as needed
        # res = await braintrustclient.fetch_summary_results(experiment_id)
        # print("Summary results:", res)

        # low_case_text = await braintrustclient.get_low_case_details(
        #     experiment_id,
        #     include_analysis=False,
        #     include_conversation=False,
        # )
        # print("Low case details size:", len(low_case_text) if isinstance(low_case_text, str) else type(low_case_text))

        # if root_span_id:
        #     conversation = await braintrustclient.get_low_case_details(
        #         experiment_id,
        #         include_analysis=False,
        #         include_conversation=True,
        #         root_span_id=root_span_id,
        #     )
        #     logger.info("Conversation fetch: %s", conversation.get("ok"))
        # else:
        #     logger.info("ROOT_SPAN_ID not set; skipping conversation fetch")

        if braintrustclient.analysis_agent_client is None:
            logger.info("analysis_agent_client not set; skipping analysis generation")
        else:
            analysis_text = await braintrustclient.get_low_case_details(
                experiment_id,
                include_analysis=True,
                include_conversation=False,
            )
            logger.info(
                "Analysis content size: %s",
                len(analysis_text) if isinstance(analysis_text, str) else type(analysis_text),
            )

    # res = await braintrustclient.compare_bad_cases_scores(experiment_ids[0], experiment_ids[1])
    # res = await braintrustclient.compare_bad_cases_scores(experiment_ids[1], experiment_ids[0], reverse=True)
    logger.info("%s", res)

if __name__ == "__main__":
    import asyncio
    asyncio.run(test())
