"""Main SDK client class."""

from __future__ import annotations

import asyncio
import json
import logging
import shutil
import subprocess
import sys
import time
import uuid
import weakref
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Awaitable, ClassVar

from prompt_auto_lab.core.liquid_entries import normalize_liquid_prompt_entries
from prompt_auto_lab.core.metrics_config import normalize_metrics_config as _shared_normalize_metrics_config
from prompt_auto_lab.sdk.config import AutoPromptLabConfig
from prompt_auto_lab.sdk.config_file import load_create_kwargs_from_file, build_create_kwargs
from prompt_auto_lab.sdk.result import OptimizationResult

logger = logging.getLogger(__name__)


def _new_session_id() -> str:
    return str(int(time.time())) + "_" + str(uuid.uuid4())

def _normalize_metrics_config(
    metrics_config: list[dict[str, Any]] | None,
) -> list[dict[str, Any]] | None:
    """Normalize metrics_config to UI-compatible schema."""
    return _shared_normalize_metrics_config(metrics_config)


def _raise_unsupported_copilot_for_sdk() -> None:
    raise ValueError(
        "model_client_type=Copilot is not supported by AutoPromptLab SDK. "
        "Use model_client_type=Local with a full OpenAI-compatible local_endpoint instead."
    )


def _validate_create_config(
    *,
    braintrust_api_key: str,
    model_client_type: str,
    azure_endpoint: str | None,
    model_deployment: str | None,
    azure_api_version: str | None,
    papyrus_endpoint: str | None,
    papyrus_model: str | None,
    papyrus_quota_id: str | None,
    papyrus_scope: str | None,
    local_endpoint: str,
    local_model: str,
) -> str:
    """Validate create() configuration and return normalized model client type."""
    if not isinstance(braintrust_api_key, str) or not braintrust_api_key.strip():
        raise ValueError("braintrust_api_key is required and cannot be empty.")

    normalized_type = str(model_client_type or "").strip().lower()
    if normalized_type == "copilot":
        _raise_unsupported_copilot_for_sdk()

    allowed_types = {"azure", "papyrus", "local"}
    if normalized_type not in allowed_types:
        raise ValueError(
            "model_client_type must be one of: Azure, Papyrus, Local."
        )

    if normalized_type == "azure":
        if not isinstance(azure_endpoint, str) or not azure_endpoint.strip():
            raise ValueError("azure_endpoint is required when model_client_type=Azure.")
        if not isinstance(model_deployment, str) or not model_deployment.strip():
            raise ValueError("model_deployment is required when model_client_type=Azure.")
        if not isinstance(azure_api_version, str) or not azure_api_version.strip():
            raise ValueError("azure_api_version is required when model_client_type=Azure.")
    elif normalized_type == "papyrus":
        if not isinstance(papyrus_endpoint, str) or not papyrus_endpoint.strip():
            raise ValueError("papyrus_endpoint is required when model_client_type=Papyrus.")
        if not isinstance(papyrus_model, str) or not papyrus_model.strip():
            raise ValueError("papyrus_model is required when model_client_type=Papyrus.")
        if not isinstance(papyrus_scope, str) or not papyrus_scope.strip():
            raise ValueError("papyrus_scope is required when model_client_type=Papyrus.")
    elif normalized_type == "local":
        if not isinstance(local_endpoint, str) or not local_endpoint.strip():
            raise ValueError("local_endpoint is required when model_client_type=Local.")
        if not isinstance(local_model, str) or not local_model.strip():
            raise ValueError("local_model is required when model_client_type=Local.")

    return normalized_type.title()


class LabStatus(Enum):
    """Status of the AutoPromptLab instance."""
    IDLE = "idle"
    RUNNING = "running"
    STOPPING = "stopping"
    COMPLETED = "completed"
    FAILED = "failed"


class AutoPromptLab:
    """
    Main SDK class for running prompt optimization workflows.
    
    This class provides a high-level interface to the AutoPromptLab
    optimization framework.
    
    Two main use cases:
    1. Start a new optimization experiment
    2. Resume from a checkpoint
    """
    _active_instances: ClassVar[weakref.WeakSet] = weakref.WeakSet()
    
    def __init__(
        self,
        config: AutoPromptLabConfig | None = None,
        *,
        init_model_client: bool = True,
    ) -> None:
        """
        Initialize AutoPromptLab with configuration.
        
        Args:
            config: SDK configuration.
        """
        self._config = config or AutoPromptLabConfig()
        self._orchestrator = None  # Lazy initialization
        self._event_callback: Callable[[str, str, str], Awaitable[None]] | None = None
        self._running_task: asyncio.Task | None = None
        self._status: LabStatus = LabStatus.IDLE
        self._last_result: OptimizationResult | None = None
        self._current_run_id: str | None = None
        self._session_id: str | None = None
        self._user_id: str | None = None
        self._user_branch: str | None = None
        self._session_branch: str | None = None
        self._resolved_config_path: Path | None = None
        self._config_bundle: dict[str, Any] = {}
        self._create_kwargs_snapshot: dict[str, Any] = {}
        self._model_client: Any = None
        self._model_think_client: Any = None
        self._clients: Any = None  # ModelClients dataclass (main/think/analysis)
        self._analysis_agent_client: Any = None
        self._run_root_file_handler: logging.Handler | None = None
        self._run_noisy_file_handler: logging.Handler | None = None
        self._run_noisy_logger_propagate: dict[str, bool] = {}
        
        # Ensure session directory exists
        self._config.session_dir.mkdir(parents=True, exist_ok=True)
        AutoPromptLab._active_instances.add(self)
        if init_model_client:
            self._init_model_clients()

    def _get_session_picasso_repo_path(self) -> Path:
        """Return the effective Picasso working copy path for this SDK instance."""
        from prompt_auto_lab.core.picasso_workspace import get_session_picasso_path

        user_id = str(self._user_id or "sdk-user").strip() or "sdk-user"
        session_id = str(self._session_id or "")
        if not session_id:
            raise ValueError("session_id is required to resolve Picasso working path")
        return get_session_picasso_path(user_id, session_id)
    
    @property
    def status(self) -> LabStatus:
        """Current status of the optimization run."""
        return self._status
    
    @property
    def is_running(self) -> bool:
        """Whether an optimization is currently running."""
        return self._status == LabStatus.RUNNING
    
    @property
    def last_result(self) -> OptimizationResult | None:
        """Result from the last optimization run."""
        return self._last_result

    @property
    def session_id(self) -> str | None:
        """Resolved session id when created from session-aware entrypoints."""
        return self._session_id

    @property
    def user_id(self) -> str | None:
        """Resolved user id when created from session-aware entrypoints."""
        return self._user_id

    @property
    def session_dir(self) -> Path:
        """Effective session directory."""
        return self._config.session_dir

    @property
    def picasso_repo_path(self) -> str:
        """Resolved session-scoped Picasso working repository path."""
        return str(self._get_session_picasso_repo_path())

    @property
    def resolved_config_path(self) -> Path | None:
        """Resolved JSON config path used by create_from_file()."""
        return self._resolved_config_path

    @property
    def config_bundle(self) -> dict[str, Any]:
        """Grouped parsed config sections loaded from JSON config file."""
        return dict(self._config_bundle)

    @property
    def create_kwargs_snapshot(self) -> dict[str, Any]:
        """Final kwargs snapshot that was used to construct this SDK instance."""
        return dict(self._create_kwargs_snapshot)
    
    @classmethod
    def create(
        cls,
        *,
        # ========================================
        # Required - Local Environment (1 param)
        # ========================================
        braintrust_api_key: str,
        
        # ========================================
        # Required - Experiment Config (3 params)
        # ========================================
        dataset_name: str,
        picasso_branch: str,
        liquid_prompt_entries: list[dict],
        
        # ========================================
        # Optional - Evaluation settings
        # ========================================
        eval_project_name: str = "",
        dataset_project_name: str = "",
        max_dataset_rows: int = 300,
        dataset_sampling_method: str = "all",
        additional_cmd_args: str = "",
        conversation_mode: str = "Smart",
        environment: str = "global",
        model: str | None = None,
        picasso_env: str = "runner",
        task_id_prefix: str = "",
        
        # ========================================
        # Optional - Feature flags
        # ========================================
        add_feature_flags: str = "",
        baseline_feature_flags: str = "",
        picasso_version: str = "Baseline",
        run_leaderboard_after_optimization: bool = False,

        # ========================================
        # Optional - Optimization defaults
        # ========================================
        max_rounds: int = 10,
        termination_rule: dict[str, Any] | None = None,
        rollback_strategy: str | None = None,
        baseline_experiment_id: str | None = None,
        baseline_experiment_link: str | None = None,
        selection_strategy: str | None = None,
        pareto_max_selected: int | None = None,
        
        # ========================================
        # Optional - Braintrust
        # ========================================
        braintrust_org: str = "non-pme",
        eval_name: str = "default",
        metrics_config: list[dict[str, Any]] | None = None,
        dimension_threshold: float = 0.5,
        bad_case_limit_per_iteration: int = 10,
        badcase_deep_analysis: bool = False,
        
        # ========================================
        # Optional - Model settings
        # ========================================
        model_client_type: str = "Azure",
        azure_endpoint: str | None = None,
        model_deployment: str | None = None,
        azure_api_version: str | None = None,
        papyrus_endpoint: str | None = None,
        papyrus_model: str | None = None,
        papyrus_quota_id: str | None = None,
        papyrus_scope: str | None = None,
        local_endpoint: str = "http://localhost:23333/api/openai/v1",
        local_model: str = "",
        use_structured_output: bool = True,

        # ========================================
        # Optional - Storage
        # ========================================
        session_dir: str | Path | None = None,
        session_id: str | None = None,
        user_id: str = "sdk-user",
        
        # ========================================
        # Optional - Debug
        # ========================================
        local_debug: bool = False,
        enable_run_logging: bool = False,
    ) -> "AutoPromptLab":
        """
        Create an AutoPromptLab instance with model/eval configuration.

        Required Parameters (Local Environment):
            braintrust_api_key: Braintrust API key.
        
        Required Parameters (Experiment Config):
            dataset_name: Dataset name for evaluation.
            picasso_branch: Picasso repository branch for prompts.
            liquid_prompt_entries: Liquid template entries (list of dicts with name/path).
        
        Optional Parameters:
            eval_project_name: Braintrust project name.
            dataset_project_name: Braintrust dataset project name.
            max_dataset_rows: Maximum dataset rows for evaluation (default: 300).
            conversation_mode: Conversation mode (default: "Smart").
            environment: Evaluation environment, e.g. "global" or "china".
            model: Aipa eval model parameter.
            add_feature_flags: Additional feature flags.
            baseline_feature_flags: Baseline feature flags.
            picasso_version: Picasso version identifier (default: "Baseline").
            run_leaderboard_after_optimization: Whether to submit leaderboard job
                after optimization (default: False).
            max_rounds: Default max rounds for optimize()/resume() when not passed.
            termination_rule: Default termination rule for optimize().
            rollback_strategy: Default rollback strategy for optimize().
            baseline_experiment_id: Default baseline experiment id for optimize().
            baseline_experiment_link: Default baseline experiment link for optimize().
            selection_strategy: Selection strategy: 'greedy' or 'gepa' (default: server default).
            pareto_max_selected: Max candidates from Pareto front (1-5), used when selection_strategy='gepa'.
            braintrust_org: Braintrust organization (default: "non-pme").
            eval_name: Evaluation run name (default: "default").
            metrics_config: Metric configuration list. Each item supports:
                name, is_llm_judged, description, measurement_purpose,
                optimization_direction (higher_better/lower_better).
                If omitted, SDK derives defaults from dataset mapping when available.
            dimension_threshold: Threshold for low-score dimension filtering.
            bad_case_limit_per_iteration: Max bad cases analyzed per iteration.
            model_client_type: Model client type (Azure/Papyrus/Local).
            azure_endpoint: Azure OpenAI endpoint URL.
            model_deployment: Azure OpenAI model deployment name.
            papyrus_endpoint: Papyrus endpoint URL.
            papyrus_model: Papyrus model name.
            local_endpoint: Local OpenAI-compatible base URL.
            local_model: Local model ID.
            use_structured_output: Whether to use structured model output.
            session_dir: Root directory for session artifacts (default: ".autopromptlab/sessions").
            session_id: Optional session id. If provided, SDK creates/uses
                user-scoped session directory as session_dir.
            user_id: User id used with session_id (default: "sdk-user").
            local_debug: Enable local debug mode.
            enable_run_logging: Whether to write per-run log files under
                session_dir/logs (default: False, console-only).
        
        Returns:
            Configured AutoPromptLab instance.
        
        Example:
            lab = AutoPromptLab.create(
                # Required - Local Environment
                braintrust_api_key="sk-xxx",
                
                # Required - Experiment Config
                dataset_name="picasso_cn_chat_0923",
                picasso_branch="autopromptlab2/main",
                liquid_prompt_entries=[
                    {"name": "BasePrompt", "path": "BasePrompt/Prompt.liquid"}
                ],
                
                # Optional
                session_dir="./sessions",
            )
        """
        model_client_type_normalized = _validate_create_config(
            braintrust_api_key=braintrust_api_key,
            model_client_type=model_client_type,
            azure_endpoint=azure_endpoint,
            model_deployment=model_deployment,
            azure_api_version=azure_api_version,
            papyrus_endpoint=papyrus_endpoint,
            papyrus_model=papyrus_model,
            papyrus_quota_id=papyrus_quota_id,
            papyrus_scope=papyrus_scope,
            local_endpoint=local_endpoint,
            local_model=local_model,
        )

        normalized_metrics_config = _normalize_metrics_config(metrics_config)
        if normalized_metrics_config is None:
            from prompt_auto_lab.core.orchestrator_builder import get_default_metrics_config_from_dataset
            normalized_metrics_config = get_default_metrics_config_from_dataset(dataset_name)
        normalized_liquid_entries = normalize_liquid_prompt_entries(liquid_prompt_entries)
        if not normalized_liquid_entries:
            raise ValueError("liquid_prompt_entries must contain at least one valid {path} entry.")

        effective_session_dir = Path(session_dir) if session_dir else Path(".autopromptlab/sessions")
        resolved_session_id_raw = str(session_id).strip() if session_id is not None else ""
        resolved_session_id: str | None = resolved_session_id_raw or None
        normalized_user_id: str | None = None
        if resolved_session_id is not None:
            from prompt_auto_lab.core.user_store import ensure_user_record, get_user_session_dir

            normalized_user_id = str(user_id or "sdk-user").strip() or "sdk-user"
            ensure_user_record(normalized_user_id)
            user_session_dir = get_user_session_dir(normalized_user_id, resolved_session_id)
            user_session_dir.mkdir(parents=True, exist_ok=True)
            effective_session_dir = user_session_dir

        config = AutoPromptLabConfig(
            session_dir=effective_session_dir,
            braintrust_api_key=braintrust_api_key,
            braintrust_org=braintrust_org,
            eval_name=eval_name,
            metrics_config=normalized_metrics_config,
            dimension_threshold=dimension_threshold,
            bad_case_limit_per_iteration=bad_case_limit_per_iteration,
            badcase_deep_analysis=badcase_deep_analysis,
            eval_project_name=eval_project_name,
            dataset_project_name=dataset_project_name,
            dataset_name=dataset_name,
            max_dataset_rows=max_dataset_rows,
            dataset_sampling_method=dataset_sampling_method,
            additional_cmd_args=additional_cmd_args,
            conversation_mode=conversation_mode,
            environment=environment,
            model=model,
            picasso_env=picasso_env,
            task_id_prefix=task_id_prefix,
            picasso_branch=picasso_branch,
            liquid_prompt_entries=normalized_liquid_entries,
            add_feature_flags=add_feature_flags,
            baseline_feature_flags=baseline_feature_flags,
            picasso_version=picasso_version,
            run_leaderboard_after_optimization=run_leaderboard_after_optimization,
            max_rounds=int(max_rounds),
            termination_rule=termination_rule,
            rollback_strategy=rollback_strategy,
            baseline_experiment_id=baseline_experiment_id,
            baseline_experiment_link=baseline_experiment_link,
            selection_strategy=selection_strategy,
            pareto_max_selected=pareto_max_selected,
            model_client_type=model_client_type_normalized,
            azure_endpoint=azure_endpoint,
            model_deployment=model_deployment,
            azure_api_version=azure_api_version,
            papyrus_endpoint=papyrus_endpoint,
            papyrus_model=papyrus_model,
            papyrus_quota_id=papyrus_quota_id,
            papyrus_scope=papyrus_scope,
            local_endpoint=local_endpoint,
            local_model=local_model,
            use_structured_output=use_structured_output,
            local_debug=local_debug,
            enable_run_logging=enable_run_logging,
        )
        lab = cls(config=config)
        lab._session_id = resolved_session_id
        lab._user_id = normalized_user_id
        return lab

    @classmethod
    def create_from_file(
        cls,
        config_path: str | Path,
        *,
        session_id: str | None = None,
        load_env_secrets: bool = True,
        **overrides: Any,
    ) -> "AutoPromptLab":
        """Create SDK instance from JSON config file.

        The JSON file uses grouped sections such as
        ``picasso_config``/``aipa_config``/
        ``braintrust_config``/``model_config``/``runtime_config``/
        ``optimize_config``. ``overrides`` are applied last.

        Session handling is built-in:
        - user_id is fixed to ``sdk-user``
        - if ``session_id`` is missing, SDK auto-generates one
        - session_dir resolves to the user-scoped session directory
        """
        resolved_config_path = Path(config_path).expanduser().resolve()
        config_bundle = load_create_kwargs_from_file(
            resolved_config_path,
            load_env_secrets=load_env_secrets,
        )
        resolved_session_id_raw = str(session_id).strip() if session_id else ""
        resolved_session_id = resolved_session_id_raw or _new_session_id()
        normalized_user_id = "sdk-user"

        create_kwargs = build_create_kwargs(config_bundle)
        create_kwargs.update({k: v for k, v in overrides.items() if v is not None})
        create_kwargs["session_id"] = resolved_session_id
        create_kwargs["user_id"] = normalized_user_id
        lab = cls.create(**create_kwargs)
        lab._resolved_config_path = resolved_config_path
        lab._config_bundle = config_bundle
        lab._create_kwargs_snapshot = dict(create_kwargs)
        return lab
    
    def _build_orchestrator(
        self,
        run_dir: Path | None = None,
        user_proxy_mode_override: str | None = None,
        user_input_callback_override: Callable[..., str | Awaitable[str]] | None = None,
    ):
        """Build orchestrator lazily when needed.
        
        Args:
            run_dir: Directory for this run's outputs. If None, uses session_dir directly.
            user_proxy_mode_override: Optional UserProxy mode override for this run.
            user_input_callback_override: Optional user input callback override for this run.
        """
        # Note: We rebuild orchestrator for each run to use the correct run_dir
        from prompt_auto_lab.core.orchestrator_builder import (
            OrchestratorBuilder,
            OrchestratorConfig,
            ModelClientConfig,
            EvalClientConfig,
            BraintrustConfig,
        )
        
        cfg = self._config
        
        # Use run_dir if provided, otherwise use session_dir as fallback
        output_dir = run_dir / "outputs" if run_dir else cfg.session_dir

        model_client_type_normalized = str(cfg.model_client_type or "Azure").strip() or "Azure"
        if model_client_type_normalized.lower() == "copilot":
            _raise_unsupported_copilot_for_sdk()
        btql_model_endpoint = None
        btql_model_name = None
        if model_client_type_normalized.lower() == "papyrus":
            btql_model_endpoint = cfg.papyrus_endpoint
            btql_model_name = cfg.papyrus_model
        elif model_client_type_normalized.lower() == "local":
            btql_model_endpoint = cfg.local_endpoint
            btql_model_name = cfg.local_model or None
        
        # Build configuration
        orch_config = OrchestratorConfig(
            model=ModelClientConfig(
                client_type=model_client_type_normalized,
                azure_endpoint=cfg.azure_endpoint,
                model_deployment=cfg.model_deployment,
                azure_api_version=cfg.azure_api_version,
                papyrus_endpoint=cfg.papyrus_endpoint,
                papyrus_model=cfg.papyrus_model,
                papyrus_quota_id=cfg.papyrus_quota_id,
                papyrus_scope=cfg.papyrus_scope,
                local_endpoint=cfg.local_endpoint,
                local_model=cfg.local_model,
                use_structured_output=cfg.use_structured_output,
            ),
            prebuilt_clients=self._clients,
            prebuilt_model_client=self._model_client,
            prebuilt_model_think_client=self._model_think_client,
            eval=EvalClientConfig(
                output_dir=str(output_dir),
                eval_name=cfg.eval_name,
                dataset_name=cfg.dataset_name,
                local_debug=cfg.local_debug,
                max_dataset_rows=cfg.max_dataset_rows,
                dataset_sampling_method=cfg.dataset_sampling_method,
                additional_cmd_args=cfg.additional_cmd_args,
                conversation_mode=cfg.conversation_mode,
                environment=cfg.environment,
                add_feature_flags=cfg.add_feature_flags,
                baseline_feature_flags=cfg.baseline_feature_flags,
                picasso_version=cfg.picasso_version,
                model=cfg.model,
                picasso_env=cfg.picasso_env,
                task_id_prefix=cfg.task_id_prefix,
                eval_project_name=cfg.eval_project_name,
                dataset_project_name=cfg.dataset_project_name,
                braintrust_org=cfg.braintrust_org,
            ),
            braintrust=BraintrustConfig(
                api_key=cfg.braintrust_api_key or "",
                metrics_config=cfg.metrics_config,
                dimension_threshold=cfg.dimension_threshold,
                bad_case_limit_per_iteration=cfg.bad_case_limit_per_iteration,
                output_dir=str(output_dir),
                eval_name=cfg.eval_name,
                braintrust_org=cfg.braintrust_org,
                environment=cfg.environment,
                local_debug=cfg.local_debug,
                model_client_type=model_client_type_normalized.lower(),
                model_endpoint=btql_model_endpoint,
                model_name=btql_model_name,
                badcase_deep_analysis=cfg.badcase_deep_analysis,
            ),
            picasso_repo_path=self._get_session_picasso_repo_path(),
            session_dir=cfg.session_dir,
            user_proxy_mode=user_proxy_mode_override or cfg.user_proxy_mode or ("callback" if cfg.user_input_callback else "console"),
            user_input_callback=user_input_callback_override if user_input_callback_override is not None else cfg.user_input_callback,
            user_proxy_timeout=cfg.user_proxy_timeout,
            prefer_native_user_proxy=cfg.prefer_native_user_proxy,
            analysis_agent_client=self._analysis_agent_client,
        )
        
        self._orchestrator = OrchestratorBuilder.build(
            orch_config,
            event_callback=self._event_callback,
        )
        
        return self._orchestrator

    def _init_model_clients(self) -> None:
        """Initialize model clients at SDK construction time (fail-fast)."""
        from prompt_auto_lab.core.orchestrator_builder import ModelClientBuilder, ModelClientConfig

        if str(self._config.model_client_type or "").strip().lower() == "copilot":
            _raise_unsupported_copilot_for_sdk()

        model_cfg = ModelClientConfig(
            client_type=self._config.model_client_type,
            azure_endpoint=self._config.azure_endpoint,
            model_deployment=self._config.model_deployment,
            azure_api_version=self._config.azure_api_version,
            papyrus_endpoint=self._config.papyrus_endpoint,
            papyrus_model=self._config.papyrus_model,
            papyrus_quota_id=self._config.papyrus_quota_id,
            papyrus_scope=self._config.papyrus_scope,
            local_endpoint=self._config.local_endpoint,
            local_model=self._config.local_model,
            use_structured_output=self._config.use_structured_output,
        )
        clients = ModelClientBuilder.build(model_cfg)
        self._clients = clients
        # Backward-compat aliases
        self._model_client = clients.main
        self._model_think_client = clients.think

        # Build AnalysisAgentClient once; reuse across runs
        try:
            from prompt_auto_lab.clients.picasso.analysis_agent import AnalysisAgentClient
            self._analysis_agent_client = AnalysisAgentClient(
                store_dir=str(self._config.session_dir),
                metrics_config=self._config.metrics_config,
            )
        except Exception as exc:
            logger.warning("Could not initialize AnalysisAgentClient: %s", exc)
    
    @classmethod
    def from_orchestrator(
        cls,
        orchestrator,
        config: AutoPromptLabConfig | None = None,
    ) -> "AutoPromptLab":
        """
        Create an AutoPromptLab instance from an existing Orchestrator.
        
        Use this when you need full control over the Orchestrator configuration.
        
        Args:
            orchestrator: Pre-configured Orchestrator instance.
            config: Optional SDK configuration.
        
        Returns:
            AutoPromptLab instance wrapping the orchestrator.
        """
        instance = cls(config=config, init_model_client=False)
        instance._orchestrator = orchestrator
        return instance
    
    @staticmethod
    def _ensure_console_logging(*, level: int) -> None:
        root_logger = logging.getLogger()
        root_logger.setLevel(level)

        has_console_handler = any(
            isinstance(handler, logging.StreamHandler) and not isinstance(handler, logging.FileHandler)
            for handler in root_logger.handlers
        )
        if has_console_handler:
            return

        stream = sys.stdout
        if hasattr(stream, "reconfigure"):
            try:
                stream.reconfigure(encoding="utf-8", errors="replace")
            except Exception:
                pass

        formatter = logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")
        stream_handler = logging.StreamHandler(stream)
        stream_handler.setFormatter(formatter)
        root_logger.addHandler(stream_handler)

    def _teardown_run_logging(self) -> None:
        root_handler = self._run_root_file_handler
        if root_handler is not None:
            root_logger = logging.getLogger()
            if root_handler in root_logger.handlers:
                root_logger.removeHandler(root_handler)
            root_handler.close()
            self._run_root_file_handler = None

        noisy_handler = self._run_noisy_file_handler
        if noisy_handler is not None:
            for logger_name, previous_propagate in list(self._run_noisy_logger_propagate.items()):
                noisy_logger = logging.getLogger(logger_name)
                if noisy_handler in noisy_logger.handlers:
                    noisy_logger.removeHandler(noisy_handler)
                noisy_logger.propagate = previous_propagate
            noisy_handler.close()
            self._run_noisy_file_handler = None
            self._run_noisy_logger_propagate = {}

    @classmethod
    def _teardown_run_logging_for_path(cls, target: Path) -> None:
        resolved_target = target.resolve()
        for instance in list(cls._active_instances):
            session_dir = instance._config.session_dir.resolve()
            if session_dir == resolved_target or resolved_target in session_dir.parents:
                instance._teardown_run_logging()

    def _setup_run_logging(self, logs_dir: Path) -> tuple[Path | None, Path | None]:
        """
        Setup logging for the current run.

        When ``enable_run_logging`` is True, creates two log files:
        - run_{timestamp}.log - Main log (INFO level)
        - run_{timestamp}.verbose.log - Verbose log (includes httpx, autogen)
        Otherwise, logging remains console-only.
        
        Args:
            logs_dir: Directory for log files when file logging is enabled.

        Returns:
            Tuple of (log_path, verbose_log_path), or (None, None) when
            file logging is disabled.
        """
        from datetime import datetime, timezone
        from prompt_auto_lab.core.logging_setup import resolve_log_level
        
        log_level = resolve_log_level(default=self._config.log_level)
        self._ensure_console_logging(level=log_level)
        self._suppress_noisy_debug_logs()
        self._teardown_run_logging()

        if not self._config.enable_run_logging:
            return None, None

        logs_dir.mkdir(parents=True, exist_ok=True)
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
        log_path = logs_dir / f"run_{timestamp}.log"
        noisy_log_path = logs_dir / f"run_{timestamp}.verbose.log"
        formatter = logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")

        root_handler = logging.FileHandler(log_path, encoding="utf-8")
        root_handler.setFormatter(formatter)
        logging.getLogger().addHandler(root_handler)
        self._run_root_file_handler = root_handler

        noisy_handler = logging.FileHandler(noisy_log_path, encoding="utf-8")
        noisy_handler.setFormatter(formatter)
        noisy_logger_names = ("httpx", "autogen_core.events")
        self._run_noisy_logger_propagate = {}
        for logger_name in noisy_logger_names:
            noisy_logger = logging.getLogger(logger_name)
            self._run_noisy_logger_propagate[logger_name] = noisy_logger.propagate
            noisy_logger.addHandler(noisy_handler)
            noisy_logger.propagate = False
        self._run_noisy_file_handler = noisy_handler

        return log_path, noisy_log_path

    def _apply_optimize_defaults(self, kwargs: dict[str, Any]) -> None:
        """Fill optimize() kwargs from config when the caller did not override them."""
        cfg = self._config
        if "termination_rule" not in kwargs and cfg.termination_rule is not None:
            kwargs["termination_rule"] = cfg.termination_rule
        if "rollback_strategy" not in kwargs and cfg.rollback_strategy:
            kwargs["rollback_strategy"] = cfg.rollback_strategy
        if "baseline_experiment_id" not in kwargs and cfg.baseline_experiment_id:
            kwargs["baseline_experiment_id"] = cfg.baseline_experiment_id
        if "baseline_experiment_link" not in kwargs and cfg.baseline_experiment_link:
            kwargs["baseline_experiment_link"] = cfg.baseline_experiment_link
        if "selection_strategy" not in kwargs and cfg.selection_strategy:
            kwargs["selection_strategy"] = cfg.selection_strategy
        if "pareto_max_selected" not in kwargs and cfg.pareto_max_selected is not None:
            kwargs["pareto_max_selected"] = cfg.pareto_max_selected

    def _prepare_run_paths(self, *, run_id: str) -> tuple[Path, Path, Path]:
        """Create per-session directories needed by optimize()/resume()."""
        run_dir = self._config.session_dir
        checkpoint_dir = run_dir / "checkpoints"
        outputs_dir = run_dir / "outputs"
        logs_dir = run_dir / "logs"

        dirs_to_create = [checkpoint_dir, outputs_dir]
        if self._config.enable_run_logging:
            dirs_to_create.append(logs_dir)

        for path in dirs_to_create:
            path.mkdir(parents=True, exist_ok=True)

        self._current_run_id = run_id
        return run_dir, checkpoint_dir, logs_dir

    @staticmethod
    def _log_run_start(run_id: str, log_path: Path | None) -> None:
        """Emit a consistent run-start log message."""
        if log_path is not None:
            logger.info("Run %s started, logs: %s", run_id, log_path)
            return
        logger.info("Run %s started", run_id)

    @staticmethod
    def _suppress_noisy_debug_logs() -> None:
        """Suppress low-level debug logs from third-party SDKs."""
        noisy_logger_names = (
            "openai",
            "openai._base_client",
            "httpx",
            "httpcore",
            "autogen_core",
            "autogen_core.events",
            "autogen_agentchat",
        )
        for name in noisy_logger_names:
            logging.getLogger(name).setLevel(logging.WARNING)

    async def _ensure_picasso_repo_ready(self) -> None:
        """Ensure the effective Picasso working repo exists and is usable."""
        repo_path = self._get_session_picasso_repo_path()

        if repo_path.exists() and (repo_path / ".git").exists():
            self._populate_default_picasso_commit_id(repo_path)
            return

        if repo_path.exists():
            if repo_path.is_dir():
                shutil.rmtree(repo_path)
            else:
                repo_path.unlink()

        from prompt_auto_lab.core.picasso_workspace import copy_picasso_repo

        logger.info(
            "Session Picasso repo not ready, cloning to %s (branch: %s)",
            repo_path,
            self._config.picasso_branch,
        )
        await copy_picasso_repo(
            repo_path,
            branch=self._config.picasso_branch,
        )
        logger.info("Session Picasso repo cloned: %s", repo_path)
        self._populate_default_picasso_commit_id(repo_path)

    def _populate_default_picasso_commit_id(self, repo_path: Path) -> None:
        """Fill picasso_commit_id from the working repo HEAD when the user omitted it."""
        existing_commit_id = str(self._config.picasso_commit_id or "").strip()
        if existing_commit_id:
            self._config.picasso_commit_id = existing_commit_id
            return

        self._config.picasso_commit_id = None
        try:
            result = subprocess.run(
                ["git", "rev-parse", "HEAD"],
                cwd=str(repo_path),
                capture_output=True,
                text=True,
                timeout=10,
            )
        except Exception as exc:
            logger.warning("Failed to set picasso_commit_id from HEAD: %s", exc)
            return

        if result.returncode != 0:
            logger.warning("Failed to resolve HEAD commit: %s", result.stderr.strip())
            return

        head_commit = result.stdout.strip()
        if not head_commit:
            logger.warning("Failed to set picasso_commit_id from HEAD: empty rev-parse output")
            return

        self._config.picasso_commit_id = head_commit
        logger.info(
            "No picasso_commit_id provided. Using latest commit on branch %s: %s",
            self._config.picasso_branch,
            head_commit[:7],
        )

    async def _create_session_branch(self) -> None:
        """Create an isolated session branch from the user's picasso branch.

        After calling this method, ``self._config.picasso_branch`` is swapped to
        the new session branch so that all downstream operations use it.  The
        original branch is preserved in ``self._user_branch``.

        Raises:
            RuntimeError: If branch creation or push fails.
        """
        from prompt_auto_lab.core.picasso_workspace import create_session_branch

        repo_path = self._get_session_picasso_repo_path()
        session_id = self._session_id or ""
        user_branch = self._config.picasso_branch

        session_branch = await create_session_branch(repo_path, session_id, user_branch)
        self._user_branch = user_branch
        self._session_branch = session_branch
        self._config.picasso_branch = session_branch

    def on_event(
        self,
        callback: Callable[[str, str, str], Awaitable[None]],
    ) -> "AutoPromptLab":
        """
        Register a callback for optimization events.
        
        The callback receives (event_type, agent_name, message) for each event.
        
        Args:
            callback: Async callback function.
        
        Returns:
            Self for method chaining.
        
        Example:
            async def handle_event(event_type, agent, message):
                print(f"[{agent}] {message}")
            
            lab.on_event(handle_event)
        """
        self._event_callback = callback
        if self._orchestrator:
            self._orchestrator._broadcast_event = callback
        return self

    def on_user_input(
        self,
        callback: Callable[..., str | Awaitable[str]],
    ) -> "AutoPromptLab":
        """
        Register user input callback for SDK callback-mode UserProxy.

        The callback receives user proxy context args from AutoGen and should
        return the user reply text.
        """
        self._config.user_input_callback = callback
        if not self._config.user_proxy_mode:
            self._config.user_proxy_mode = "callback"
        if self._orchestrator:
            self._orchestrator._user_input_callback = callback
            self._orchestrator._user_proxy_mode = "callback"
        return self
    
    async def optimize(
        self,
        *,
        max_rounds: int | None = None,
        checkpoint_each_turn: bool = True,
        resume_from: str | Path | None = None,
        **kwargs: Any,
    ) -> OptimizationResult:
        """
        Run a prompt optimization workflow.
        
        Task type is "picasso" (Aipa + Braintrust evaluation).
        Checkpoints, outputs, and logs are saved directly under session_dir.
        
        Args:
            max_rounds: Maximum optimization iterations. If omitted, uses config max_rounds.
            checkpoint_each_turn: If True (default), save checkpoint after each agent turn.
            resume_from: Path to checkpoint file to resume from.
            **kwargs: Additional task-specific parameters.
        
        Returns:
            OptimizationResult with optimization outcomes.
        
        Example:
            # Start new experiment (uses checkpoint_dir from create())
            result = await lab.optimize(max_rounds=10)
            
            # Resume from a checkpoint
            result = await lab.optimize(
                max_rounds=10,
                resume_from="./checkpoints/checkpoint_0005_Analyzer.json",
            )
        """
        if self._status == LabStatus.RUNNING:
            raise RuntimeError("Optimization already running. Call stop() first.")
        if "dataset_path" in kwargs:
            raise TypeError("dataset_path is no longer supported by the SDK optimize() API.")
        await self._ensure_picasso_repo_ready()

        # Create isolated session branch so user's original branch is never modified
        if resume_from is None:
            await self._create_session_branch()

        # Use shared UserInputBuilder
        from prompt_auto_lab.core.orchestrator_builder import (
            UserInputBuilder,
            ProtocolClientBuilder,
            finalize_aipa_task,
            submit_leaderboard_job,
        )
        
        cfg = self._config
        orchestrator = None
        # Use the session branch created by _create_session_branch(), falling back to picasso_branch
        aipa_eval_branch = getattr(self, "_session_branch", None) or cfg.picasso_branch
        aipa_branch_created = False
        user_proxy_mode_override = cfg.user_proxy_mode or ("callback" if cfg.user_input_callback else "console")
        user_input_callback_override: Callable[..., str | Awaitable[str]] | None = cfg.user_input_callback

        async def _stop_aware_input_callback(*cb_args: Any, **cb_kwargs: Any) -> str:
            callback = cfg.user_input_callback
            if callback is not None:
                result = callback(*cb_args, **cb_kwargs)
                if asyncio.iscoroutine(result):
                    result = await result
                text = "" if result is None else str(result)
            else:
                prompt = "User input> "
                if cb_args and isinstance(cb_args[0], str):
                    prompt = cb_args[0]
                text = await asyncio.to_thread(input, prompt)

            if text.strip().lower() == "stop":
                logger.info("Received stop command from user input; cancelling optimization workflow.")
                running = self._running_task
                if running is not None and not running.done():
                    running.cancel()
            return text

        if user_proxy_mode_override in ("console", "callback") or cfg.user_input_callback is not None:
            # Route console/callback through one stop-aware callback so typing "stop" can cancel immediately.
            user_proxy_mode_override = "callback"
            user_input_callback_override = _stop_aware_input_callback

        resolved_max_rounds = int(max_rounds) if max_rounds is not None else int(cfg.max_rounds)
        self._apply_optimize_defaults(kwargs)

        # Task type is always "picasso" for SDK
        # liquid_prompt_entries comes from config (required in create())
        user_input = UserInputBuilder.build(
            task_type="picasso",
            dataset_name=cfg.dataset_name,
            metrics_config=cfg.metrics_config,
            eval_name=cfg.eval_name,
            eval_project_name=cfg.eval_project_name,
            dataset_project_name=cfg.dataset_project_name,
            braintrust_org=cfg.braintrust_org,
            picasso_version=cfg.picasso_version,
            conversation_mode=cfg.conversation_mode,
            max_dataset_rows=cfg.max_dataset_rows,
            branch=cfg.picasso_branch,
            commit_id=cfg.picasso_commit_id,
            liquid_prompt_entries=cfg.liquid_prompt_entries,
            run_leaderboard_after_optimization=cfg.run_leaderboard_after_optimization,
            **kwargs,
        )
        
        # Directory structure (flat under session_dir, no run_ subdirectory):
        # session_dir/
        # ├── checkpoints/    # State snapshots
        # ├── outputs/        # Intermediate outputs
        # └── logs/           # Run logs (one log file per optimize/resume)
        from datetime import datetime
        
        current_run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        run_dir, checkpoint_dir, logs_dir = self._prepare_run_paths(run_id=current_run_id)

        user_input["run_id"] = current_run_id
        
        # Setup logging to logs_dir
        log_path, verbose_log_path = self._setup_run_logging(logs_dir)
        self._log_run_start(current_run_id, log_path)

        try:
            # Build orchestrator with this run's output directory
            orchestrator = self._build_orchestrator(
                run_dir=run_dir,
                user_proxy_mode_override=user_proxy_mode_override,
                user_input_callback_override=user_input_callback_override,
            )

            # Update status
            self._status = LabStatus.RUNNING
            self._running_task = asyncio.current_task()

            await orchestrator.run(
                user_input=user_input,
                max_rounds=resolved_max_rounds,
                use_console=True,
                checkpoint_dir=checkpoint_dir,
                checkpoint_each_turn=checkpoint_each_turn,
                resume_from=resume_from,
            )
            
            # Extract latest state (will be normalized again when building OptimizationResult).
            state = await orchestrator.shared_state.read()

            try:
                await submit_leaderboard_job(
                    enabled=bool(user_input.get("run_leaderboard_after_optimization", False)),
                    github_eval_client=orchestrator.aipa,
                    aipa_eval_branch=aipa_eval_branch,
                    iteration=state.get("iteration") if isinstance(state.get("iteration"), dict) else None,
                    logger_=logger,
                    notify=None,
                )
            except Exception:
                # Keep optimization success even if leaderboard submit fails.
                logger.warning("Failed to trigger leaderboard eval in SDK optimize.", exc_info=True)
            
            logger.info("Run %s completed successfully", current_run_id)
            self._last_result = self._build_optimization_result(
                success=True,
                session_id=self._session_id,
                run_dir=run_dir,
                state=state,
            )
            self._persist_result(run_dir)
            self._status = LabStatus.COMPLETED
            return self._last_result
        
        except asyncio.CancelledError:
            logger.info("Optimization cancelled by user")
            cancelled_state: dict[str, Any] | None = None
            try:
                if orchestrator is not None and hasattr(orchestrator, "shared_state"):
                    cancelled_state = await orchestrator.shared_state.read()
            except BaseException:
                logger.debug("Failed to read shared_state on cancellation.", exc_info=True)
            self._last_result = self._build_optimization_result(
                success=False,
                session_id=self._session_id,
                run_dir=run_dir,
                state=cancelled_state,
                stop_reason="Cancelled by user",
            )
            self._persist_result(run_dir)
            self._status = LabStatus.IDLE
            return self._last_result
        
        except Exception as e:
            logger.exception("Optimization failed")
            failed_state: dict[str, Any] | None = None
            try:
                if orchestrator is not None and hasattr(orchestrator, "shared_state"):
                    failed_state = await orchestrator.shared_state.read()
            except Exception:
                logger.debug("Failed to read shared_state on failure.", exc_info=True)
            self._last_result = self._build_optimization_result(
                success=False,
                session_id=self._session_id,
                run_dir=run_dir,
                state=failed_state,
                stop_reason=str(e),
            )
            self._persist_result(run_dir)
            self._status = LabStatus.FAILED
            return self._last_result
        finally:
            aipa_client = getattr(orchestrator, "aipa", None) if orchestrator is not None else None
            await asyncio.shield(
                finalize_aipa_task(
                    github_eval_client=aipa_client,
                    aipa_eval_branch=aipa_eval_branch,
                    aipa_eval_base_branch=aipa_eval_branch,
                    aipa_branch_created=False,
                    logger_=logger,
                    notify=None,
                    skip_cleanup_on_main=True,
                )
            )
            self._teardown_run_logging()
            self._running_task = None
    
    async def stop(self, timeout: float = 10.0) -> bool:
        """
        Stop the currently running optimization.
        
        This will gracefully cancel the optimization and save a checkpoint
        if checkpointing is enabled.
        
        Args:
            timeout: Maximum time to wait for graceful shutdown in seconds.
        
        Returns:
            True if stopped successfully, False if not running or timeout.
        
        Example:
            # In another task or signal handler
            await lab.stop()
        """
        if self._status != LabStatus.RUNNING:
            logger.warning("stop() called but no optimization running")
            return False
        
        if self._running_task is None:
            logger.warning("stop() called but no task to cancel")
            return False
        
        self._status = LabStatus.STOPPING
        logger.info("Stopping optimization...")
        
        # Cancel the running task
        self._running_task.cancel()
        
        try:
            await asyncio.wait_for(
                asyncio.shield(self._running_task),
                timeout=timeout,
            )
        except asyncio.TimeoutError:
            logger.warning("Timeout waiting for optimization to stop")
            return False
        except asyncio.CancelledError:
            pass  # Expected
        
        self._status = LabStatus.IDLE
        self._running_task = None
        logger.info("Optimization stopped")
        return True
    
    async def resume(
        self,
        checkpoint_path: str | Path | None = None,
        **kwargs: Any,
    ) -> OptimizationResult:
        """
        Resume optimization from a checkpoint.

        Two ways to resume:
        1. Auto-resume: No arguments - uses latest checkpoint from session_dir
        2. By path: Specify exact checkpoint_path to resume from

        Args:
            checkpoint_path: Path to checkpoint file.
            **kwargs: Override any parameters from the original run.

        Returns:
            OptimizationResult from the resumed run.

        Example:
            # Auto-resume from latest checkpoint
            result = await lab.resume()

            # Resume from specific checkpoint file
            result = await lab.resume(checkpoint_path="./runs/checkpoints/checkpoint_0005.json")
        """
        from prompt_auto_lab.orchestrator.orchestrator import find_latest_checkpoint
        await self._ensure_picasso_repo_ready()

        sess_dir = self._config.session_dir

        # Determine checkpoint path
        if checkpoint_path is None:
            # Look for checkpoints in session_dir/checkpoints/ first, then session_dir/
            checkpoint_dir = sess_dir / "checkpoints"
            if checkpoint_dir.exists():
                checkpoint_path = find_latest_checkpoint(checkpoint_dir)
            else:
                checkpoint_path = find_latest_checkpoint(sess_dir)

            if checkpoint_path is None:
                raise ValueError(
                    "No checkpoint found. "
                    "Run optimize() first or provide a checkpoint_path."
                )

        checkpoint_path = Path(checkpoint_path)
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

        logger.info("Resuming from checkpoint: %s", checkpoint_path)

        max_rounds_raw = kwargs.pop("max_rounds", None)
        max_rounds = int(max_rounds_raw) if max_rounds_raw is not None else int(self._config.max_rounds)

        # Call optimize with resume_from
        return await self.optimize(
            max_rounds=max_rounds,
            checkpoint_each_turn=True,
            resume_from=checkpoint_path,
            **kwargs,
        )

    def list_checkpoints(self) -> list[dict[str, Any]]:
        """
        List available checkpoints.

        Returns:
            List of checkpoint info dicts with keys: path, agent_turn, agent_name, saved_at

        Example:
            checkpoints = lab.list_checkpoints()
            for cp in checkpoints:
                print(f"{cp['agent_turn']}: {cp['agent_name']} at {cp['saved_at']}")
        """
        import json as json_module

        sess_dir = self._config.session_dir

        # Checkpoints in session_dir/checkpoints/
        checkpoint_dir = sess_dir / "checkpoints"
        target_dir = checkpoint_dir if checkpoint_dir.exists() else sess_dir

        checkpoints = []
        for checkpoint_file in target_dir.glob("checkpoint_*.json"):
            try:
                name = checkpoint_file.stem
                parts = name.split("_")
                if len(parts) >= 3:
                    agent_turn = int(parts[1])
                    agent_name = "_".join(parts[2:])

                    # Read saved_at from file
                    data = json_module.loads(checkpoint_file.read_text(encoding="utf-8"))
                    saved_at = data.get("saved_at", "")

                    checkpoints.append({
                        "path": str(checkpoint_file),
                        "agent_turn": agent_turn,
                        "agent_name": agent_name,
                        "saved_at": saved_at,
                    })
            except (ValueError, json_module.JSONDecodeError):
                continue

        # Sort by agent_turn descending (latest first)
        checkpoints.sort(key=lambda c: c["agent_turn"], reverse=True)
        return checkpoints

    async def delete_task(
        self,
        *,
        delete_branches: bool = True,
        branches_to_delete: list[str] | None = None,
    ) -> dict[str, Any]:
        """
        Delete task artifacts and related generated Picasso branches.

        Args:
            delete_branches: Whether to delete generated Picasso branches.
            branches_to_delete: Optional explicit list of Picasso branches to delete.
                When omitted, generated branches are discovered from task state.

        Returns:
            Dict with deletion summary.
        """
        if self._status == LabStatus.RUNNING:
            raise RuntimeError("Cannot delete running task. Stop it first.")

        sess_dir = self._config.session_dir
        session_repo_path = (
            self._get_session_picasso_repo_path()
            if self._session_id is not None
            else None
        )
        if not sess_dir.exists() and (
            session_repo_path is None or not Path(session_repo_path).exists()
        ):
            return {
                "status": "not_found",
                "run_dir_deleted": False,
                "session_repo_deleted": False,
                "deleted_branches": [],
                "failed_branches": [],
                "skipped_branches": [],
            }

        from prompt_auto_lab.core.orchestrator_builder import (
            ProtocolClientBuilder,
            collect_candidate_branches_from_task_dir,
            delete_task_artifacts,
        )
        self._teardown_run_logging()

        git_client = None
        if delete_branches:
            try:
                git_client = ProtocolClientBuilder.build_git_client(self._get_session_picasso_repo_path())
            except Exception as e:
                logger.warning("Failed to initialize GitClient for branch cleanup: %s", e)

        # Use user's original branch as the protected branch for cleanup
        original_branch = self._user_branch or self._config.picasso_branch

        # Resolve branch cleanup set.
        # - If caller provides branches_to_delete, use that list.
        # - Otherwise, auto-discover generated branches from task artifacts.
        explicit_branches: set[str] | None = None
        if branches_to_delete is not None:
            explicit_branches = {
                branch.strip()
                for branch in branches_to_delete
                if isinstance(branch, str) and branch.strip()
            }
        else:
            explicit_branches = collect_candidate_branches_from_task_dir(
                sess_dir,
                original_branch=original_branch,
                repo_path=session_repo_path,
            )

        # Always attempt to include session branch cleanup when known.
        if self._session_branch:
            explicit_branches.add(self._session_branch)

        cleanup_result = await delete_task_artifacts(
            task_dir=sess_dir,
            original_branch=original_branch,
            session_repo_path=session_repo_path,
            delete_branches=delete_branches and git_client is not None,
            git_client=git_client,
            branches_to_delete=explicit_branches,
            logger_=logger,
        )
        run_dir_deleted = bool(cleanup_result.get("task_dir_deleted"))
        session_repo_deleted = bool(cleanup_result.get("session_repo_deleted"))
        deleted_branches = list(cleanup_result.get("deleted_branches", []))
        failed_branches = list(cleanup_result.get("failed_branches", []))
        skipped_branches = list(cleanup_result.get("skipped_branches", []))

        self._current_run_id = None

        return {
            "status": "deleted" if run_dir_deleted else "partial",
            "run_dir_deleted": run_dir_deleted,
            "session_repo_deleted": session_repo_deleted,
            "deleted_branches": deleted_branches,
            "failed_branches": failed_branches,
            "skipped_branches": skipped_branches,
        }
    
    async def get_state(self) -> dict[str, Any]:
        """
        Get the current shared state.
        
        Returns:
            Current state dictionary, or empty dict if not running.
        """
        if self._orchestrator is None:
            return {}
        
        if not hasattr(self._orchestrator, "shared_state"):
            return {}
        
        return await self._orchestrator.shared_state.read()
    
    def get_progress(self) -> dict[str, Any]:
        """
        Get current optimization progress.
        
        Returns:
            Dict with iteration, max_rounds, best_score, status.
        """
        return {
            "status": self._status.value,
            "is_running": self.is_running,
            "last_result": self._last_result,
        }

    def get_iterations(self) -> dict[str, Any]:
        """Return the ``iteration`` map from the last result's state.

        The iteration map is a dict keyed by iteration identifier
        (``"root"`` for baseline, ``"1"``, ``"2"``, … for optimisation
        rounds).  Each value contains candidates, scores, branch info, etc.

        Returns:
            The iteration dict, or an empty dict if no result is available.

        Example::

            iters = lab.get_iterations()
            for key, data in iters.items():
                print(key, data.get("overall_score"))
        """
        result = self._last_result
        if result is None:
            return {}
        iteration = result.state.get("iteration")
        return dict(iteration) if isinstance(iteration, dict) else {}

    def get_baseline(self) -> dict[str, Any] | None:
        """Return the baseline (``root``) iteration from the last result.

        The baseline is the evaluation of the **original prompt** before any
        optimisation happens.  It is stored under the ``"root"`` key in the
        iteration map.

        Returns:
            The root iteration dict (with ``candidates``, ``overall_score``,
            ``dimension_metrics``, etc.), or ``None`` if no baseline exists.

        Example::

            baseline = lab.get_baseline()
            if baseline:
                print("Baseline score:", baseline.get("overall_score"))
        """
        return self.get_iterations().get("root")

    def get_result(
        self,
        refresh: bool = True,
    ) -> OptimizationResult | None:
        """
        Get an OptimizationResult for the current session.

        Lookup order when *refresh=True*:
        1. ``session_dir/result.json`` (persisted by ``optimize()`` / ``resume()``).
        2. Rebuild from ``version.jsonl`` state log (legacy / fallback).
        3. In-memory ``last_result`` cache.

        Args:
            refresh: If False, return cached last_result only. If True, try to
                reload from the persisted ``result.json`` first.

        Returns:
            Parsed OptimizationResult, or None if nothing is available.
        """
        if not refresh:
            return self._last_result

        sess_dir = self._config.session_dir

        # 1. Try persisted result.json
        result_file = sess_dir / "result.json"
        if result_file.exists():
            try:
                loaded = OptimizationResult.load(result_file)
                self._last_result = loaded
                return loaded
            except Exception:
                logger.warning("Failed to load result.json, falling back to version.jsonl", exc_info=True)

        # 2. Fallback: rebuild from version.jsonl
        persisted_state = self._load_latest_version_state(sess_dir)
        if not isinstance(persisted_state, dict):
            return self._last_result

        result = self._build_optimization_result(
            success=True,
            session_id=self._session_id,
            run_dir=sess_dir,
            state=persisted_state,
        )
        self._last_result = result
        return result

    # ── Result persistence helpers ────────────────────────────────

    RESULT_FILENAME: str = "result.json"
    """Default filename for the persisted OptimizationResult."""

    def _persist_result(self, run_dir: Path) -> None:
        """Save ``_last_result`` to ``run_dir/result.json`` (best-effort)."""
        if self._last_result is None:
            return
        try:
            self._last_result.save(run_dir / self.RESULT_FILENAME)
            logger.info("Persisted result to %s", run_dir / self.RESULT_FILENAME)
        except Exception:
            logger.warning("Failed to persist result.json", exc_info=True)

    @staticmethod
    def load_result(
        path: str | Path | None = None,
        *,
        session_id: str | None = None,
        user_id: str = "sdk-user",
    ) -> OptimizationResult:
        """Load a previously saved ``OptimizationResult``.

        Three ways to locate the result file (checked in order):

        1. **By session_id** — resolves the path automatically via the
           user-store directory layout so you never need to know the
           underlying file path::

               result = AutoPromptLab.load_result(session_id="my-session")

        2. **By directory** — ``result.json`` is appended automatically::

               result = AutoPromptLab.load_result("./sessions/abc")

        3. **By exact file path**::

               result = AutoPromptLab.load_result("./sessions/abc/result.json")

        Args:
            path: Path to a ``result.json`` file **or** a directory that
                contains one.  Can be omitted when *session_id* is given.
            session_id: Session identifier.  When provided the path is
                resolved automatically using the user-store layout
                (``{store}/users/{user_id}/sessions/{session_id}/result.json``).
            user_id: User id used together with *session_id* (default:
                ``"sdk-user"``).

        Returns:
            The reconstructed ``OptimizationResult``.

        Raises:
            FileNotFoundError: If the result file does not exist.
            ValueError: If neither *path* nor *session_id* is provided.
        """
        if session_id is not None:
            from prompt_auto_lab.core.user_store import get_user_session_dir

            p = get_user_session_dir(str(user_id or "sdk-user"), session_id) / "result.json"
        elif path is not None:
            p = Path(path)
            if p.is_dir():
                p = p / "result.json"
        else:
            raise ValueError("Provide either 'path' or 'session_id'.")

        return OptimizationResult.load(p)

    @staticmethod
    def list_all_results(
        base_dir: str | Path | None = None,
        *,
        user_id: str = "sdk-user",
    ) -> list[dict[str, Any]]:
        """Discover all persisted ``result.json`` files under a directory tree.

        This lets you browse past sessions without having to remember paths::

            results = AutoPromptLab.list_all_results()
            for r in results:
                print(r["session_id"], r["final_score"], r["success"])

            # Load the full result by session_id — no path needed
            full = AutoPromptLab.load_result(session_id=r["session_id"])

        If *base_dir* is omitted the method scans the user-store sessions
        directory for the given *user_id*
        (``{store}/users/{user_id}/sessions/``).

        Args:
            base_dir: Root directory to scan.  All ``result.json`` files
                found recursively underneath are returned.  When ``None``,
                the user-store sessions directory is used.
            user_id: User id used when *base_dir* is ``None`` (default:
                ``"sdk-user"``).

        Returns:
            A list of summary dicts sorted by ``session_id`` descending
            (newest first).  Each dict contains:

            - ``session_id`` — inferred from the parent directory name
            - ``path`` — absolute path to the ``result.json`` file
            - ``success`` — whether optimization succeeded
            - ``final_score`` — final evaluation score (may be ``None``)
            - ``iterations`` — number of iterations completed (derived from ``state["iteration"]``)
            - ``stop_reason`` — why the run stopped (may be ``None``)
        """
        if base_dir is None:
            from prompt_auto_lab.core.user_store import get_user_dir

            base = get_user_dir(str(user_id or "sdk-user")) / "sessions"
        else:
            base = Path(base_dir)

        if not base.exists():
            return []

        entries: list[dict[str, Any]] = []
        for result_file in base.rglob("result.json"):
            try:
                data = json.loads(result_file.read_text(encoding="utf-8"))
            except Exception:
                continue

            # Infer session_id from nearest parent directory
            session_id_inferred = result_file.parent.name

            entries.append({
                "session_id": session_id_inferred,
                "path": str(result_file),
                "success": data.get("success"),
                "final_score": data.get("final_score"),
                "iterations": AutoPromptLab._extract_iteration_count(
                    data.get("state") if isinstance(data.get("state"), dict) else {}
                ),
                "stop_reason": data.get("stop_reason"),
            })

        # Sort newest first by session_id; fall back to empty
        entries.sort(key=lambda e: e.get("session_id") or "", reverse=True)
        return entries

    def _build_optimization_result(
        self,
        *,
        success: bool,
        session_id: str | None,
        run_dir: Path | None = None,
        state: dict[str, Any] | None = None,
        stop_reason: str | None = None,
    ) -> OptimizationResult:
        """
        Build OptimizationResult from persisted version.jsonl + shared_state fallback.

        Priority:
        1) Latest matching record from version.jsonl.
        2) Provided shared_state snapshot.
        """
        state_fallback = state if isinstance(state, dict) else {}
        persisted_state = self._load_latest_version_state(run_dir)
        result_state = persisted_state if isinstance(persisted_state, dict) else state_fallback

        best_metrics = self._extract_best_metrics_from_state(result_state)
        final_score = best_metrics.get("overall_score")
        if final_score is None:
            final_score = self._extract_final_score(result_state)

        best_version = best_metrics.get("best_version")
        if best_version is None and result_state.get("best_iteration_idx") is not None:
            best_version = str(result_state.get("best_iteration_idx"))
        if best_version is None and result_state.get("version") is not None:
            best_version = str(result_state.get("version"))

        final_stop_reason = stop_reason
        if not final_stop_reason:
            state_reason = result_state.get("stop_reason")
            if isinstance(state_reason, str) and state_reason:
                final_stop_reason = state_reason
            else:
                fallback_reason = state_fallback.get("stop_reason")
                if isinstance(fallback_reason, str) and fallback_reason:
                    final_stop_reason = fallback_reason

        return OptimizationResult(
            success=success,
            session_id=session_id,
            best_version=best_version,
            final_score=final_score,
            stop_reason=final_stop_reason,
            dimension_scores=best_metrics.get("dimension_scores"),
            best_branch=best_metrics.get("branch"),
            best_commit_id=best_metrics.get("commit_id"),
            best_candidate_id=best_metrics.get("candidate_id"),
            experiment_id=best_metrics.get("experiment_id"),
            experiment_link=best_metrics.get("experiment_link"),
            state=result_state,
        )

    def _load_latest_version_state(self, run_dir: Path | None) -> dict[str, Any] | None:
        """Load the latest state from version.jsonl, preferring persisted snapshots."""
        candidates: list[Path] = []
        if run_dir is not None:
            candidates.extend(
                [
                    run_dir / "logs" / "version.jsonl",
                    run_dir / "version.jsonl",
                    run_dir / "outputs" / "version.jsonl",
                ]
            )
        if run_dir != self._config.session_dir:
            candidates.append(self._config.session_dir / "version.jsonl")

        for path in candidates:
            if not path.exists():
                continue
            try:
                last_payload: dict[str, Any] | None = None
                with open(path, "r", encoding="utf-8") as f:
                    for line in f:
                        line = line.strip()
                        if not line:
                            continue
                        payload = json.loads(line)
                        if isinstance(payload, dict):
                            last_payload = payload
                if last_payload is not None:
                    return last_payload
            except Exception:
                logger.warning("Failed reading latest version state from %s", path, exc_info=True)
        return None

    def _to_float(self, value: Any) -> float | None:
        if isinstance(value, (int, float)):
            return float(value)
        return None

    def _normalize_dimension_scores(self, raw: Any) -> dict[str, float]:
        if not isinstance(raw, dict):
            return {}
        result: dict[str, float] = {}
        for key, value in raw.items():
            if isinstance(value, (int, float)):
                result[str(key)] = float(value)
            elif isinstance(value, dict):
                score = value.get("score")
                if isinstance(score, (int, float)):
                    result[str(key)] = float(score)
        return result

    @staticmethod
    def _extract_iteration_count(state: dict[str, Any]) -> int:
        iteration = state.get("iteration")
        if not isinstance(iteration, dict):
            return 0
        non_root = [k for k in iteration.keys() if str(k) != "root"]
        if non_root:
            return len(non_root)
        return 1 if "root" in iteration else 0

    def _extract_best_metrics_from_state(self, state: dict[str, Any]) -> dict[str, Any]:
        """
        Extract best metrics from iteration candidates in persisted state.

        This follows UI metrics behavior: choose the candidate with highest overall_score.
        """
        iteration = state.get("iteration")
        best: dict[str, Any] | None = None
        best_score: float | None = None

        if isinstance(iteration, dict):
            for node_id, node_data in iteration.items():
                if not isinstance(node_data, dict):
                    continue

                candidates = node_data.get("candidates")
                if isinstance(candidates, dict) and candidates:
                    for candidate_id, candidate_data in candidates.items():
                        if not isinstance(candidate_data, dict):
                            continue
                        score = self._to_float(candidate_data.get("overall_score"))
                        if score is None:
                            continue
                        if best_score is None or score > best_score:
                            best_score = score
                            best = {
                                "best_version": str(node_id),
                                "candidate_id": str(candidate_id),
                                "branch": candidate_data.get("branch") or "",
                                "commit_id": candidate_data.get("commit_id") or "",
                                "overall_score": score,
                                "dimension_scores": self._normalize_dimension_scores(
                                    candidate_data.get("dimension_metrics") or candidate_data.get("metrics")
                                ),
                                "experiment_id": candidate_data.get("experiment_id") or "",
                                "experiment_link": candidate_data.get("experiment_link") or "",
                            }

                node_score = self._to_float(node_data.get("overall_score"))
                if node_score is not None and (best_score is None or node_score > best_score):
                    best_score = node_score
                    best = {
                        "best_version": str(node_id),
                        "candidate_id": "",
                        "branch": node_data.get("analysis_branch") or node_data.get("branch") or "",
                        "commit_id": node_data.get("analysis_commit_id") or node_data.get("commit_id") or "",
                        "overall_score": node_score,
                        "dimension_scores": self._normalize_dimension_scores(
                            node_data.get("dimension_metrics") or node_data.get("metrics")
                        ),
                        "experiment_id": node_data.get("analysis_experiment_id") or node_data.get("experiment_id") or "",
                        "experiment_link": node_data.get("experiment_link") or "",
                    }

        if best is not None:
            return best

        fallback_score = self._extract_final_score(state)
        return {
            "best_version": str(state.get("version")) if state.get("version") is not None else None,
            "candidate_id": "",
            "branch": "",
            "commit_id": "",
            "overall_score": fallback_score,
            "dimension_scores": self._normalize_dimension_scores(state.get("dimension_scores")),
            "experiment_id": state.get("experiment_id") or "",
            "experiment_link": state.get("experiment_link") or "",
        }
    
    def _extract_final_score(self, state: dict[str, Any]) -> float | None:
        """Extract the final score from the state."""
        metrics = state.get("metrics", {})
        if isinstance(metrics, dict):
            # Try common score keys
            for key in ["score", "overall_score", "avg_score", "accuracy"]:
                if key in metrics:
                    return float(metrics[key])
        return None
    
    async def evaluate(
        self,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """
        Evaluate a prompt without optimization (single round).
        
        Args:
            **kwargs: Additional parameters.
        
        Returns:
            Evaluation metrics dictionary.
        """
        if "dataset_path" in kwargs:
            raise TypeError("dataset_path is no longer supported by the SDK evaluate() API.")
        # Single evaluation run (max_rounds=1, no optimization)
        result = await self.optimize(
            max_rounds=1,
            **kwargs,
        )
        
        return result


# Convenience function for quick usage
async def optimize_prompt(
    *,
    # Required - Local Environment
    braintrust_api_key: str,
    # Required - Experiment Config
    dataset_name: str,
    picasso_branch: str,
    liquid_prompt_entries: list[dict],
    # Optional
    max_rounds: int = 10,
    **kwargs: Any,
) -> OptimizationResult:
    """
    Convenience function to run prompt optimization.

    This is a shortcut for creating an AutoPromptLab instance
    and running optimization in one call.

    Args:
        braintrust_api_key: Braintrust API key (required).
        dataset_name: Dataset name for evaluation (required).
        picasso_branch: Picasso repository branch for prompts (required).
        liquid_prompt_entries: Liquid template entries (required).
        max_rounds: Maximum optimization iterations.
        **kwargs: Additional parameters passed to create().

    Returns:
        OptimizationResult with optimization outcomes.

    Example:
        from prompt_auto_lab import optimize_prompt

        result = asyncio.run(optimize_prompt(
            braintrust_api_key="sk-xxx",
            dataset_name="picasso_cn_chat_0923",
            picasso_branch="autopromptlab2/main",
            liquid_prompt_entries=[{"path": "...", "can_edit": True}],
            max_rounds=10,
        ))
    """
    lab = AutoPromptLab.create(
        braintrust_api_key=braintrust_api_key,
        dataset_name=dataset_name,
        picasso_branch=picasso_branch,
        liquid_prompt_entries=liquid_prompt_entries,
        **kwargs,
    )
    return await lab.optimize(max_rounds=max_rounds)
