"""
Shared builder utilities for creating Orchestrator and its dependencies.

This module provides a unified way to construct the Orchestrator with all
required clients, used by both the API (sessions.py) and SDK (lab.py).
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import random
import re
import shutil
import stat
import subprocess
import time
import urllib.parse
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Awaitable, Literal
from prompt_auto_lab.core.metrics_config import (
    build_default_metrics_config as _shared_build_default_metrics_config,
    get_eval_key_metrics_from_dataset as _shared_get_eval_key_metrics_from_dataset,
    get_default_metrics_config_from_dataset as _shared_get_default_metrics_config_from_dataset,
    normalize_metrics_config as _shared_normalize_metrics_config,
)

logger = logging.getLogger(__name__)


# ============================================================================
# Configuration Data Classes
# ============================================================================

@dataclass
class ModelClients:
    """Named container for all model client roles."""

    main: Any | None = None
    think: Any | None = None
    analysis: Any | None = None


@dataclass
class ModelClientConfig:
    """Configuration for model client creation."""
    
    client_type: str = "Azure"
    """Model client type: Azure, Papyrus, Copilot, Local."""
    
    # Azure specific
    azure_endpoint: str | None = None
    model_deployment: str | None = None
    azure_api_version: str | None = None
    
    # Papyrus specific
    papyrus_endpoint: str | None = None
    papyrus_model: str | None = None
    papyrus_quota_id: str | None = None
    papyrus_scope: str | None = None
    
    # Copilot specific
    copilot_model: str = ""
    
    # Local specific
    local_endpoint: str = "http://localhost:23333/api/openai/v1"
    local_model: str = ""
    
    # Structured output
    use_structured_output: bool = True


@dataclass
class EvalClientConfig:
    """Configuration for evaluation client creation."""
    
    # Paths
    github_cli_path: str = ""
    output_dir: str | Path = ""
    
    # Eval settings
    eval_name: str | None = None
    dataset_name: str = ""
    local_debug: bool = False
    dataset_sampling_method: str | None = None
    max_dataset_rows: int | None = None
    aipa_eval_branch: str = "main"
    conversation_mode: str = "DefaultMode"
    additional_cmd_args: str | None = None
    
    # Feature flags
    add_feature_flags: str | None = None
    baseline_feature_flags: str | None = None
    picasso_version: str | None = None
    environment: str | None = None
    model: str | None = None
    picasso_env: str | None = None
    task_id_prefix: str | None = None
    
    # Project settings
    eval_project_name: str | None = None
    dataset_project_name: str | None = None
    braintrust_org: str = "non-pme"


@dataclass
class BraintrustConfig:
    """Configuration for Braintrust client creation."""
    
    api_key: str = ""
    metrics_config: list[dict[str, Any]] | None = None
    # Backward compatibility: old key_metrics path.
    key_metrics: list[str] | None = None
    dimension_threshold: float = 0.5
    bad_case_limit_per_iteration: int = 10
    output_dir: str | Path = ""
    local_debug: bool = False
    eval_name: str | None = None
    braintrust_org: str = "non-pme"
    environment: str | None = None
    
    # Model info for BTQL agent
    model_client_type: str = "azure"
    model_endpoint: str | None = None
    model_name: str | None = None
    
    # Optional callback
    log_callback: Callable | None = None
    badcase_deep_analysis: bool = False


@dataclass
class OrchestratorConfig:
    """Full configuration for Orchestrator creation."""
    
    # Model
    model: ModelClientConfig = field(default_factory=ModelClientConfig)
    prebuilt_clients: ModelClients | None = None
    prebuilt_model_client: Any | None = None
    prebuilt_model_think_client: Any | None = None
    
    # Evaluation
    eval: EvalClientConfig = field(default_factory=EvalClientConfig)
    
    # Braintrust
    braintrust: BraintrustConfig = field(default_factory=BraintrustConfig)
    
    # Picasso / Git
    picasso_repo_path: str | Path = ""
    session_dir: str | Path = ""
    
    # User / Session
    user_id: str = ""
    session_id: str = ""
    user_proxy_mode: Literal["web", "console", "callback"] | None = None
    user_input_callback: Callable[..., Any] | None = None
    user_proxy_timeout: float = 3600.0
    prefer_native_user_proxy: bool = False
    analysis_agent_client: Any | None = None


# ============================================================================
# Model Client Builders
# ============================================================================

class ModelClientBuilder:
    """Builder for model clients."""
    
    @staticmethod
    def build_azure_clients(
        deployment: str | None = None,
        endpoint: str | None = None,
        token_provider: Callable | None = None,
        api_version: str | None = None,
    ):
        """
        Build Azure OpenAI model clients.
        
        Returns:
            Tuple of (model_client, model_think_client)
        """
        from prompt_auto_lab.core.auth_token import get_azure_token_provider
        from prompt_auto_lab.core.model_clients import (
            build_azure_model_client,
            build_azure_model_thought_client,
        )
        
        if token_provider is None:
            token_provider = get_azure_token_provider("https://cognitiveservices.azure.com/.default")
        
        azure_endpoint = endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
        if not azure_endpoint:
            raise ValueError("Azure endpoint is required. Pass it explicitly or set AZURE_OPENAI_ENDPOINT.")
        if not deployment:
            raise ValueError("Azure model deployment name is required.")
        if not api_version:
            raise ValueError("Azure API version is required.")
        
        model_client = build_azure_model_client(
            token_provider=token_provider,
            deployment=deployment,
            endpoint=azure_endpoint,
            api_version=api_version,
        )
        model_think_client = build_azure_model_thought_client(
            token_provider=token_provider,
            endpoint=azure_endpoint,
            api_version=api_version,
        )
        
        return model_client, model_think_client
    
    @staticmethod
    def build_papyrus_clients(
        endpoint: str | None = None,
        model: str | None = None,
        token_provider: Callable | None = None,
        quota_id: str | None = None,
        scope: str | None = None,
    ):
        """
        Build Papyrus model clients.
        
        Returns:
            Tuple of (model_client, model_think_client)
        """
        from prompt_auto_lab.core.auth_token import get_azure_token_provider
        from prompt_auto_lab.core.model_clients import build_papyrus_model_client
        
        papyrus_scope = scope or os.getenv("PAPYRUS_VERIFY_SCOPES")
        if not papyrus_scope:
            raise ValueError("Papyrus scope is required. Pass it explicitly or set PAPYRUS_VERIFY_SCOPES.")
        if token_provider is None:
            token_provider = get_azure_token_provider(papyrus_scope)
        
        papyrus_endpoint = endpoint or os.getenv("PAPYRUS_ENDPOINT")
        if not papyrus_endpoint:
            raise ValueError("Papyrus endpoint is required. Pass it explicitly or set PAPYRUS_ENDPOINT.")
        papyrus_model = model or os.getenv("PAPYRUS_MODEL")
        if not papyrus_model:
            raise ValueError("Papyrus model is required. Pass it explicitly or set PAPYRUS_MODEL.")
        papyrus_quota_id = quota_id or os.getenv("PAPYRUS_QUOTA_ID", "")
        
        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,
        )
        
        return model_client, model_think_client
    
    @staticmethod
    def build_copilot_clients(model: str = ""):
        """
        Build Copilot proxy clients.
        
        Returns:
            Tuple of (model_client, model_think_client)
        """
        from prompt_auto_lab.core.model_clients import build_copilot_proxy_client
        
        model_client = build_copilot_proxy_client(model=model)
        model_think_client = build_copilot_proxy_client(model=model)
        
        return model_client, model_think_client
    
    @staticmethod
    def build_local_clients(
        model: str,
        endpoint: str = "http://localhost:23333/api/openai/v1",
    ):
        """
        Build local model clients.
        
        Returns:
            Tuple of (model_client, model_think_client)
        """
        from prompt_auto_lab.core.model_clients import build_local_model_client
        
        base_url = endpoint.rstrip("/")
        
        model_client = build_local_model_client(model=model, base_url=base_url)
        model_think_client = build_local_model_client(model=model, base_url=base_url)
        
        return model_client, model_think_client

    @classmethod
    def build(cls, config: ModelClientConfig):
        """
        Build model clients based on configuration.
        
        Returns:
            ModelClients(main, think, analysis)
        """
        client_type = config.client_type.lower()
        
        if client_type == "papyrus":
            main, think = cls.build_papyrus_clients(
                endpoint=config.papyrus_endpoint,
                model=config.papyrus_model,
                quota_id=config.papyrus_quota_id,
                scope=config.papyrus_scope,
            )
        elif client_type == "copilot":
            main, think = cls.build_copilot_clients(model=config.copilot_model)
        elif client_type == "local":
            if not config.local_model:
                raise ValueError("Local model ID is required.")
            main, think = cls.build_local_clients(
                model=config.local_model,
                endpoint=config.local_endpoint,
            )
        else:
            # Default: Azure
            main, think = cls.build_azure_clients(
                deployment=config.model_deployment,
                endpoint=config.azure_endpoint,
                api_version=config.azure_api_version,
            )

        return ModelClients(main=main, think=think, analysis=main)


# ============================================================================
# Protocol Client Builders
# ============================================================================

class ProtocolClientBuilder:
    """Builder for protocol clients (eval, braintrust, git, etc.)."""
    
    @staticmethod
    def build_eval_client(config: EvalClientConfig, backend=None):
        """
        Build GitHubActionEvalClient.
        
        Args:
            config: Evaluation client configuration
            backend: Optional backend, will create InMemoryEvalBackend if None
        
        Returns:
            GitHubActionEvalClient instance
        """
        from prompt_auto_lab.core.protocol import (
            GitHubActionEvalClient,
            InMemoryEvalBackend,
        )
        
        if backend is None:
            backend = InMemoryEvalBackend()
        
        return GitHubActionEvalClient(
            backend=backend,
            github_cli_path=config.github_cli_path,
            eval_name=config.eval_name,
            dataset_name=config.dataset_name,
            output_dir=str(config.output_dir),
            local_debug=config.local_debug,
            add_feature_flags=config.add_feature_flags,
            baseline_feature_flags=config.baseline_feature_flags,
            aipa_eval_branch=config.aipa_eval_branch,
            conversation_mode=config.conversation_mode,
            dataset_sampling_method=config.dataset_sampling_method,
            max_dataset_rows=config.max_dataset_rows,
            additional_cmd_args=config.additional_cmd_args,
            picasso_version=config.picasso_version,
            eval_project_name=config.eval_project_name,
            dataset_project_name=config.dataset_project_name,
            braintrust_org=config.braintrust_org,
            environment=config.environment,
            model=config.model,
            picasso_env=config.picasso_env,
            task_id_prefix=config.task_id_prefix,
        )
    
    @staticmethod
    def build_braintrust_client(
        config: BraintrustConfig,
        backend=None,
        model_client=None,
        analysis_model_client=None,
        analysis_agent_client=None,
    ):
        """
        Build BraintrustClient.
        
        Args:
            config: Braintrust configuration
            backend: Optional backend, will create InMemoryEvalBackend if None
            model_client: Optional model client for BTQL agent
        
        Returns:
            BraintrustClient instance
        """
        from prompt_auto_lab.core.protocol import (
            BraintrustClient,
            InMemoryEvalBackend,
        )
        
        if backend is None:
            backend = InMemoryEvalBackend()
        
        metrics_config = _shared_normalize_metrics_config(config.metrics_config)
        if metrics_config is None:
            metrics_config = _shared_build_default_metrics_config(config.key_metrics)

        return BraintrustClient(
            backend=backend,
            api_key=config.api_key,
            metrics_config=metrics_config,
            dimension_threshold=config.dimension_threshold,
            bad_case_limit_per_iteration=config.bad_case_limit_per_iteration,
            output_dir=str(config.output_dir),
            local_debug=config.local_debug,
            eval_name=config.eval_name,
            braintrust_org=config.braintrust_org,
            environment=config.environment,
            model_client_type=config.model_client_type,
            model_endpoint=config.model_endpoint,
            model_name=config.model_name,
            log_callback=config.log_callback,
            model_client=model_client,
            analysis_model_client=analysis_model_client,
            analysis_agent_client=analysis_agent_client,
            badcase_deep_analysis=config.badcase_deep_analysis,
        )
    
    @staticmethod
    def build_picasso_client(repo_path: str | Path):
        """Build PicassoClient."""
        from prompt_auto_lab.core.protocol import PicassoClient
        return PicassoClient(repo_path=repo_path)
    
    @staticmethod
    def build_git_client(repo_path: str | Path, session_dir: str | Path | None = None):
        """Build GitClient."""
        from prompt_auto_lab.core.protocol import GitClient
        return GitClient(repo_path=repo_path, store_dir=session_dir)
    
    @staticmethod
    def build_code_executors():
        """
        Build code executors.
        
        Returns:
            Tuple of (liquid_executor, python_executor)
        """
        from prompt_auto_lab.core.protocol import (
            SimpleLiquidCodeExecutor,
            LocalPythonCodeExecutor,
        )
        return SimpleLiquidCodeExecutor(), LocalPythonCodeExecutor()


# ============================================================================
# Orchestrator Builder
# ============================================================================

class OrchestratorBuilder:
    """
    Main builder for creating fully configured Orchestrator instances.
    
    This class provides a unified way to construct the Orchestrator
    with all required dependencies, used by both API and SDK.
    """
    
    @staticmethod
    def build(
        config: OrchestratorConfig,
        event_callback: Callable[[str, str, str, dict | None], Awaitable[None]] | None = None,
    ):
        """
        Build a fully configured Orchestrator.
        
        Args:
            config: Complete orchestrator configuration
            event_callback: Optional callback for broadcasting events
        
        Returns:
            Configured Orchestrator instance
        """
        from prompt_auto_lab.core.protocol import InMemoryEvalBackend
        from prompt_auto_lab.orchestrator.orchestrator import Orchestrator
        
        # Build model clients (or reuse prebuilt ones when provided)
        if config.prebuilt_clients is not None:
            model_client = config.prebuilt_clients.main
            model_think_client = config.prebuilt_clients.think
            analysis_model_client = config.prebuilt_clients.analysis or model_client
        else:
            analysis_model_client = None
        has_model_client = config.prebuilt_model_client is not None
        has_model_think_client = config.prebuilt_model_think_client is not None
        if config.prebuilt_clients is None and has_model_client and has_model_think_client:
            model_client = config.prebuilt_model_client
            model_think_client = config.prebuilt_model_think_client
            analysis_model_client = model_client
        elif config.prebuilt_clients is None and (has_model_client or has_model_think_client):
            raise ValueError(
                "Both prebuilt_model_client and prebuilt_model_think_client must be provided together."
            )
        elif config.prebuilt_clients is None:
            built_clients = ModelClientBuilder.build(config.model)
            model_client = built_clients.main
            model_think_client = built_clients.think
            analysis_model_client = built_clients.analysis or model_client
        
        # Build shared backend
        backend = InMemoryEvalBackend()
        
        # Build eval client
        eval_client = ProtocolClientBuilder.build_eval_client(config.eval, backend)
        
        # Build braintrust client
        braintrust_client = ProtocolClientBuilder.build_braintrust_client(
            config.braintrust,
            backend,
            model_client,
            analysis_model_client,
            config.analysis_agent_client,
        )
        
        # Build picasso and git clients
        picasso_client = ProtocolClientBuilder.build_picasso_client(config.picasso_repo_path)
        git_client = ProtocolClientBuilder.build_git_client(
            config.picasso_repo_path,
            config.session_dir,
        )
        
        # Build code executors
        liquid_executor, python_executor = ProtocolClientBuilder.build_code_executors()
        
        # Determine structured output setting
        use_structured_output = config.model.use_structured_output
        if config.model.client_type.lower() in ("local", "copilot"):
            use_structured_output = False
        
        # Create orchestrator
        orchestrator = Orchestrator(
            model_client=model_client,
            model_think_client=model_think_client,
            aipa=eval_client,
            braintrust=braintrust_client,
            picasso=picasso_client,
            git_agent=git_client,
            liquid_executor=liquid_executor,
            python_executor=python_executor,
            store_dir=config.session_dir,  # Orchestrator uses store_dir internally
            use_structured_output=use_structured_output,
        )

        # Configure user proxy mode per entrypoint (API/SDK/CLI) while keeping defaults.
        if config.user_proxy_mode is not None:
            orchestrator._user_proxy_mode = config.user_proxy_mode
        if config.user_input_callback is not None:
            orchestrator._user_input_callback = config.user_input_callback
        orchestrator._user_proxy_timeout = config.user_proxy_timeout
        orchestrator._prefer_native_user_proxy = config.prefer_native_user_proxy
        
        # Set event callback if provided
        if event_callback:
            orchestrator._broadcast_event = event_callback
        
        # Set session info if provided
        if config.session_id:
            orchestrator._session_id = config.session_id
        
        return orchestrator


# ============================================================================
# User Input Builder
# ============================================================================

class UserInputBuilder:
    """Builder for constructing user_input dict for orchestrator.run()."""
    
    @staticmethod
    def build(
        task_type: str = "picasso",
        *,
        # Basic settings
        goal: str = "Optimize a base prompt for higher evaluation score.",
        run_id: str | None = None,
        
        # Dataset / Eval settings
        dataset_name: str = "",
        key_metrics: list[str] | None = None,
        eval_name: str | None = None,
        eval_project_name: str | None = None,
        dataset_project_name: str | None = None,
        braintrust_org: str = "non-pme",
        
        # Picasso settings
        picasso_version: str | None = None,
        conversation_mode: str | None = None,
        max_dataset_rows: int | None = None,
        
        # Git / Prompt settings
        branch: str = "main",
        commit_id: str | None = None,
        liquid_prompt_entries: list[dict] | None = None,
        original_prompt_file_path: str | None = None,
        
        # Termination
        termination_rule: dict | None = None,
        
        # Skip eval (if experiment already exists)
        baseline_experiment_id: str | None = None,
        baseline_experiment_link: str | None = None,
        
        # Extra kwargs
        **kwargs,
    ) -> dict[str, Any]:
        """
        Build user_input dict for orchestrator.run().
        
        Args:
            task_type: Task type (picasso, simple, etc.)
            goal: Optimization goal description
            run_id: Unique run identifier (auto-generated if not provided)
            ... (other parameters)
        
        Returns:
            Dict suitable for passing to orchestrator.run(user_input=...)
        """
        if run_id is None:
            run_id = str(int(time.time()))
        
        if termination_rule is None:
            termination_rule = {"type": "no_progress", "rounds": 3}
        
        user_input = {
            "task_type": task_type,
            "goal": goal,
            "run_id": run_id,
            "dataset_name": dataset_name,
            "key_metrics": key_metrics,
            "eval_name": eval_name,
            "eval_project_name": eval_project_name,
            "dataset_project_name": dataset_project_name,
            "braintrust_org": braintrust_org,
            "picasso_version": picasso_version,
            "conversation_mode": conversation_mode,
            "max_dataset_rows": max_dataset_rows,
            "branch": branch,
            "commit_id": commit_id,
            "liquid_prompt_entries": liquid_prompt_entries,
            "original_prompt_file_path": original_prompt_file_path,
            "termination_rule": termination_rule,
        }
        
        # Map baseline experiment info to orchestrator-compatible keys.
        if baseline_experiment_id:
            user_input["experiment_id"] = baseline_experiment_id
        if baseline_experiment_link:
            user_input["experiment_link"] = baseline_experiment_link
        
        # Add any extra kwargs
        user_input.update(kwargs)
        
        return user_input


# ============================================================================
# Helper Functions
# ============================================================================

def get_eval_key_metrics_from_dataset(dataset_name: str) -> list[str] | None:
    """Derive key_metrics from dataset_name using shared core mapping."""
    return _shared_get_eval_key_metrics_from_dataset(dataset_name)


def get_default_metrics_config_from_dataset(dataset_name: str) -> list[dict[str, Any]] | None:
    """Build default metrics_config from shared dataset mapping."""
    return _shared_get_default_metrics_config_from_dataset(dataset_name)


def parse_max_dataset_rows(value: Any) -> int | None:
    """Parse max_dataset_rows from various input types."""
    if isinstance(value, int):
        return value
    elif isinstance(value, str):
        try:
            return int(value.strip())
        except ValueError:
            return None
    return None


def parse_key_metrics(value: Any) -> list[str] | None:
    """Parse key_metrics from various input types."""
    if isinstance(value, list):
        cleaned = [m.strip() for m in value if isinstance(m, str) and m.strip()]
        return cleaned if cleaned else None
    elif isinstance(value, str):
        cleaned = [m.strip() for m in value.split(",") if m.strip()]
        return cleaned if cleaned else None
    return None


def _parse_github_nwo(remote_url: str) -> str:
    """Parse ``owner/repo`` from a GitHub remote URL.

    Supports HTTPS (``https://github.com/owner/repo.git``),
    SSH (``git@github.com:owner/repo.git``), and
    ``ssh://`` style URLs.
    """
    url = remote_url.strip()
    if url.endswith(".git"):
        url = url[:-4]
    m = re.search(r"github\.com[/:]([^/]+/[^/]+)$", url)
    if m:
        return m.group(1)
    raise ValueError(f"Cannot parse GitHub owner/repo from remote URL: {remote_url!r}")


# Default picasso repository NWO (owner/repo)
_DEFAULT_AIPA_EVAL_REPO = "infinity-microsoft/picasso"


async def _resolve_repo_nwo() -> str:
    """Resolve ``owner/repo`` for the aipa-evals repository."""
    return _DEFAULT_AIPA_EVAL_REPO


async def _gh_api(
    args: list[str],
    *,
    cwd: str | Path | None = None,
    timeout: float = 30,
) -> subprocess.CompletedProcess[str]:
    """Run a ``gh api`` command asynchronously."""
    cmd = ["gh", "api", *args]

    def _run() -> subprocess.CompletedProcess[str]:
        return subprocess.run(
            cmd,
            cwd=str(cwd) if cwd else None,
            capture_output=True,
            text=True,
            timeout=timeout,
        )

    return await asyncio.to_thread(_run)


async def setup_aipa_eval_branch(
    aipa_git_client: Any,
    *,
    base_branch: str,
    run_id: str,
    logger_: logging.Logger | None = None,
) -> tuple[str, bool, str | None]:
    """Create a new aipa-eval branch on the *remote* via GitHub API.

    All operations are remote-only (``gh api``) so the local working
    directory is never affected.

    Returns:
        (new_branch_name, success, error_message)
    """
    log = logger_ or logger
    try:
        repo_nwo = await _resolve_repo_nwo()

        # 1. Get the SHA of the base branch from GitHub.
        ref_result = await _gh_api(
            [f"repos/{repo_nwo}/git/ref/heads/{base_branch}", "--jq", ".object.sha"],
        )
        if ref_result.returncode != 0:
            error_msg = (
                f"Failed to resolve SHA for base branch '{base_branch}': "
                f"{ref_result.stderr.strip()}"
            )
            log.error(error_msg)
            return base_branch, False, error_msg

        base_sha = ref_result.stdout.strip()
        if not base_sha:
            error_msg = f"Empty SHA returned for base branch '{base_branch}'"
            log.error(error_msg)
            return base_branch, False, error_msg

        # 2. Create the new branch on the remote.
        random_suffix = "".join(
            random.choices("0123456789abcdefghijklmnopqrstuvwxyz", k=6)
        )
        new_branch_name = f"autopromptlab/aipa_eval/{run_id}_{random_suffix}"

        create_result = await _gh_api(
            [
                f"repos/{repo_nwo}/git/refs",
                "-f", f"ref=refs/heads/{new_branch_name}",
                "-f", f"sha={base_sha}",
            ],
        )
        if create_result.returncode != 0:
            error_msg = (
                f"Failed to create remote branch '{new_branch_name}': "
                f"{create_result.stderr.strip()}"
            )
            log.error(error_msg)
            return base_branch, False, error_msg

        log.info(
            "Created remote aipa-eval branch: %s (from %s @ %s)",
            new_branch_name,
            base_branch,
            base_sha[:7],
        )
        return new_branch_name, True, None
    except Exception as e:
        error_msg = f"Unexpected error setting up aipa-eval branch: {e}"
        log.error(error_msg)
        return base_branch, False, error_msg


async def cleanup_aipa_eval_branch(
    aipa_git_client: Any,
    *,
    branch_name: str,
    base_branch: str,
    logger_: logging.Logger | None = None,
) -> bool:
    """Delete a temporary aipa-eval branch on the *remote* via GitHub API.

    Only remote-only; the local working directory is never touched.
    """
    log = logger_ or logger
    try:
        if not branch_name.startswith("autopromptlab/aipa_eval/"):
            log.info("Skipping cleanup of branch %s - not an auto-created branch", branch_name)
            return True

        repo_nwo = await _resolve_repo_nwo()

        delete_result = await _gh_api(
            ["-X", "DELETE", f"repos/{repo_nwo}/git/refs/heads/{branch_name}"],
        )
        if delete_result.returncode != 0:
            log.warning(
                "Failed to delete remote branch %s: %s",
                branch_name,
                delete_result.stderr.strip(),
            )
            return False

        log.info("Cleaned up remote aipa-eval branch: %s", branch_name)
        return True
    except Exception as e:
        log.error("Error cleaning up aipa-eval branch: %s", e)
        return False


async def finalize_aipa_task(
    *,
    github_eval_client: Any,
    aipa_eval_branch: str,
    aipa_eval_base_branch: str,
    aipa_branch_created: bool,
    logger_: logging.Logger | None = None,
    notify: Callable[[str], Awaitable[None]] | None = None,
    skip_cleanup_on_main: bool = True,
    aipa_git_client: Any = None,  # Deprecated, kept for compatibility
) -> None:
    """
    Finalize aipa task lifecycle safely.

    Steps:
    1) Kill running eval subprocess if any.
    2) Cancel all running workflows on the eval branch.
    3) Cleanup temporary eval branch if created.
    """
    log = logger_ or logger

    # 1) Kill any running eval subprocess
    try:
        if github_eval_client is not None and hasattr(github_eval_client, "kill_subprocess"):
            github_eval_client.kill_subprocess()
    except Exception:
        log.debug("Failed to kill eval subprocess during cleanup.", exc_info=True)

    should_cleanup = True
    if skip_cleanup_on_main and aipa_eval_branch == "main":
        should_cleanup = False
        log.info("Skipping cleanup for main branch: %s", aipa_eval_branch)
        if notify is not None:
            await notify(f"Skipping cleanup for main branch: {aipa_eval_branch}")

    if not should_cleanup:
        return

    # 2) Cancel all running workflows on active branches
    # Use _active_branches + _last_branch from the eval client when available (more accurate
    # than aipa_eval_branch which is the session branch, not the candidate branch).
    branches_to_cancel: set[str] = set()
    if github_eval_client is not None:
        active = getattr(github_eval_client, '_active_branches', set())
        last = getattr(github_eval_client, '_last_branch', None)
        branches_to_cancel = set(active)
        if last:
            branches_to_cancel.add(last)
    if not branches_to_cancel:
        branches_to_cancel = {aipa_eval_branch}
    for branch in branches_to_cancel:
        try:
            log.info("Cancelling all running workflows on branch: %s", branch)
            cancel_result = await github_eval_client.cancel_all_runs_on_branch(branch)
            if cancel_result.get("cancelled_count", 0) > 0:
                log.info("Cancelled %s workflow runs on branch %s", cancel_result["cancelled_count"], branch)
            if cancel_result.get("errors"):
                log.warning("Cancel errors on branch %s: %s", branch, cancel_result["errors"])
        except Exception:
            log.warning("Failed to cancel running workflows on branch: %s", branch, exc_info=True)

    # 3) Cleanup temporary branch
    if aipa_branch_created:
        await cleanup_aipa_eval_branch(
            aipa_git_client,
            branch_name=aipa_eval_branch,
            base_branch=aipa_eval_base_branch,
            logger_=log,
        )


def _extract_best_candidate(node_data: dict[str, Any]) -> dict[str, Any]:
    """Extract best candidate in one iteration node."""
    candidates = node_data.get("candidates")
    if isinstance(candidates, dict) and candidates:
        best_candidate: dict[str, Any] | None = None
        best_score: float | None = None
        best_candidate_id: str | None = None
        for cand_id, cand in candidates.items():
            if not isinstance(cand, dict):
                continue
            cand_score = cand.get("overall_score")
            if isinstance(cand_score, (int, float)):
                if best_score is None or cand_score > best_score:
                    best_score = float(cand_score)
                    best_candidate = cand
                    best_candidate_id = str(cand_id)
        if best_candidate is not None:
            return {
                "overall_score": best_candidate.get("overall_score"),
                "branch": best_candidate.get("branch") or "",
                "candidate_id": best_candidate_id or "",
            }
    return {
        "overall_score": node_data.get("overall_score"),
        "branch": node_data.get("analysis_branch") or "",
        "candidate_id": "",
    }


def _find_best_branch(iteration: dict[str, Any]) -> tuple[str | None, str | None, float | None]:
    """Find best branch across iteration tree."""
    best_node_id: str | None = None
    best_branch: str | None = None
    best_score: float | None = None
    for node_id, node_data in iteration.items():
        if not isinstance(node_data, dict):
            continue
        candidate_info = _extract_best_candidate(node_data)
        overall_score = candidate_info.get("overall_score")
        if isinstance(overall_score, (int, float)):
            score = float(overall_score)
            if best_score is None or score > best_score:
                best_score = score
                best_node_id = str(node_id)
                best_branch = str(candidate_info.get("branch") or "")
    return best_node_id, best_branch, best_score


async def submit_leaderboard_job(
    *,
    enabled: bool,
    github_eval_client: Any,
    aipa_eval_branch: str,
    iteration: dict[str, Any] | None,
    logger_: logging.Logger | None = None,
    notify: Callable[[str], Awaitable[None]] | None = None,
) -> dict[str, Any]:
    """
    Submit leaderboard eval for the best optimized branch.

    Returns:
      {
        "triggered": bool,
        "dashboard_url": str | None,
        "reason": str
      }
    """
    log = logger_ or logger
    if not enabled:
        log.info("Skipping leaderboard eval (not enabled in config)")
        return {"triggered": False, "dashboard_url": None, "reason": "disabled"}

    if not isinstance(iteration, dict) or not iteration:
        log.info("No iteration data found, skipping leaderboard eval")
        return {"triggered": False, "dashboard_url": None, "reason": "no_iteration"}

    best_node_id, best_branch, best_score = _find_best_branch(iteration)
    if best_node_id == "root":
        log.info("Best node is root (baseline), skipping leaderboard eval")
        return {"triggered": False, "dashboard_url": None, "reason": "best_is_root"}
    if not best_branch:
        log.info("No best branch found in iteration tree, skipping leaderboard eval")
        return {"triggered": False, "dashboard_url": None, "reason": "no_best_branch"}

    log.info("Triggering leaderboard eval with best branch: %s (score: %s)", best_branch, best_score)
    if notify is not None:
        await notify(f"Running leaderboard evaluation with best performing branch: {best_branch}")
        encoded_branch = urllib.parse.quote(aipa_eval_branch, safe="")
        workflow_url = (
            "https://github.com/infinity-microsoft/picasso/actions/workflows/"
            f"aipa-evals-picasso-eval.yml?query=branch%3A{encoded_branch}"
        )
        await notify(f"You can check the status according to the latest running from this URL: {workflow_url}")

    eval_result = await github_eval_client.submit_eval(
        picasso_current_branch=best_branch,
        task_id=None,
        experiment_suffix=None,
        run_leaderboard=True,
    )
    dashboard_url: str | None = None
    if eval_result and "output_log" in eval_result:
        dashboard_result = await github_eval_client.parse_eval_dashboard(eval_result["output_log"])
        if dashboard_result.get("success"):
            dashboard_url = dashboard_result.get("dashboard_url")
            log.info("Leaderboard eval dashboard URL: %s", dashboard_url)
            if notify is not None and dashboard_url:
                await notify(f"Leaderboard evaluation completed! View results: {dashboard_url}")
        else:
            log.warning("Failed to parse Machine Evals Dashboard URL from leaderboard eval output")
            if notify is not None:
                await notify("Leaderboard evaluation triggered but dashboard URL not available")

    return {"triggered": True, "dashboard_url": dashboard_url, "reason": "submitted"}


def _extract_checkpoint_turn(path: Path) -> int:
    """Extract checkpoint turn number from filename."""
    parts = path.stem.split("_")
    if len(parts) < 2:
        return -1
    try:
        return int(parts[1])
    except ValueError:
        return -1


def load_latest_state_from_task_dir(task_dir: str | Path) -> dict[str, Any] | None:
    """
    Load latest task state from checkpoints or version files.

    Search order:
    1) {task_dir}/checkpoints/checkpoint_*.json
    2) {task_dir}/checkpoint_*.json (legacy)
    3) {task_dir}/version.jsonl
    4) {task_dir}/outputs/version.jsonl
    """
    task_path = Path(task_dir)
    checkpoint_candidates: list[Path] = []

    checkpoint_dir = task_path / "checkpoints"
    if checkpoint_dir.exists():
        checkpoint_candidates = sorted(
            checkpoint_dir.glob("checkpoint_*.json"),
            key=_extract_checkpoint_turn,
            reverse=True,
        )
    else:
        checkpoint_candidates = sorted(
            task_path.glob("checkpoint_*.json"),
            key=_extract_checkpoint_turn,
            reverse=True,
        )

    for checkpoint in checkpoint_candidates:
        try:
            payload = json.loads(checkpoint.read_text(encoding="utf-8"))
            state = payload.get("state")
            if isinstance(state, dict):
                return state
        except Exception:
            continue

    for version_file in (task_path / "version.jsonl", task_path / "outputs" / "version.jsonl"):
        if not version_file.exists():
            continue
        try:
            lines = [ln for ln in version_file.read_text(encoding="utf-8").splitlines() if ln.strip()]
            if not lines:
                continue
            payload = json.loads(lines[-1])
            if isinstance(payload, dict):
                return payload
        except Exception:
            continue
    return None


def collect_candidate_branches_from_state(
    state: dict[str, Any] | None,
    *,
    original_branch: str | None = None,
) -> set[str]:
    """Collect candidate branches from an iteration state snapshot."""
    if not isinstance(state, dict):
        return set()
    iteration_data = state.get("iteration")
    if not isinstance(iteration_data, dict):
        return set()

    branches: set[str] = set()
    for node_data in iteration_data.values():
        if not isinstance(node_data, dict):
            continue
        candidates = node_data.get("candidates")
        if not isinstance(candidates, dict):
            continue
        for candidate_data in candidates.values():
            if not isinstance(candidate_data, dict):
                continue
            branch = candidate_data.get("branch")
            if not isinstance(branch, str) or not branch:
                continue
            if original_branch and branch == original_branch:
                continue
            branches.add(branch)
    return branches


def collect_candidate_branches_from_task_dir(
    task_dir: str | Path,
    *,
    original_branch: str | None = None,
    repo_path: str | Path | None = None,
) -> set[str]:
    """Collect candidate branches from repo folder first, then task artifacts as fallback."""
    ga_pattern = re.compile(r"^autopromptlab/ga/[^_]+_\d+_[^_]+$")
    session_pattern = re.compile(r"^autopromptlab/session/[a-zA-Z0-9_-]{1,20}$")

    repo_candidates: list[Path] = []
    if repo_path:
        repo_candidates.append(Path(repo_path))
    task_path = Path(task_dir)
    repo_candidates.append(task_path / "picasso")
    repo_candidates.append(task_path)

    seen_repo_paths: set[Path] = set()
    for candidate in repo_candidates:
        candidate_resolved = candidate.resolve(strict=False)
        if candidate_resolved in seen_repo_paths:
            continue
        seen_repo_paths.add(candidate_resolved)

        if not candidate.exists():
            continue

        try:
            probe = subprocess.run(
                ["git", "rev-parse", "--is-inside-work-tree"],
                cwd=str(candidate),
                capture_output=True,
                text=True,
            )
            if probe.returncode != 0 or probe.stdout.strip().lower() != "true":
                continue

            result = subprocess.run(
                [
                    "git",
                    "for-each-ref",
                    "--format=%(refname:short)",
                    "refs/heads/autopromptlab/ga",
                    "refs/heads/autopromptlab/session",
                ],
                cwd=str(candidate),
                capture_output=True,
                text=True,
            )
        except Exception as e:
            logger.warning("Failed to list managed branches from repo %s: %s", candidate, e)
            continue

        if result.returncode != 0:
            logger.warning(
                "Failed to list managed branches from repo %s: %s",
                candidate,
                result.stderr.strip(),
            )
            continue

        branches: set[str] = set()
        for line in result.stdout.splitlines():
            branch = line.strip()
            if not branch:
                continue
            if original_branch and branch == original_branch:
                continue
            if ga_pattern.match(branch) or session_pattern.match(branch):
                branches.add(branch)
        return branches

    state = load_latest_state_from_task_dir(task_dir)
    return collect_candidate_branches_from_state(state, original_branch=original_branch)


async def delete_picasso_branches(
    *,
    git_client: Any,
    branches_to_delete: set[str],
    checkout_branch: str | None = None,
    logger_: logging.Logger | None = None,
) -> tuple[list[str], list[str], list[str]]:
    """Delete branches via GitClient.git_delete_branch()."""
    log = logger_ or logger
    deleted_branches: list[str] = []
    failed_branches: list[str] = []
    skipped_branches: list[str] = []

    if checkout_branch and checkout_branch not in branches_to_delete:
        try:
            checkout_result = await git_client.git_checkout(checkout_branch)
            if not checkout_result.get("ok"):
                log.warning(
                    "Failed to checkout branch %s before cleanup: %s",
                    checkout_branch,
                    checkout_result.get("error"),
                )
        except Exception as e:
            log.warning("Error checking out branch %s before cleanup: %s", checkout_branch, e)

    for branch in sorted(branches_to_delete):
        try:
            result = await git_client.git_delete_branch(branch)
            if result.get("skipped"):
                skipped_branches.append(branch)
                log.info("Skipped branch deletion %s: %s", branch, result.get("message"))
            elif result.get("ok"):
                deleted_branches.append(branch)
            else:
                failed_branches.append(branch)
                log.warning("Failed to delete branch %s: %s", branch, result.get("error"))
        except Exception as e:
            failed_branches.append(branch)
            log.warning("Error deleting branch %s: %s", branch, e)
    return deleted_branches, failed_branches, skipped_branches


def delete_task_directory(
    task_dir: str | Path,
    *,
    logger_: logging.Logger | None = None,
) -> bool:
    """Delete task directory recursively (Windows-safe)."""
    log = logger_ or logger
    path = Path(task_dir)
    if not path.exists():
        return False
    try:
        def _on_rm_error(func, target, exc_info):
            os.chmod(target, stat.S_IWRITE)
            func(target)

        shutil.rmtree(path, onerror=_on_rm_error)
        return True
    except Exception as e:
        log.warning("Failed to delete task directory %s: %s", path, e)
        return False


def delete_detached_session_repo(
    session_repo_path: str | Path | None,
    *,
    task_dir: str | Path,
    logger_: logging.Logger | None = None,
) -> bool:
    """
    Delete a disposable session repo only when it is not already covered by task_dir.

    Returns True when the repo path no longer exists after cleanup.
    """
    log = logger_ or logger
    if not session_repo_path:
        return False

    repo_path = Path(session_repo_path)
    if not repo_path.exists():
        return False

    task_dir_path = Path(task_dir)
    repo_resolved = repo_path.resolve(strict=False)
    task_dir_resolved = task_dir_path.resolve(strict=False)

    # If the session repo lives inside task_dir, task_dir deletion already covers it.
    if repo_resolved == task_dir_resolved or task_dir_resolved in repo_resolved.parents:
        return not repo_path.exists()

    deleted = delete_task_directory(repo_path, logger_=log)
    if not deleted and repo_path.exists():
        log.warning("Failed to delete detached session repo %s", repo_path)
        return False
    return not repo_path.exists()


async def delete_task_artifacts(
    *,
    task_dir: str | Path,
    original_branch: str | None,
    session_repo_path: str | Path | None = None,
    delete_branches: bool = True,
    git_client: Any | None = None,
    branches_to_delete: set[str] | None = None,
    logger_: logging.Logger | None = None,
) -> dict[str, Any]:
    """
    Delete task artifacts (task directory + generated branches).

    Returns:
      {
        "task_dir_deleted": bool,
        "session_repo_deleted": bool,
        "deleted_branches": list[str],
        "failed_branches": list[str],
        "skipped_branches": list[str],
      }
    """
    log = logger_ or logger
    branches = (
        set(branches_to_delete)
        if isinstance(branches_to_delete, set)
        else collect_candidate_branches_from_task_dir(
            task_dir,
            original_branch=original_branch,
            repo_path=session_repo_path,
        )
    )

    deleted_branches: list[str] = []
    failed_branches: list[str] = []
    skipped_branches: list[str] = []
    if delete_branches and branches:
        if git_client is None:
            log.warning("No git client provided; skip deleting %s branches.", len(branches))
        else:
            deleted_branches, failed_branches, skipped_branches = await delete_picasso_branches(
                git_client=git_client,
                branches_to_delete=branches,
                checkout_branch=original_branch,
                logger_=log,
            )

    task_dir_deleted = delete_task_directory(task_dir, logger_=log)
    session_repo_deleted = delete_detached_session_repo(
        session_repo_path,
        task_dir=task_dir,
        logger_=log,
    )
    return {
        "task_dir_deleted": task_dir_deleted,
        "session_repo_deleted": session_repo_deleted,
        "deleted_branches": deleted_branches,
        "failed_branches": failed_branches,
        "skipped_branches": skipped_branches,
    }
