import asyncio
import contextlib
import json
import os
from collections.abc import Mapping
from pathlib import Path
from typing import Any, NoReturn, Optional

import click

from synth_ai.cli.lib.train_cfgs import find_train_cfgs_in_cwd, validate_train_cfg
from synth_ai.core.config.errors import format_error_message, get_required_value
from synth_ai.core.env import get_synth_and_env_keys, mask_str
from synth_ai.core.paths import print_paths_formatted
from synth_ai.core.telemetry import flush_logger, log_error, log_info
from synth_ai.core.urls import BACKEND_URL_API
from synth_ai.data import extract_outcome_reward
from synth_ai.sdk.api.train.builders import (
    build_prompt_learning_payload,
    build_rl_payload,
    build_sft_payload,
)
from synth_ai.sdk.api.train.context_learning import ContextLearningJob
from synth_ai.sdk.api.train.graphgen import GraphGenJob
from synth_ai.sdk.api.train.graphgen_models import load_graphgen_taskset
from synth_ai.sdk.api.train.local_api import check_local_api_health
from synth_ai.sdk.api.train.utils import (
    TrainError,
    ensure_api_base,
    http_get,
    http_post,
    limit_jsonl_examples,
    mask_value,
    post_multipart,
    preview_json,
    sleep,
    validate_sft_jsonl,
)
from synth_ai.sdk.streaming import (
    CLIHandler,
    GraphGenHandler,
    JobStreamer,
    LossCurveHandler,
    PromptLearningHandler,
    StreamConfig,
    StreamEndpoints,
    StreamType,
)

# Constants for prompt learning event types
_PROMPT_LEARNING_EVENT_BEST_PROMPT = "prompt.learning.best.prompt"
_PROMPT_LEARNING_EVENT_FINAL_RESULTS = "prompt.learning.final.results"
_PROMPT_LEARNING_EVENT_VALIDATION_SCORED = "prompt.learning.validation.scored"
_PROMPT_LEARNING_EVENT_GEPA_COMPLETE = "prompt.learning.gepa.complete"
_PROMPT_LEARNING_EVENT_MIPRO_COMPLETE = "prompt.learning.mipro.complete"
_PROMPT_LEARNING_EVENT_GEPA_NEW_BEST = "prompt.learning.gepa.new_best"
_PROMPT_LEARNING_EVENT_PHASE_CHANGED = "prompt.learning.phase.changed"
_PROMPT_LEARNING_EVENT_PROGRESS = "prompt.learning.progress"
_PROMPT_LEARNING_EVENT_STREAM_CONNECTED = "prompt.learning.stream.connected"

# Constants for formatting
_MAX_TEXT_REPLACEMENTS_DISPLAY = 3  # Max number of text replacements to show in output
_RESULTS_FILE_MAX_EVENTS = 10000  # Max events to fetch for results file generation


def _format_text_replacements(
    obj: dict[str, Any] | None, max_display: int = _MAX_TEXT_REPLACEMENTS_DISPLAY
) -> list[str]:
    """Extract and format text replacements from a candidate object.

    Args:
        obj: Candidate object dictionary containing text_replacements
        max_display: Maximum number of replacements to display

    Returns:
        List of formatted lines showing role and replacement text
    """
    lines = []
    if not obj or not isinstance(obj, dict):
        return lines

    text_replacements = obj.get("text_replacements", [])
    if not text_replacements or not isinstance(text_replacements, list):
        return lines

    for replacement in text_replacements[:max_display]:
        if isinstance(replacement, dict):
            new_text = replacement.get("new_text", "")
            role = replacement.get("apply_to_role", "system")
            if new_text:
                lines.append(f"  [{role.upper()}]: {new_text}")
                lines.append("")

    return lines


def _extract_reward_value(
    payload: Any, *, fallback_keys: Optional[list[str]] = None
) -> Optional[float]:
    if not isinstance(payload, dict):
        return None
    for key in ("outcome_objectives", "objectives"):
        objectives = payload.get(key)
        if isinstance(objectives, dict):
            reward_val = objectives.get("reward")
            if isinstance(reward_val, (int, float)) and not isinstance(reward_val, bool):
                return float(reward_val)
    reward_val = payload.get("reward")
    if isinstance(reward_val, (int, float)) and not isinstance(reward_val, bool):
        return float(reward_val)
    reward_val = extract_outcome_reward(payload)
    if reward_val is not None and not isinstance(reward_val, bool):
        return float(reward_val)
    if fallback_keys:
        for key in fallback_keys:
            value = payload.get(key)
            if isinstance(value, (int, float)) and not isinstance(value, bool):
                return float(value)
    return None


def _default_backend() -> str:
    """Return the default backend URL from centralized urls module."""
    return BACKEND_URL_API


_DEFAULT_SFT_HIDDEN_EVENTS = {
    "sft.created",
    "sft.pricing.check.requested",
    "sft.pricing.check.allowed",
    "sft.stage",
    "snapshot.fetch",
    "hatchet.preflight",
    "hatchet.submission.attempt",
    "hatchet.submission.result",
    "sft.running",
    "sft.status",
    "sft.worker.alive",
    "sft.dispatch.selected",
    "sft.config.prepared",
    "sft.strategy.selected",
    "sft.training.args",
}

_DEFAULT_RL_HIDDEN_SUBSTRINGS = {"modal", "hatchet"}

_DEFAULT_PROMPT_LEARNING_HIDDEN_EVENTS = {
    "prompt.learning.policy.tokens",
    "mipro.bootstrap.progress",  # Hide individual bootstrap seed scores
    "mipro.tpe.rankings",  # Hide verbose TPE rankings
    "mipro.tpe.selected",  # Hide TPE selection details
    "mipro.tpe.update",  # Hide TPE density updates
    "mipro.trial.duplicate",  # Hide duplicate trial messages
    "mipro.trial.started",  # Hide individual trial start messages (too verbose with instructions)
    "mipro.trial.minibatch",  # Hide minibatch completion (only show full eval)
    "mipro.trial.complete",  # Hide individual trial completion
    "mipro.iteration.skip_generation",  # Hide skip generation messages
    "mipro.budget.update",  # Hide verbose budget updates (progress handler shows summary)
    "mipro.instruction.proposed",  # Hide proposed instructions (shown in results/logs only)
    "gepa.transformation.proposed",  # Hide proposed transformations (shown in results/logs only)
    # Note: mipro.stage_proposer.called is shown so users know instruction generation is happening
}


def _load_toml_config(config_path: Path) -> dict[str, Any]:
    """Load TOML config file."""
    try:
        import tomli  # type: ignore[import-untyped]
    except ImportError:
        # Fallback to tomllib for Python 3.11+
        try:
            import tomllib as tomli
        except ImportError:
            return {}

    try:
        with open(config_path, "rb") as f:
            return tomli.load(f)
    except Exception:
        return {}


def parse_results_folder(config_path: Path) -> Path:
    """Parse results_folder from TOML config and validate it exists.

    Checks both [prompt_learning] and top-level sections.
    Raises ClickException if missing or invalid.
    """
    config = _load_toml_config(config_path)

    # Check prompt_learning section first
    pl_section = config.get("prompt_learning", {})
    if isinstance(pl_section, dict):
        results_folder = pl_section.get("results_folder")
        if results_folder:
            results_folder_str = str(results_folder).strip()
            # Resolve relative to config file's directory if path is relative
            if not Path(results_folder_str).is_absolute():
                config_dir = config_path.parent.resolve()
                results_path = (config_dir / results_folder_str).resolve()
            else:
                results_path = Path(results_folder_str).expanduser().resolve()

            # Validate that the folder exists or can be created
            try:
                results_path.mkdir(parents=True, exist_ok=True)
            except (OSError, PermissionError) as e:
                raise click.ClickException(
                    f"Could not create results folder: {results_path}\n"
                    f"  Error: {e}\n"
                    f"  Config: {config_path}\n"
                    f"  TOML results_folder: {results_folder}"
                ) from e

            return results_path

    # Check top-level section
    results_folder = config.get("results_folder")
    if results_folder:
        results_folder_str = str(results_folder).strip()
        # Resolve relative to config file's directory if path is relative
        if not Path(results_folder_str).is_absolute():
            config_dir = config_path.parent.resolve()
            results_path = (config_dir / results_folder_str).resolve()
        else:
            results_path = Path(results_folder_str).expanduser().resolve()

        # Validate that the folder exists or can be created
        try:
            results_path.mkdir(parents=True, exist_ok=True)
        except (OSError, PermissionError) as e:
            raise click.ClickException(
                f"Could not create results folder: {results_path}\n"
                f"  Error: {e}\n"
                f"  Config: {config_path}\n"
                f"  TOML results_folder: {results_folder}"
            ) from e

        return results_path

    # Missing - raise error
    raise click.ClickException(
        f"Missing required 'results_folder' field in TOML config: {config_path}\n"
        f"  Please add 'results_folder = \"path/to/results\"' to [prompt_learning] section or top-level.\n"
        f"  Paths can be relative (to config file directory) or absolute."
    )


def parse_display_config(config_path: Path) -> dict[str, Any]:
    """Parse [display] section from TOML config."""
    config = _load_toml_config(config_path)
    display_section = config.get("display", {})

    # Also extract termination_config for max limits
    termination_section = config.get("termination_config", {})
    # Also check prompt_learning.termination_config
    pl_section = config.get("prompt_learning", {})
    if isinstance(pl_section, dict):
        pl_termination = pl_section.get("termination_config", {})
        if isinstance(pl_termination, dict):
            # Merge with top-level termination_config (top-level takes precedence)
            termination_section = {**pl_termination, **termination_section}

    return {
        "local_backend": display_section.get("local_backend", False),
        "tui": display_section.get("tui", False),
        "show_curve": display_section.get("show_curve", True),
        "verbose_summary": display_section.get("verbose_summary", True),
        "show_trial_results": display_section.get("show_trial_results", True),
        "show_transformations": display_section.get("show_transformations", False),
        "show_validation": display_section.get("show_validation", True),
        "max_tokens": termination_section.get("max_tokens"),
        "max_time_seconds": termination_section.get("max_time_seconds"),
        "max_rollouts": termination_section.get("max_rollouts"),
    }


def _build_stream_components(
    stream_format: str,
    *,
    hidden_event_types: set[str] | None = None,
    hidden_event_substrings: set[str] | None = None,
) -> tuple[StreamConfig, list]:
    """Return stream configuration and handlers for the requested format."""
    if stream_format == "chart":
        config = StreamConfig(
            enabled_streams={StreamType.STATUS, StreamType.EVENTS, StreamType.METRICS},
            event_types={
                "sft.progress",
                "sft.training.started",
                "sft.training.finish",
                "sft.validation.summary",
                "rl.train.step",
                "rl.train.started",
                "rl.train.completed",
                "workflow.completed",
                "workflow.failed",
            },
            metric_names={"train.loss"},
        )
        handlers = [LossCurveHandler()]
    else:
        config = StreamConfig.default()
        handlers = [
            CLIHandler(
                hidden_event_types=hidden_event_types or set(),
                hidden_event_substrings=hidden_event_substrings or set(),
            )
        ]
    return config, handlers


def _validate_openai_key_if_provider_is_openai(cfg_path: Path) -> None:
    """Validate that OPENAI_API_KEY is set if the provider is OpenAI.

    For prompt learning jobs, checks if policy.provider is 'openai' and raises
    a ClickException if OPENAI_API_KEY is not set in the environment.
    """
    cfg = _load_toml_config(cfg_path)

    # Check prompt_learning section
    pl_section = cfg.get("prompt_learning", {})
    if not isinstance(pl_section, dict):
        return

    policy = pl_section.get("policy", {})
    if not isinstance(policy, dict):
        return

    provider = policy.get("provider", "").lower()

    if provider == "openai":
        openai_key = os.environ.get("OPENAI_API_KEY", "").strip()
        if not openai_key:
            raise click.ClickException(
                "OPENAI_API_KEY is required when using provider='openai'.\n"
                "Please set OPENAI_API_KEY in your shell environment."
            )


# Module-level logging to track import and registration
import logging as _logging  # noqa: E402
import sys  # noqa: E402

_logger = _logging.getLogger(__name__)
_logger.debug("[TRAIN_MODULE] Module synth_ai.cli.train imported")


@click.command()
@click.argument("cfg_path", required=False, type=click.Path(exists=True, path_type=Path))
@click.option("--task-url", default=None, help="Override task app base URL (RL only)")
@click.option(
    "--dataset",
    "dataset_path",
    type=click.Path(),
    default=None,
    help="Override dataset JSONL path (SFT)",
)
@click.option("--model", default=None, help="Override model identifier")
@click.option(
    "--allow-experimental",
    "allow_experimental",
    is_flag=True,
    flag_value=True,
    default=None,
    help="Allow experimental models (overrides SDK_EXPERIMENTAL env)",
)
@click.option(
    "--no-allow-experimental",
    "allow_experimental",
    is_flag=True,
    flag_value=False,
    help="Disallow experimental models (overrides SDK_EXPERIMENTAL env)",
)
@click.option("--idempotency", default=None, help="Idempotency-Key header for job creation")
@click.option("--dry-run", is_flag=True, hidden=True, help="Deprecated: no-op")
@click.option("--poll/--no-poll", default=True, help="Poll job status until terminal state")
@click.option(
    "--poll-timeout", default=3600.0, type=float, help="Maximum seconds to poll before timing out"
)
@click.option("--poll-interval", default=5.0, type=float, help="Seconds between poll attempts")
@click.option(
    "--stream-format",
    type=click.Choice(["cli", "chart"]),
    default="cli",
    show_default=True,
    help="Streaming output style (cli = line updates, chart = live loss panel)",
)
@click.option(
    "--examples",
    "examples_limit",
    type=int,
    default=None,
    help="Limit SFT training to the first N examples",
)
@click.option(
    "--backend",
    "backend_override",
    default=None,
    help="Backend base URL (e.g., http://localhost:8000). Overrides BACKEND_BASE_URL env var.",
)
@click.option(
    "--local-backend",
    is_flag=True,
    default=None,
    help="Use local backend (localhost:8000). Overrides TOML [display].local_backend",
)
@click.option(
    "--tui",
    is_flag=True,
    default=None,
    help="Enable live TUI dashboard. Overrides TOML [display].tui",
)
@click.option(
    "--show-curve",
    is_flag=True,
    default=None,
    help="Show optimization curve at end. Overrides TOML [display].show_curve",
)
@click.option(
    "--verbose-summary",
    is_flag=True,
    default=None,
    help="Show detailed final summary. Overrides TOML [display].verbose_summary",
)
@click.option(
    "--type",
    "train_type_override",
    type=click.Choice(["prompt", "rl", "sft", "graphgen", "context_learning"]),
    default=None,
    help="Explicitly set training type. Required for GraphGen (uses JSON datasets).",
)
@click.option(
    "--rollout-budget",
    "rollout_budget",
    type=int,
    default=None,
    help="Rollout budget for GraphGen optimization (default: 100)",
)
@click.option(
    "--proposer-effort",
    "proposer_effort",
    type=click.Choice(["low", "medium", "high"]),
    default=None,
    help="Proposer effort level for GraphGen (default: medium)",
)
def train(
    cfg_path: Path | None,
    task_url: str | None,
    dataset_path: str | None,
    model: str | None,
    allow_experimental: bool | None,
    idempotency: str | None,
    dry_run: bool,
    poll: bool,
    poll_timeout: float,
    poll_interval: float,
    stream_format: str,
    examples_limit: int | None,
    backend_override: str | None,
    local_backend: bool | None,
    tui: bool | None,
    show_curve: bool | None,
    verbose_summary: bool | None,
    train_type_override: str | None,
    rollout_budget: int | None,
    proposer_effort: str | None,
) -> None:
    """Interactive launcher for RL / SFT / Prompt Learning / GraphGen / Context Learning jobs."""
    import traceback

    ctx: dict[str, Any] = {
        "cfg_path": str(cfg_path) if cfg_path else None,
        "poll": poll,
        "poll_timeout": poll_timeout,
        "poll_interval": poll_interval,
        "stream_format": stream_format,
        "backend_override": backend_override,
    }
    log_info("train_command invoked", ctx=ctx)

    # Wrap entire function in try-except to catch ALL exceptions
    try:
        # Log entry point IMMEDIATELY - this should always appear
        sys.stderr.write("[TRAIN_CMD] Starting train command\n")
        sys.stderr.flush()
        click.echo(f"[TRAIN_CMD] Args: cfg_path={cfg_path}, poll={poll}", err=True)
        click.echo(f"[TRAIN_CMD] Python executable: {sys.executable}", err=True)
        click.echo(f"[TRAIN_CMD] Working directory: {os.getcwd()}", err=True)

        # Handle GraphGen specially - it uses JSON datasets, not TOML configs
        if train_type_override == "graphgen":
            # For GraphGen, dataset_path is required and cfg_path is ignored
            if not dataset_path:
                raise click.ClickException(
                    "GraphGen requires --dataset flag with path to JSON dataset file.\n"
                    "Usage: synth-ai train --type graphgen --dataset my_tasks.json"
                )
            train_type = train_type_override
            click.echo(f"[TRAIN_CMD] GraphGen mode: using dataset {dataset_path}", err=True)
        else:
            # Non-GraphGen: use TOML config
            if not cfg_path:
                available_cfgs = find_train_cfgs_in_cwd()
                if len(available_cfgs) == 1:
                    train_type, cfg_path_str, _ = available_cfgs[0]
                    cfg_path = Path(cfg_path_str)
                    print(f"Automatically selected {train_type} training config at", cfg_path)
                else:
                    if len(available_cfgs) == 0:
                        print("No training config found in cwd.")
                        print("Validate your training config: synth-ai train-cfg check [CFG_PATH]")
                    else:
                        print("Multiple training configs found. Please specify which one to use:")
                        print_paths_formatted(available_cfgs)
                    print("Usage: synth-ai train --config [CFG_PATH]")
                    return None

            train_type = train_type_override or validate_train_cfg(cfg_path)

        synth_api_key, _ = get_synth_and_env_keys()

        # Resolve backend URL with priority: --backend flag > BACKEND_BASE_URL env > default
        if backend_override:
            # CLI flag takes highest precedence
            backend_base = ensure_api_base(backend_override.strip())
            click.echo(f"Backend base: {backend_base} (from --backend flag)")
        else:
            # Check BACKEND_BASE_URL after evaluating overrides
            backend_base_url_env = os.environ.get("BACKEND_BASE_URL", "").strip()
            backend_override_env = os.environ.get("BACKEND_OVERRIDE", "").strip()

            # Debug: Show what env vars are set
            click.echo(
                f"🔍 DEBUG: BACKEND_BASE_URL={backend_base_url_env or '(not set)'}", err=True
            )
            click.echo(
                f"🔍 DEBUG: BACKEND_OVERRIDE={backend_override_env or '(not set)'}", err=True
            )

            # Use _default_backend() to respect BACKEND_BASE_URL env var
            backend_raw = _default_backend()
            click.echo(f"🔍 DEBUG: _default_backend() returned: {backend_raw}", err=True)
            backend_base = ensure_api_base(backend_raw)

            # Assertion: Validate backend URL is what we expect
            if backend_base_url_env:
                expected_backend = ensure_api_base(backend_base_url_env)
                if backend_base != expected_backend:
                    raise click.ClickException(
                        f"Backend URL mismatch! Expected: {expected_backend}, Got: {backend_base}. "
                        f"BACKEND_BASE_URL={backend_base_url_env} but resolved to {backend_base}. "
                        f"This indicates BACKEND_BASE_URL is not being respected.\n"
                        f"💡 Solutions:\n"
                        f"   1. Set BACKEND_BASE_URL=http://localhost:8000 in your shell\n"
                        f"   2. Use --backend http://localhost:8000 flag (requires package rebuild)\n"
                        f"   3. Set BACKEND_OVERRIDE=http://localhost:8000 in your shell\n"
                        f"   4. Set SYNTH_BACKEND_URL_OVERRIDE=local and LOCAL_BACKEND_URL=http://localhost:8000"
                    )

            click.echo(f"Backend base: {backend_base} (key {mask_str(synth_api_key)})")
            if backend_base_url_env:
                click.echo(f"  (from BACKEND_BASE_URL={backend_base_url_env})")

        # Skip TOML-based validation for GraphGen (uses JSON datasets)
        if train_type != "graphgen" and cfg_path:
            _validate_openai_key_if_provider_is_openai(cfg_path)

        match train_type:
            case "prompt":
                if not cfg_path:
                    raise click.ClickException("Prompt Learning requires a TOML config file.")
                handle_prompt_learning(
                    cfg_path=cfg_path,
                    backend_base=backend_base,
                    synth_key=synth_api_key,
                    task_url_override=task_url,
                    allow_experimental=allow_experimental,
                    dry_run=dry_run,
                    poll=poll,
                    poll_timeout=poll_timeout,
                    poll_interval=poll_interval,
                    stream_format=stream_format,
                )
            case "context_learning":
                if not cfg_path:
                    raise click.ClickException(
                        "Context Learning requires a TOML config file.\n"
                        "Usage: synth-ai train --type context_learning --config my_context.toml"
                    )
                handle_context_learning(
                    cfg_path=cfg_path,
                    backend_base=backend_base,
                    synth_key=synth_api_key,
                    poll=poll,
                    stream_format=stream_format,
                )
            case "rl":
                if not cfg_path:
                    raise click.ClickException("RL requires a TOML config file.")
                handle_rl(
                    cfg_path=cfg_path,
                    backend_base=backend_base,
                    synth_key=synth_api_key,
                    task_url_override=task_url,
                    model_override=model,
                    idempotency=idempotency,
                    allow_experimental=allow_experimental,
                    dry_run=dry_run,
                    poll=poll,
                    poll_timeout=poll_timeout,
                    poll_interval=poll_interval,
                    stream_format=stream_format,
                )
            case "sft":
                if not cfg_path:
                    raise click.ClickException("SFT requires a TOML config file.")
                dataset_override_path = (
                    Path(dataset_path).expanduser().resolve() if dataset_path else None
                )
                handle_sft(
                    cfg_path=cfg_path,
                    backend_base=backend_base,
                    synth_key=synth_api_key,
                    dataset_override=dataset_override_path,
                    allow_experimental=allow_experimental,
                    dry_run=dry_run,
                    poll=poll,
                    poll_timeout=poll_timeout,
                    poll_interval=poll_interval,
                    stream_format=stream_format,
                    examples_limit=examples_limit,
                )
            case "graphgen":
                if not dataset_path:
                    raise click.ClickException("GraphGen requires a dataset path.")
                graphgen_dataset_path = Path(dataset_path).expanduser().resolve()
                handle_graphgen(
                    dataset_path=graphgen_dataset_path,
                    backend_base=backend_base,
                    synth_key=synth_api_key,
                    policy_models=model,
                    rollout_budget=rollout_budget,
                    proposer_effort=proposer_effort,
                    poll=poll,
                    poll_timeout=poll_timeout,
                    poll_interval=poll_interval,
                    stream_format=stream_format,
                )
    except Exception as e:
        ctx["error"] = type(e).__name__
        log_error("train_command failed", ctx=ctx)
        click.echo(f"[TRAIN_CMD] FATAL ERROR: {e}", err=True)
        traceback.print_exc(file=sys.stderr)
        raise
    finally:
        flush_logger()


def handle_context_learning(
    *,
    cfg_path: Path,
    backend_base: str,
    synth_key: str,
    poll: bool,
    stream_format: str,
) -> None:
    """Submit and stream a Context Learning job.

    Context Learning is SSE-first; polling flags are ignored.
    """
    if not poll:
        click.echo("Note: --no-poll is ignored for context learning (SSE streaming only).")

    click.echo("\n=== Submitting Context Learning Job ===")
    try:
        job = ContextLearningJob.from_config(
            cfg_path,
            backend_url=backend_base,
            api_key=synth_key,
        )
        result = job.submit()
    except Exception as e:
        raise click.ClickException(str(e)) from e

    click.echo("\n✓ Job created:")
    click.echo(f"  Context Learning Job ID: {result.job_id}")
    click.echo(f"  Status: {result.status}")

    click.echo("\n=== Streaming Job Progress ===")
    if stream_format == "chart":
        click.echo("Chart stream format is not supported for context learning; using CLI output.")

    try:
        final_status = job.stream_until_complete()
    except Exception as e:
        raise click.ClickException(str(e)) from e

    status = final_status.get("status") if isinstance(final_status, dict) else "unknown"
    click.echo(f"\nFinal status: {status}")
    click.echo(preview_json(final_status, limit=600))

    if status in {"succeeded", "completed"}:
        click.echo("\n=== Best Preflight Script ===")
        try:
            best = job.download_best_script()
            if best.preflight_script:
                click.echo(best.preflight_script[:2000])
                if len(best.preflight_script) > 2000:
                    click.echo(f"\n... (truncated, {len(best.preflight_script)} chars total)")
        except Exception as e:
            click.echo(f"⚠️  Could not download best script: {e}")


def _wait_for_training_file(
    backend_base: str, api_key: str, file_id: str, *, timeout: float = 10.0
) -> None:
    """Wait for training file to be visible after upload.

    Reduced from 120s to 10s because:
    - POST response already confirms file is uploaded
    - Backend now forces read-your-writes consistency
    - By job creation time, replica lag has resolved
    - Quick sanity check only, not critical path
    """
    url = f"{backend_base.rstrip('/')}/files/{file_id}"
    headers = {"Authorization": f"Bearer {api_key}"}
    elapsed = 0.0
    interval = 2.0
    first_check = True
    while True:
        resp = http_get(url, headers=headers, timeout=30.0)
        if resp.status_code == 200:
            try:
                data = resp.json()
            except json.JSONDecodeError:
                data = {}
            status = str(
                data.get("status") or data.get("state") or data.get("storage_state") or "ready"
            ).lower()
            if first_check:
                click.echo(f"File uploaded successfully (id={file_id}, status={status})")
                first_check = False
            if status in {"ready", "uploaded", "stored", "complete"}:
                click.echo(f"✓ Training file ready (status={status})")
                return
            # Show progress for processing states
            if status in {"processing", "pending", "validating"}:
                click.echo(
                    f"  Waiting for file processing... (status={status}, {elapsed:.0f}s elapsed)"
                )
        elif resp.status_code == 404:
            # Keep polling; object may not be visible yet
            if first_check:
                click.echo(f"Waiting for file {file_id} to become visible...")
                first_check = False
        elif resp.status_code in {401, 403}:
            # Auth errors won't resolve by polling - fail immediately
            try:
                error_body = resp.json()
            except json.JSONDecodeError:
                error_body = resp.text[:400]
            click.echo("\n[ERROR] Authentication failed when checking training file:")
            click.echo(f"  URL: {url}")
            click.echo(f"  Status: {resp.status_code}")
            click.echo(f"  Response: {error_body}")
            click.echo(f"  API key: {mask_value(api_key)}")
            raise click.ClickException(
                f"Authentication error ({resp.status_code}). "
                "Check that your SYNTH_API_KEY is valid and has permission to access this organization's files."
            )
        else:
            # Other errors - show details but keep polling
            try:
                error_body = resp.json()
            except json.JSONDecodeError:
                error_body = resp.text[:400]
            click.echo(f"[WARN] Unexpected response checking file {file_id}:")
            click.echo(f"  URL: {url}")
            click.echo(f"  Status: {resp.status_code}")
            click.echo(f"  Response: {error_body}")

        if elapsed >= timeout:
            raise click.ClickException(
                f"Training file {file_id} not ready after {timeout:.0f}s (last status: {resp.status_code})"
            )
        sleep(interval)
        elapsed += interval


def handle_rl(
    *,
    cfg_path: Path,
    backend_base: str,
    synth_key: str,
    task_url_override: str | None,
    model_override: str | None,
    idempotency: str | None,
    allow_experimental: bool | None,
    dry_run: bool,
    poll: bool,
    poll_timeout: float,
    poll_interval: float,
    stream_format: str,
) -> None:
    ctx: dict[str, Any] = {
        "cfg_path": str(cfg_path),
        "backend_base": backend_base,
        "task_url_override": task_url_override,
        "poll": poll,
    }
    log_info("handle_rl invoked", ctx=ctx)
    overrides: dict[str, Any] = {
        "backend": backend_base,
        "task_url": task_url_override,
        "model": model_override,
    }
    build = build_rl_payload(
        config_path=cfg_path,
        task_url=task_url_override or os.environ.get("TASK_APP_URL", ""),
        overrides=overrides,
        idempotency=idempotency,
        allow_experimental=allow_experimental,
    )

    # Backend-side verification: try ALL org environment keys against /health and /task_info
    verify_url = f"{backend_base}/rl/verify_task_app"
    verify_headers = {"Authorization": f"Bearer {synth_key}", "Content-Type": "application/json"}
    try:
        vresp = http_post(
            verify_url, headers=verify_headers, json_body={"endpoint_base_url": build.task_url}
        )
        try:
            parsed_json = vresp.json()
        except json.JSONDecodeError:
            parsed_json = None

        if isinstance(parsed_json, Mapping):
            vjs: dict[str, Any] = dict(parsed_json)
        else:
            vjs = {
                "status": vresp.status_code,
                "text": (vresp.text or "")[:400],
            }
            if parsed_json is not None:
                vjs["body"] = parsed_json
    except Exception as _ve:
        raise click.ClickException(
            f"Task app verification call failed: {type(_ve).__name__}: {_ve}"
        ) from _ve
    if vresp.status_code is not None and vresp.status_code >= 400:
        click.echo("Task app verification error:\n" + preview_json(vjs, limit=800))
        raise click.ClickException(f"Verification failed with status {vresp.status_code}")
    if not bool(vjs.get("any_ok")):
        click.echo("Task app verification failed; no auth combination succeeded. Full report:")
        click.echo(preview_json(vjs, limit=1200))
        raise click.ClickException("Task app verification failed (auth)")
    else:
        # Print concise summary
        try:
            cands = vjs.get("candidates_first15") or []
            attempts_raw = vjs.get("attempts")
            attempts: list[Mapping[str, Any]] = (
                [a for a in attempts_raw if isinstance(a, Mapping)]
                if isinstance(attempts_raw, list)
                else []
            )
            statuses = [attempt.get("status") for attempt in attempts]
            click.echo(f"Verification OK (candidates={cands}, statuses={statuses})")
        except (KeyError, ValueError, AttributeError):
            # Parsing verification summary failed, but verification itself succeeded
            click.echo("Verification OK")

    env_key = get_required_value(
        "environment_api_key",
        env_value=os.environ.get("ENVIRONMENT_API_KEY"),
    )
    os.environ["ENVIRONMENT_API_KEY"] = env_key

    click.echo("Performing task app health check…")
    health = check_local_api_health(build.task_url, env_key)
    if not health.ok:
        click.echo(f"Task app health check failed: {health.detail}")
        raise click.ClickException("Aborting due to failing health check")
    else:
        click.echo("Task app healthy")

    create_url = f"{backend_base}/rl/jobs"
    headers = {"Authorization": f"Bearer {synth_key}", "Content-Type": "application/json"}
    if build.idempotency:
        headers["Idempotency-Key"] = build.idempotency

    click.echo(f"POST {create_url}")
    click.echo("Payload preview:\n" + preview_json(build.payload, limit=800))

    resp = http_post(create_url, headers=headers, json_body=build.payload)
    try:
        js = resp.json()
    except json.JSONDecodeError as e:
        click.echo(f"⚠️  Failed to parse JSON response: {e}")
        js = {"status": resp.status_code, "text": resp.text[:400]}
    click.echo(f"Response {resp.status_code}: {preview_json(js, limit=400)}")
    if resp.status_code not in (200, 201):
        raise click.ClickException("Job creation failed")
    job_id = js.get("job_id") or js.get("id")
    if not job_id:
        raise click.ClickException("Response missing job id")

    if not poll:
        click.echo(f"Created job {job_id} (polling disabled)")
        return

    click.echo("\n=== Streaming Job Progress ===")

    # Enable metrics for prompt learning
    if stream_format == "chart":
        config = StreamConfig(
            enabled_streams={StreamType.STATUS, StreamType.EVENTS, StreamType.METRICS},
            event_types={
                "prompt.learning.progress",
                "prompt.learning.gepa.start",
                "prompt.learning.gepa.complete",
            },
            metric_names={"gepa.transformation.mean_reward"},
        )
        handlers = [LossCurveHandler()]
        click.echo("Using live chart (metric=gepa.transformation.mean_reward)")
    else:
        config = StreamConfig(
            enabled_streams={StreamType.STATUS, StreamType.EVENTS, StreamType.METRICS},
            metric_names={"gepa.transformation.mean_reward"},
        )
        handlers = [CLIHandler(hidden_event_substrings=_DEFAULT_RL_HIDDEN_SUBSTRINGS)]

    streamer = JobStreamer(
        base_url=backend_base,
        api_key=synth_key,
        job_id=job_id,
        endpoints=StreamEndpoints.rl(job_id),
        config=config,
        handlers=handlers,
        interval_seconds=poll_interval,
        timeout_seconds=poll_timeout,
    )
    final_status = asyncio.run(streamer.stream_until_terminal())
    click.echo(f"Final status: {final_status.get('status', 'unknown')}")
    click.echo(preview_json(final_status, limit=600))


def handle_sft(
    *,
    cfg_path: Path,
    backend_base: str,
    synth_key: str,
    dataset_override: Path | None,
    allow_experimental: bool | None,
    dry_run: bool,
    poll: bool,
    poll_timeout: float,
    poll_interval: float,
    stream_format: str,
    examples_limit: int | None,
) -> None:
    ctx: dict[str, Any] = {
        "cfg_path": str(cfg_path),
        "backend_base": backend_base,
        "dataset_override": str(dataset_override) if dataset_override else None,
        "poll": poll,
    }
    log_info("handle_sft invoked", ctx=ctx)
    try:
        build = build_sft_payload(
            config_path=cfg_path,
            dataset_override=dataset_override,
            allow_experimental=allow_experimental,
        )
    except TrainError as exc:
        _raise_sft_usage_error(exc)

    limited_path: Path | None = None

    try:
        if examples_limit is not None:
            limited_path = limit_jsonl_examples(build.train_file, examples_limit)
            click.echo(
                f"Using first {examples_limit} examples from {build.train_file} -> {limited_path}"
            )
            build.train_file = limited_path

        click.echo("Validating training dataset…")
        validate_sft_jsonl(build.train_file)
        if build.validation_file and build.validation_file.suffix == ".jsonl":
            click.echo("Validating validation dataset…")
            validate_sft_jsonl(build.validation_file)

        upload_url = f"{backend_base.rstrip('/')}/files"
        click.echo("\n=== Uploading Training Data ===")
        click.echo(f"Dataset: {build.train_file}")
        click.echo(f"Destination: {upload_url}")
        resp = post_multipart(
            upload_url, api_key=synth_key, file_field="file", file_path=build.train_file
        )
        js = (
            resp.json()
            if resp.headers.get("content-type", "").startswith("application/json")
            else {}
        )
        if resp.status_code is not None and resp.status_code >= 400 or "id" not in js:
            click.echo("\n[ERROR] Training file upload failed:")
            click.echo(f"  URL: {upload_url}")
            click.echo(f"  Status: {resp.status_code}")
            click.echo(f"  Response: {js or resp.text[:400]}")
            click.echo(f"  File: {build.train_file}")
            raise click.ClickException(
                f"Training file upload failed with status {resp.status_code}"
            )
        train_file_id = js["id"]
        click.echo(f"✓ Training file uploaded (id={train_file_id})")
        val_file_id = None
        if build.validation_file:
            click.echo(f"Uploading validation dataset: {build.validation_file}")
            vresp = post_multipart(
                upload_url,
                api_key=synth_key,
                file_field="file",
                file_path=build.validation_file,
            )
            vjs = (
                vresp.json()
                if vresp.headers.get("content-type", "").startswith("application/json")
                else {}
            )
            if vresp.status_code is not None and vresp.status_code < 400 and "id" in vjs:
                val_file_id = vjs["id"]
                click.echo(f"✓ Validation file uploaded (id={val_file_id})")
            else:
                click.echo(
                    f"[WARN] Validation upload failed ({vresp.status_code}): {vjs or vresp.text[:200]}"
                )
        payload = dict(build.payload)
        payload["training_file_id"] = train_file_id
        if val_file_id:
            payload.setdefault("metadata", {}).setdefault("effective_config", {}).setdefault(
                "data", {}
            )["validation_files"] = [val_file_id]

        click.echo("\n=== Checking File Processing Status ===")
        try:
            _wait_for_training_file(backend_base, synth_key, train_file_id)
        except click.ClickException as exc:
            click.echo(f"[WARN] File readiness check failed: {exc}")
            click.echo("Proceeding anyway - backend will validate file during job creation...")

        click.echo("\n=== Creating Training Job ===")
        click.echo("Job payload preview:")
        click.echo(preview_json(payload, limit=800))

        create_url = f"{backend_base}/learning/jobs"
        headers = {"Authorization": f"Bearer {synth_key}", "Content-Type": "application/json"}
        click.echo(f"\nPOST {create_url}")
        resp = http_post(create_url, headers=headers, json_body=payload)
        js = (
            resp.json()
            if resp.headers.get("content-type", "").startswith("application/json")
            else {}
        )
        if resp.status_code not in (200, 201):
            click.echo("\n[ERROR] Job creation failed:")
            click.echo(f"  URL: {create_url}")
            click.echo(f"  Status: {resp.status_code}")
            click.echo(f"  Response: {preview_json(js, limit=600)}")
            raise click.ClickException(f"Job creation failed with status {resp.status_code}")
        job_id = js.get("job_id") or js.get("id")
        if not job_id:
            raise click.ClickException("Response missing job id")
        click.echo(f"✓ Job created (id={job_id})")

        click.echo("\n=== Starting Training Job ===")
        start_url = f"{backend_base}/learning/jobs/{job_id}/start"
        click.echo(f"POST {start_url}")
        start_resp = http_post(start_url, headers=headers, json_body={})
        if start_resp.status_code not in (200, 201):
            click.echo(f"[WARN] Job start returned status {start_resp.status_code}")
        else:
            click.echo("✓ Job started")

        if not poll:
            click.echo(f"Started job {job_id} (polling disabled)")
            return

        click.echo("\n=== Streaming Job Progress ===")
        config, handlers = _build_stream_components(
            stream_format, hidden_event_types=_DEFAULT_SFT_HIDDEN_EVENTS
        )
        if stream_format == "chart":
            click.echo("Using live loss chart (metric=train.loss)")
        streamer = JobStreamer(
            base_url=backend_base,
            api_key=synth_key,
            job_id=job_id,
            endpoints=StreamEndpoints.learning(job_id),
            config=config,
            handlers=handlers,
            interval_seconds=poll_interval,
            timeout_seconds=poll_timeout,
        )
        final_status = asyncio.run(streamer.stream_until_terminal())
        status = final_status.get("status") if isinstance(final_status, dict) else "unknown"
        click.echo(f"Final status: {status}")
        click.echo(preview_json(final_status, limit=600))
    finally:
        if limited_path is not None:
            with contextlib.suppress(OSError):
                limited_path.unlink(missing_ok=True)
            # Clean up empty parent directory if possible
            with contextlib.suppress(OSError):
                limited_path.parent.rmdir()


def handle_graphgen(
    *,
    dataset_path: Path,
    backend_base: str,
    synth_key: str,
    policy_models: str | None,
    rollout_budget: int | None,
    proposer_effort: str | None,
    poll: bool,
    poll_timeout: float,
    poll_interval: float,
    stream_format: str,
) -> None:
    """Handle GraphGen workflow optimization job creation and streaming.

    GraphGen uses JSON dataset files and auto-generates task apps.
    """
    ctx: dict[str, Any] = {
        "dataset_path": str(dataset_path),
        "backend_base": backend_base,
        "poll": poll,
    }
    log_info("handle_graphgen invoked", ctx=ctx)

    # Load dataset
    click.echo(f"Loading GraphGen dataset from: {dataset_path}")
    try:
        dataset = load_graphgen_taskset(dataset_path)
    except FileNotFoundError as e:
        raise click.ClickException(f"Dataset file not found: {dataset_path}") from e
    except ValueError as e:
        raise click.ClickException(f"Invalid GraphGen dataset format: {e}") from e

    problem_spec = None
    try:
        raw_dataset = json.loads(dataset_path.read_text())
        problem_spec = raw_dataset.get("problem_spec") or raw_dataset.get("initial_prompt")
    except Exception:
        problem_spec = None

    click.echo(f"Dataset loaded: {dataset.metadata.name}")
    click.echo(f"  Tasks: {len(dataset.tasks)}")
    click.echo(f"  Gold outputs: {len(dataset.gold_outputs)}")
    click.echo(f"  Verifier mode: {dataset.verifier_config.mode}")

    # Create GraphGen job
    if not policy_models:
        raise click.ClickException("policy_models is required")
    job = GraphGenJob.from_dataset(
        dataset=dataset,
        policy_models=policy_models,
        rollout_budget=rollout_budget or 100,
        proposer_effort=proposer_effort or "medium",  # type: ignore
        problem_spec=problem_spec,
        backend_url=backend_base,
        api_key=synth_key,
        auto_start=True,
    )

    click.echo("\n=== Submitting GraphGen Job ===")
    click.echo(f"Policy models: {job.config.policy_models}")
    click.echo(f"Rollout budget: {job.config.rollout_budget}")
    click.echo(f"Proposer effort: {job.config.proposer_effort}")

    try:
        result = job.submit()
    except RuntimeError as e:
        raise click.ClickException(str(e)) from e

    click.echo("\n✓ Job created:")
    click.echo(f"  GraphGen Job ID: {result.graphgen_job_id}")
    click.echo(f"  Status: {result.status}")

    if not poll:
        click.echo(f"\nCreated job {result.graphgen_job_id} (polling disabled)")
        return

    click.echo("\n=== Streaming Job Progress ===")

    # Build stream handlers
    if stream_format == "chart":
        handlers = [LossCurveHandler()]
        click.echo("Using live loss chart (metric=gepa.transformation.mean_reward)")
    else:
        handlers = [GraphGenHandler()]

    # Stream until complete
    try:
        final_status = job.stream_until_complete(
            timeout=poll_timeout,
            interval=poll_interval,
            handlers=handlers,
        )
    except TimeoutError as e:
        raise click.ClickException(str(e)) from e

    status = final_status.get("status") if isinstance(final_status, dict) else "unknown"
    click.echo(f"\nFinal status: {status}")
    click.echo(preview_json(final_status, limit=600))

    # Download and display best prompt if succeeded
    if status == "succeeded" or status == "completed":
        click.echo("\n=== Best Optimized Prompt ===")
        try:
            prompt = job.download_prompt()
            if prompt:
                click.echo(prompt[:2000])
                if len(prompt) > 2000:
                    click.echo(f"\n... (truncated, {len(prompt)} chars total)")
        except Exception as e:
            click.echo(f"⚠️  Could not download prompt: {e}")


def _raise_sft_usage_error(exc: TrainError) -> NoReturn:
    message = str(exc).strip()
    lower_msg = message.lower()
    context = "Preparing SFT training job payload"
    impact = "Cannot submit training job without a valid dataset path"

    if "dataset not specified" in lower_msg:
        raise click.UsageError(
            format_error_message(
                summary="Dataset path required",
                context=context,
                problem="No dataset path was provided via config or CLI",
                impact=impact,
                solutions=[
                    (
                        'Add [job].data = "/path/to/data.jsonl" to the config',
                        "Persist the dataset path in the TOML file",
                    ),
                    (
                        "Re-run with --dataset /path/to/data.jsonl",
                        "Override the dataset path from the CLI",
                    ),
                    (
                        "Use an absolute path accessible from the current working directory",
                        "Relative paths are resolved from the shell cwd",
                    ),
                ],
            )
        ) from exc

    if "dataset not found" in lower_msg:
        raise click.UsageError(
            format_error_message(
                summary="Dataset path not found",
                context=context,
                problem=message,
                impact=impact,
                solutions=[
                    (
                        "Verify the dataset path exists on disk",
                        "Double-check spelling and that the file hasn't moved",
                    ),
                    (
                        "Provide an absolute path to the dataset file",
                        "Avoid relying on relative paths that resolve incorrectly",
                    ),
                    (
                        "Sync the dataset to this machine before running the CLI",
                        "Remote paths must be accessible locally",
                    ),
                ],
            )
        ) from exc

    raise click.ClickException(message) from exc


def _save_verbose_log_file(
    events: list[dict[str, Any]],
    log_file: Path,
    algorithm_name: str,
    job_id: str,
    append_summary: bool = False,
) -> None:
    """Save a verbose log file with all events in chronological order, including summary.

    If append_summary is True, only append the summary section (events were already streamed live).
    """
    import json
    from datetime import datetime

    try:
        lines = []
        if not append_summary:
            # Full log file with header and all events
            lines.append("=" * 80)
            lines.append(f"{algorithm_name} PROMPT LEARNING VERBOSE LOG")
            lines.append("=" * 80)
            lines.append(f"Job ID: {job_id}")
            lines.append(f"Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            lines.append(f"Total Events: {len(events)}")
            lines.append("=" * 80)
            lines.append("")

        # Sort events by timestamp if available
        def get_timestamp(event: dict[str, Any]) -> str:
            return event.get("timestamp", event.get("created_at", ""))

        sorted_events = sorted(events, key=get_timestamp)

        # Only include events if not appending summary (events were already streamed live)
        if not append_summary:
            for idx, event in enumerate(sorted_events, 1):
                if not isinstance(event, dict):
                    continue

                event_type = event.get("type", "unknown")
                timestamp = event.get("timestamp") or event.get("created_at", "")
                level = event.get("level", "info")
                message = event.get("message", "")
                data = event.get("data", {})

                lines.append(f"[{idx}] {timestamp} [{level.upper()}] {event_type}")
                if message:
                    lines.append(f"  Message: {message}")
                if data:
                    # Format data nicely (truncate very long values)
                    formatted_data = {}
                    for key, value in data.items():
                        if isinstance(value, dict | list):
                            # Convert to JSON string, truncate if too long
                            json_str = json.dumps(value, indent=2)
                            if len(json_str) > 1000:
                                json_str = json_str[:1000] + "... (truncated)"
                            formatted_data[key] = json_str
                        elif isinstance(value, str) and len(value) > 500:
                            formatted_data[key] = value[:500] + "... (truncated)"
                        else:
                            formatted_data[key] = value

                    if formatted_data:
                        lines.append(f"  Data: {json.dumps(formatted_data, indent=2)}")
                lines.append("")

        # Add summary table and chart at the end (always included)
        if append_summary:
            lines.append("\n\n")
        lines.append("=" * 80)
        lines.append("FINAL SUMMARY")
        lines.append("=" * 80)

        try:
            from synth_ai.sdk.api.train.summary import _generate_summary_text

            # Extract optimization curve from events
            optimization_curve = None
            trial_scores = []
            for event in sorted_events:
                if isinstance(event, dict):
                    event_type = event.get("type", "")
                    if event_type in ("prompt.learning.trial.complete", "mipro.new_incumbent"):
                        data = event.get("data", {})
                        trial_num = data.get("trial") or data.get("trial_num")
                        score = data.get("score") or data.get("minibatch_score")
                        if trial_num is not None and score is not None:
                            trial_scores.append((trial_num, score))

            if trial_scores:
                best_so_far = {}
                for trial_num, score in sorted(trial_scores):
                    if trial_num not in best_so_far or score > best_so_far[trial_num]:
                        best_so_far[trial_num] = score
                optimization_curve = sorted(best_so_far.items())

            summary_text, curve_text = _generate_summary_text(
                events=sorted_events,
                algorithm=algorithm_name.lower() if algorithm_name else None,
                optimization_curve=optimization_curve,
            )
            if summary_text:
                lines.append(summary_text)
            if curve_text:
                lines.append("")
                lines.append(curve_text)
        except Exception as e:
            lines.append(f"⚠️  Could not generate summary: {e}")

        lines.append("=" * 80)
        lines.append("END OF LOG")
        lines.append("=" * 80)

        # Write to file (append if summary-only mode)
        mode = "a" if append_summary else "w"
        with open(log_file, mode, encoding="utf-8") as f:
            if append_summary:
                f.write("\n")
            f.write("\n".join(lines))

    except Exception as e:
        click.echo(f"⚠️  Could not save verbose log file: {e}")


def _save_prompt_learning_results_locally(
    *,
    backend_base: str,
    api_key: str,
    job_id: str,
    config_path: Path,
    results_folder: Path,
) -> None:
    """Fetch events and generate results file locally after prompt learning completes."""
    from datetime import datetime

    try:
        # Fetch all events
        url = f"{backend_base}/prompt-learning/online/jobs/{job_id}/events?limit={_RESULTS_FILE_MAX_EVENTS}"
        headers = {"Authorization": f"Bearer {api_key}"}
        resp = http_get(url, headers=headers, timeout=30.0)

        if resp.status_code != 200:
            click.echo(
                f"⚠️  Could not fetch events to generate results file (status={resp.status_code})"
            )
            return

        data = resp.json()
        # Handle both list response (backend) and dict response (legacy compatibility)
        if isinstance(data, list):
            events = data
        elif isinstance(data, dict):
            events = data.get("events", [])
            if not isinstance(events, list):
                click.echo(f"⚠️  Events field is not a list: {type(events).__name__}")
                return
        else:
            click.echo(f"⚠️  Unexpected response type: {type(data).__name__}")
            return

        if not events:
            return

        # Extract key data from events
        best_score = None
        best_prompt = None
        baseline_score = None
        attempted_candidates = []
        optimized_candidates = []
        mipro_topk_candidates = []  # Collect MIPRO top-K candidates
        proposed_instructions = []  # Collect proposed instructions from MIPRO
        proposed_transformations = []  # Collect proposed transformations from GEPA

        for event in events:
            if not isinstance(event, dict):
                continue  # Skip malformed events

            event_type = event.get("type", "")
            event_data = event.get("data", {})
            if not isinstance(event_data, dict):
                event_data = {}  # Fallback to empty dict for safety

            if event_type == _PROMPT_LEARNING_EVENT_BEST_PROMPT:
                best_score = event_data.get("best_score")
                best_prompt = event_data.get("best_prompt")
            elif event_type == _PROMPT_LEARNING_EVENT_FINAL_RESULTS:
                attempted_candidates = event_data.get("attempted_candidates", [])
                optimized_candidates = event_data.get("optimized_candidates", [])
            elif event_type == _PROMPT_LEARNING_EVENT_VALIDATION_SCORED:
                # Check if this is the baseline by checking for is_baseline flag or baseline in message
                is_baseline = event_data.get("is_baseline", False)
                if not is_baseline:
                    msg = event.get("message", "")
                    is_baseline = "baseline" in msg.lower()
                if is_baseline:
                    baseline_score = _extract_reward_value(event_data)
            elif event_type == _PROMPT_LEARNING_EVENT_GEPA_COMPLETE and best_score is None:
                best_score = event_data.get("best_score")
            elif event_type == _PROMPT_LEARNING_EVENT_MIPRO_COMPLETE:
                # MIPRO completion event includes best_prompt and best_score
                if best_score is None:
                    best_score = event_data.get("best_score")
                if best_prompt is None:
                    best_prompt = event_data.get("best_prompt")
            elif event_type == "mipro.topk.evaluated":
                # Extract MIPRO top-K candidate data with full details
                rank = event_data.get("rank")
                train_score = event_data.get("train_score")
                test_score = event_data.get("test_score")
                if rank is not None and train_score is not None and test_score is not None:
                    # Extract full instruction text (may be multi-line)
                    instruction_text = event_data.get("instruction_text", "")
                    if not instruction_text:
                        # Try to get from instruction_lines if available
                        instruction_lines = event_data.get("instruction_lines", [])
                        if instruction_lines:
                            instruction_text = "\n".join(str(line) for line in instruction_lines)

                    mipro_topk_candidates.append(
                        {
                            "rank": rank,
                            "train_score": train_score,
                            "test_score": test_score,
                            "lift_absolute": event_data.get("lift_absolute"),
                            "lift_percent": event_data.get("lift_percent"),
                            "instruction_text": instruction_text,
                            "instruction_lines": event_data.get("instruction_lines", []),
                            "demo_indices": event_data.get("demo_indices", []),
                            "stage_payloads": event_data.get("stage_payloads", {}),
                            "instruction_indices": event_data.get("instruction_indices", []),
                            "test_per_seed": event_data.get("test_per_seed", {}),
                        }
                    )
            elif event_type == "mipro.baseline.test":
                # Extract baseline test score
                if baseline_score is None:
                    baseline_score = event_data.get("test_score")
            elif event_type == "mipro.instruction.proposed":
                # Collect proposed instructions
                proposed_instructions.append(
                    {
                        "iteration": event_data.get("iteration"),
                        "stage_id": event_data.get("stage_id"),
                        "module_id": event_data.get("module_id"),
                        "instruction_id": event_data.get("instruction_id"),
                        "instruction_text": event_data.get("instruction_text", ""),
                        "instruction_lines": event_data.get("instruction_lines", []),
                        "demo_indices": event_data.get("demo_indices", []),
                        "proposal_id": event_data.get("proposal_id"),
                        "timestamp": event.get("created_at"),
                    }
                )
            elif event_type == "gepa.transformation.proposed":
                # Collect proposed transformations
                proposed_transformations.append(
                    {
                        "generation": event_data.get("generation"),
                        "mutation_type": event_data.get("mutation_type"),
                        "operator": event_data.get("operator"),
                        "transformation_id": event_data.get("transformation_id"),
                        "parent_id": event_data.get("parent_id"),
                        "transformation_text": event_data.get("transformation_text", ""),
                        "transformation_dict": event_data.get("transformation_dict", {}),
                        "mutation_params": event_data.get("mutation_params", {}),
                        "timestamp": event.get("created_at"),
                    }
                )

        # Check if we have any results to display (best_prompt, best_score, or candidates)
        has_results = bool(
            attempted_candidates or optimized_candidates or best_prompt or best_score is not None
        )
        if not has_results:
            return

        # Determine algorithm name from events
        algorithm_name = "PROMPT LEARNING"
        for event in events:
            if isinstance(event, dict):
                event_type = event.get("type", "")
                if "gepa" in event_type.lower():
                    algorithm_name = "GEPA"
                    break
                elif "mipro" in event_type.lower():
                    algorithm_name = "MIPRO"
                    break

        # Generate formatted report
        lines = []
        lines.append("=" * 80)
        lines.append(f"{algorithm_name} PROMPT LEARNING RESULTS")
        lines.append("=" * 80)
        lines.append(f"Job ID: {job_id}")
        lines.append(f"Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        lines.append("")
        if baseline_score is not None:
            lines.append(f"📊 Baseline Score: {baseline_score:.4f} ({baseline_score * 100:.1f}%)")
        if best_score is not None:
            lines.append(f"🏆 Best Score:     {best_score:.4f} ({best_score * 100:.1f}%)")
        if baseline_score is not None and best_score is not None:
            improvement = (
                ((best_score - baseline_score) / baseline_score) * 100 if baseline_score > 0 else 0
            )
            lines.append(
                f"📈 Improvement:    {improvement:+.1f}% relative ({(best_score - baseline_score) * 100:+.1f} pp absolute)"
            )
        lines.append("=" * 80)
        lines.append("")

        # Add best prompt if available
        if best_prompt and isinstance(best_prompt, dict):
            lines.append("🏆 BEST PROMPT")
            lines.append("-" * 80)
            sections = best_prompt.get("sections", [])
            if not isinstance(sections, list):
                sections = []
            for sec in sections:
                if not isinstance(sec, dict):
                    continue
                role = sec.get("role", "unknown")
                content = sec.get("content", "")
                lines.append(f"\n[{role.upper()}]:")
                lines.append(content)
            lines.append("")

        # Add optimized candidates
        if optimized_candidates and isinstance(optimized_candidates, list):
            lines.append("=" * 80)
            lines.append(f"✨ TOP OPTIMIZED CANDIDATES ({len(optimized_candidates)})")
            lines.append("=" * 80)
            lines.append("")

            for idx, cand in enumerate(optimized_candidates):
                if not isinstance(cand, dict):
                    continue
                candidate_score = cand.get("score") or {}
                reward_val = _extract_reward_value(candidate_score)
                if reward_val is None:
                    reward_val = _extract_reward_value(cand)
                if reward_val is None:
                    reward_val = 0.0
                prompt_length = candidate_score.get("prompt_length", 0)
                payload_kind = cand.get("payload_kind", "unknown")

                # Try score.instance_scores first, then cand.instance_scores (explicit check)
                instance_scores = (
                    candidate_score.get("instance_scores")
                    if "instance_scores" in candidate_score
                    else cand.get("instance_scores")
                )
                n_eval = (
                    len(instance_scores)
                    if instance_scores and isinstance(instance_scores, list)
                    else 0
                )

                lines.append(
                    f"[{idx + 1}] Reward: {reward_val:.4f} | Length: {prompt_length} | Type: {payload_kind} | N: {n_eval}"
                )
                lines.append("-" * 80)

                obj = cand.get("object")
                if obj and isinstance(obj, dict) and payload_kind == "transformation":
                    # For transformations, text_replacements are nested in data
                    data_obj = obj.get("data", {})
                    replacement_lines = _format_text_replacements(data_obj)
                    lines.extend(replacement_lines)
                lines.append("")

        # Add MIPRO top-K candidates
        if mipro_topk_candidates and isinstance(mipro_topk_candidates, list):
            # Sort by rank
            mipro_topk_candidates.sort(key=lambda x: x.get("rank", 999))
            lines.append("=" * 80)
            lines.append(f"🎯 TOP-K CANDIDATES ({len(mipro_topk_candidates)})")
            lines.append("=" * 80)
            lines.append("")

            for cand in mipro_topk_candidates:
                rank = cand.get("rank", 0)
                train_score = cand.get("train_score", 0.0)
                test_score = cand.get("test_score", 0.0)
                lift_abs = cand.get("lift_absolute")
                lift_pct = cand.get("lift_percent")
                instruction_text = cand.get("instruction_text", "")
                instruction_lines = cand.get("instruction_lines", [])
                demo_indices = cand.get("demo_indices", [])
                instruction_indices = cand.get("instruction_indices", [])
                stage_payloads = cand.get("stage_payloads", {})
                test_per_seed = cand.get("test_per_seed", {})

                lift_str = ""
                if lift_abs is not None and lift_pct is not None:
                    lift_str = f" | Lift: {lift_abs:+.3f} ({lift_pct:+.1f}%)"

                lines.append(
                    f"[Rank {rank}] Train: {train_score:.4f} ({train_score * 100:.1f}%) | Test: {test_score:.4f} ({test_score * 100:.1f}%){lift_str}"
                )
                lines.append("-" * 80)

                # Show full instruction text (use instruction_lines if available, otherwise instruction_text)
                if instruction_lines:
                    lines.append("Instructions:")
                    for idx, instr_line in enumerate(instruction_lines, 1):
                        lines.append(f"  {idx}. {instr_line}")
                elif instruction_text:
                    # Split multi-line instructions
                    instr_parts = instruction_text.split("\n")
                    if len(instr_parts) > 1:
                        lines.append("Instructions:")
                        for idx, part in enumerate(instr_parts, 1):
                            if part.strip():
                                lines.append(f"  {idx}. {part.strip()}")
                    else:
                        lines.append(f"Instruction: {instruction_text}")

                if instruction_indices:
                    lines.append(f"Instruction Indices: {instruction_indices}")
                if demo_indices:
                    lines.append(f"Demo Indices: {demo_indices}")

                # Show per-stage breakdown if available
                if stage_payloads:
                    lines.append("Per-stage breakdown:")
                    for stage_id, payload in stage_payloads.items():
                        if isinstance(payload, dict):
                            instr_ids = payload.get("instruction_indices", [])
                            demo_ids = payload.get("demo_indices", [])
                            module_id = payload.get("module_id", "unknown")
                            lines.append(
                                f"  [{module_id}/{stage_id}] instr_ids={instr_ids} demo_ids={demo_ids}"
                            )

                # Show test per-seed scores if available
                if test_per_seed:
                    seed_scores = []
                    for seed, score in sorted(test_per_seed.items()):
                        seed_scores.append(f"{seed}: {score:.2f}")
                    if seed_scores:
                        lines.append(f"Test per-seed: {', '.join(seed_scores)}")

                lines.append("")

        # Add all proposal candidates
        if attempted_candidates and isinstance(attempted_candidates, list):
            lines.append("=" * 80)
            lines.append(f"💡 ALL PROPOSAL CANDIDATES ({len(attempted_candidates)})")
            lines.append("=" * 80)
            lines.append("")

            for idx, cand in enumerate(attempted_candidates):
                if not isinstance(cand, dict):
                    continue
                reward_val = _extract_reward_value(cand)
                if reward_val is None:
                    reward_val = 0.0
                prompt_length = cand.get("prompt_length", 0)
                tool_rate = cand.get("tool_call_rate", 0.0)
                instance_scores = cand.get("instance_scores", [])
                n_eval = len(instance_scores) if instance_scores else 0

                lines.append(
                    f"[{idx + 1}] Reward: {reward_val:.4f} | Length: {prompt_length} | Tool Rate: {tool_rate:.2f} | N: {n_eval}"
                )
                lines.append("-" * 80)

                obj = cand.get("object")
                if obj and isinstance(obj, dict):
                    # For proposals, text_replacements are at top level of object
                    replacement_lines = _format_text_replacements(obj)
                    lines.extend(replacement_lines)
                lines.append("")

        # Add proposed instructions section (MIPRO)
        if proposed_instructions and isinstance(proposed_instructions, list):
            lines.append("=" * 80)
            lines.append(f"💡 PROPOSED INSTRUCTIONS ({len(proposed_instructions)})")
            lines.append("=" * 80)
            lines.append("")

            for idx, instr in enumerate(proposed_instructions):
                if not isinstance(instr, dict):
                    continue
                iteration = instr.get("iteration", "?")
                stage_id = instr.get("stage_id", "?")
                module_id = instr.get("module_id", "?")
                instruction_id = instr.get("instruction_id", "?")
                instruction_text = instr.get("instruction_text", "")
                instruction_lines = instr.get("instruction_lines", [])
                demo_indices = instr.get("demo_indices", [])

                lines.append(
                    f"[{idx + 1}] Iteration {iteration} | Stage: {stage_id} | Module: {module_id} | ID: {instruction_id}"
                )
                if demo_indices:
                    lines.append(f"Demo Indices: {demo_indices}")
                lines.append("-" * 80)

                # Show instruction text (use instruction_lines if available, otherwise instruction_text)
                if instruction_lines:
                    for line_idx, line in enumerate(instruction_lines, 1):
                        if line.strip():
                            lines.append(f"  {line_idx}. {line.strip()}")
                elif instruction_text:
                    # Split multi-line instructions
                    instr_parts = instruction_text.split("\n")
                    if len(instr_parts) > 1:
                        for line_idx, part in enumerate(instr_parts, 1):
                            if part.strip():
                                lines.append(f"  {line_idx}. {part.strip()}")
                    else:
                        lines.append(f"  {instruction_text}")

                lines.append("")

        # Add proposed transformations section (GEPA)
        if proposed_transformations and isinstance(proposed_transformations, list):
            lines.append("=" * 80)
            lines.append(f"🧬 PROPOSED TRANSFORMATIONS ({len(proposed_transformations)})")
            lines.append("=" * 80)
            lines.append("")

            for idx, trans in enumerate(proposed_transformations):
                if not isinstance(trans, dict):
                    continue
                generation = trans.get("generation", "?")
                mutation_type = trans.get("mutation_type", "?")
                operator = trans.get("operator", "?")
                transformation_id = trans.get("transformation_id", "?")
                parent_id = trans.get("parent_id", "?")
                transformation_text = trans.get("transformation_text", "")
                transformation_dict = trans.get("transformation_dict", {})

                lines.append(
                    f"[{idx + 1}] Generation {generation} | Type: {mutation_type} | Operator: {operator}"
                )
                lines.append(f"Transformation ID: {transformation_id} | Parent ID: {parent_id}")
                lines.append("-" * 80)

                # Show transformation text
                if transformation_text:
                    lines.append("Transformation Text:")
                    lines.append(f"  {transformation_text}")

                # Show transformation dict details if available
                if transformation_dict:
                    text_replacements = transformation_dict.get("text_replacements", [])
                    if text_replacements:
                        lines.append("Text Replacements:")
                        for repl_idx, repl in enumerate(text_replacements, 1):
                            if isinstance(repl, dict):
                                apply_to = repl.get("apply_to_role", "unknown")
                                old_text = repl.get("old_text", "")[:100]
                                new_text = repl.get("new_text", "")[:200]
                                lines.append(f"  {repl_idx}. [{apply_to}]")
                                if old_text:
                                    lines.append(f"      Old: {old_text}...")
                                if new_text:
                                    lines.append(f"      New: {new_text}...")

                lines.append("")

        # Add summary table and chart before END OF REPORT
        lines.append("")
        lines.append("=" * 80)
        lines.append("FINAL SUMMARY")
        lines.append("=" * 80)

        # Generate summary table text (reuse summary.py logic)
        try:
            from synth_ai.sdk.api.train.summary import _generate_summary_text

            # Extract optimization curve from events if available
            optimization_curve = None
            # Try to extract curve from trial events
            trial_scores = []
            for event in events:
                if isinstance(event, dict):
                    event_type = event.get("type", "")
                    if event_type in ("prompt.learning.trial.complete", "mipro.new_incumbent"):
                        data = event.get("data", {})
                        trial_num = data.get("trial") or data.get("trial_num")
                        score = data.get("score") or data.get("minibatch_score")
                        if trial_num is not None and score is not None:
                            trial_scores.append((trial_num, score))

            if trial_scores:
                # Build optimization curve (best score so far at each trial)
                best_so_far = {}
                for trial_num, score in sorted(trial_scores):
                    if trial_num not in best_so_far or score > best_so_far[trial_num]:
                        best_so_far[trial_num] = score
                optimization_curve = sorted(best_so_far.items())

            summary_text, curve_text = _generate_summary_text(
                events=events,
                algorithm=algorithm_name.lower() if algorithm_name else None,
                optimization_curve=optimization_curve,
            )
            if summary_text:
                lines.append(summary_text)
            if curve_text:
                lines.append("")
                lines.append(curve_text)
        except Exception as e:
            lines.append(f"⚠️  Could not generate summary: {e}")

        lines.append("=" * 80)
        lines.append("END OF REPORT")
        lines.append("=" * 80)

        # Determine save location
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # Use results_folder from config (create if it doesn't exist)
        output_dir = results_folder
        output_dir.mkdir(parents=True, exist_ok=True)

        # Use algorithm-specific filename
        algorithm_prefix = algorithm_name.lower() if algorithm_name else "prompt_learning"
        output_file = output_dir / f"{algorithm_prefix}_results_{job_id}_{timestamp}.txt"

        with open(output_file, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))

        click.echo(f"\n📄 Results saved locally to: {output_file}")

        # Also save verbose log file with all events (append summary if log was streamed live)
        log_file = output_dir / f"{algorithm_prefix}_log_{job_id}_{timestamp}.log"
        append_summary = (
            log_file.exists()
        )  # If log file exists, it was streamed live, so just append summary
        _save_verbose_log_file(
            events, log_file, algorithm_name, job_id, append_summary=append_summary
        )
        click.echo(f"📋 Verbose log saved locally to: {log_file}")

    except (PermissionError, OSError) as e:
        click.echo(f"⚠️  Could not save results file locally: {e}")
    except Exception as e:
        click.echo(f"⚠️  Unexpected error saving results file: {e}")


def handle_prompt_learning(
    *,
    cfg_path: Path,
    backend_base: str,
    synth_key: str,
    task_url_override: str | None,
    allow_experimental: bool | None,
    dry_run: bool,
    poll: bool,
    poll_timeout: float,
    poll_interval: float,
    stream_format: str,
    display_config: dict[str, Any] | None = None,
    tui: bool = False,
    show_curve: bool = True,
    verbose_summary: bool = True,
) -> None:
    """Handle prompt learning job creation (MIPRO or GEPA)."""
    ctx: dict[str, Any] = {
        "cfg_path": str(cfg_path),
        "backend_base": backend_base,
        "task_url_override": task_url_override,
        "poll": poll,
    }
    log_info("handle_prompt_learning invoked", ctx=ctx)
    env_key = get_required_value(
        "environment_api_key",
        env_value=os.environ.get("ENVIRONMENT_API_KEY"),
    )
    os.environ["ENVIRONMENT_API_KEY"] = env_key

    overrides: dict[str, Any] = {
        "backend": backend_base,
        "task_url": task_url_override,
    }

    build = build_prompt_learning_payload(
        config_path=cfg_path,
        task_url=task_url_override,
        overrides=overrides,
        allow_experimental=allow_experimental,
    )

    # Assertion: Validate task app URL is reachable from backend perspective
    # If backend is localhost and task app is localhost, they should be able to communicate
    task_app_url = build.task_url or ""
    if backend_base.startswith("http://localhost") or backend_base.startswith("http://127.0.0.1"):
        if task_app_url.startswith("http://localhost") or task_app_url.startswith(
            "http://127.0.0.1"
        ):
            # Both are local - this should work
            pass
        else:
            click.echo(
                f"⚠️  WARNING: Backend is local ({backend_base}) but task app is remote ({task_app_url})"
            )
            click.echo(
                "   The backend may not be able to reach the task app. Consider using a tunnel or local task app."
            )

    click.echo("Performing task app health check…")
    click.echo(f"Task app URL: {build.task_url}")
    click.echo("⏳ Checking /health endpoint (timeout: 10s)...")
    health = check_local_api_health(build.task_url, env_key, timeout=10.0)
    if not health.ok:
        click.echo(f"❌ Task app health check failed: {health.detail}")
        click.echo(f"   Health status: {health.health_status}")
        click.echo(f"   Task info status: {health.task_info_status}")
        click.echo("💡 Troubleshooting:")
        click.echo("   1. Ensure the task app is running: lsof -i :8102")
        click.echo("   2. Test manually: curl -v http://127.0.0.1:8102/health")
        click.echo("   3. Check task app logs for errors")
        click.echo("   4. Restart the task app if it's hung")
        raise click.ClickException("Aborting due to failing health check")
    else:
        click.echo("Task app healthy")

    # Ensure backend_base has /api prefix
    if not backend_base.endswith("/api"):
        backend_base = ensure_api_base(backend_base)

    # Assertion: Validate backend URL before making request
    if not backend_base.startswith("http"):
        raise click.ClickException(
            f"Invalid backend URL: {backend_base}. Must start with http:// or https://"
        )

    create_url = f"{backend_base}/prompt-learning/online/jobs"
    headers = {"Authorization": f"Bearer {synth_key}", "Content-Type": "application/json"}

    click.echo(f"POST {create_url}")
    click.echo("Payload preview:\n" + preview_json(build.payload, limit=800))

    # Assertion: If using local backend, verify it's actually localhost
    if (
        os.getenv("BACKEND_BASE_URL")
        and "localhost" in os.getenv("BACKEND_BASE_URL", "").lower()
        and "localhost" not in backend_base.lower()
        and "127.0.0.1" not in backend_base
    ):
        raise click.ClickException(
            f"BACKEND_BASE_URL was set to localhost but backend_base resolved to {backend_base}. "
            f"This indicates the environment variable is not being respected."
        )

    # Increase timeout for job creation (can take longer due to validation checks)
    resp = http_post(create_url, headers=headers, json_body=build.payload, timeout=180.0)
    try:
        js = resp.json()
    except json.JSONDecodeError as e:
        click.echo(f"⚠️  Failed to parse JSON response: {e}")
        js = {"status": resp.status_code, "text": resp.text[:400]}
    click.echo(f"Response {resp.status_code}: {preview_json(js, limit=400)}")
    if resp.status_code not in (200, 201):
        raise click.ClickException("Job creation failed")
    job_id = js.get("job_id") or js.get("id")
    if not job_id:
        raise click.ClickException("Response missing job id")

    if not poll:
        click.echo(f"Created job {job_id} (polling disabled)")
        return

    algorithm = str(build.payload.get("algorithm") or "").lower()
    metric_names: set[str] | None = None
    if algorithm == "gepa":
        metric_names = {"gepa.transformation.mean_reward"}

    chart_mode = stream_format == "chart" and algorithm == "gepa"
    if stream_format == "chart" and not chart_mode:
        click.echo(
            "Chart streaming is only available for GEPA jobs; showing textual updates instead."
        )

    # Prepare log file path for real-time streaming
    results_folder = parse_results_folder(cfg_path)
    from datetime import datetime

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    algorithm_prefix = algorithm.lower() if algorithm else "prompt_learning"
    log_file = results_folder / f"{algorithm_prefix}_log_{job_id}_{timestamp}.log"

    # Write initial streaming message to log file if handler will be created
    if not chart_mode:
        try:
            log_file.parent.mkdir(parents=True, exist_ok=True)
            with open(log_file, "a", encoding="utf-8") as f:
                f.write("\n=== Streaming Job Progress ===\n")
        except Exception:
            pass  # Continue even if log file can't be written

    click.echo("\n=== Streaming Job Progress ===")

    # Create appropriate handler based on algorithm
    if algorithm == "gepa":
        if chart_mode:
            config = StreamConfig(
                enabled_streams={StreamType.STATUS, StreamType.EVENTS, StreamType.METRICS},
                event_types={
                    "prompt.learning.progress",
                    "prompt.learning.gepa.start",
                    "prompt.learning.gepa.complete",
                },
                metric_names=metric_names,
            )
            handlers = [LossCurveHandler()]
            click.echo("Using live loss chart (metric=gepa.transformation.mean_reward)")
        else:
            config = StreamConfig(
                enabled_streams={
                    StreamType.STATUS,
                    StreamType.EVENTS,
                    StreamType.METRICS,
                    StreamType.TIMELINE,
                },
                metric_names=metric_names,
                max_events_per_poll=500,  # Capture more events per poll
                deduplicate=True,  # Still deduplicate but capture more
                # Don't filter events - show all of them
                event_types=None,  # No whitelist - show all event types
                event_types_exclude=None,  # No blacklist - show all events
                event_levels=None,  # Show all levels
            )
            # Use PromptLearningHandler for enhanced event handling
            handler = PromptLearningHandler(
                show_trial_results=display_config.get("show_trial_results", True)
                if display_config
                else True,
                show_transformations=display_config.get("show_transformations", False)
                if display_config
                else False,
                show_validation=display_config.get("show_validation", True)
                if display_config
                else True,
                max_tokens=display_config.get("max_tokens") if display_config else None,
                max_time_seconds=display_config.get("max_time_seconds") if display_config else None,
                max_rollouts=display_config.get("max_rollouts") if display_config else None,
                log_file=log_file,
            )
            handlers = [handler]
    else:
        # Use PromptLearningHandler for MIPRO (same as GEPA)
        config = StreamConfig(
            enabled_streams={
                StreamType.STATUS,
                StreamType.EVENTS,
                StreamType.METRICS,
                StreamType.TIMELINE,
            },
            metric_names=metric_names,
            max_events_per_poll=500,  # Capture more events per poll
            deduplicate=True,  # Still deduplicate but capture more
            # Don't filter events - show all of them
            event_types=None,  # No whitelist - show all event types
            event_types_exclude=None,  # No blacklist - show all events
            event_levels=None,  # Show all levels
        )
        handler = PromptLearningHandler(
            show_trial_results=display_config.get("show_trial_results", True)
            if display_config
            else True,
            show_transformations=display_config.get("show_transformations", False)
            if display_config
            else False,
            show_validation=display_config.get("show_validation", True) if display_config else True,
            max_tokens=display_config.get("max_tokens") if display_config else None,
            max_time_seconds=display_config.get("max_time_seconds") if display_config else None,
            max_rollouts=display_config.get("max_rollouts") if display_config else None,
            log_file=log_file,
        )
        handlers = [handler]

    streamer = JobStreamer(
        base_url=backend_base,
        api_key=synth_key,
        job_id=job_id,
        endpoints=StreamEndpoints.prompt_learning(job_id),
        config=config,
        handlers=handlers,
        interval_seconds=poll_interval,
        timeout_seconds=poll_timeout,
    )
    final_status = asyncio.run(streamer.stream_until_terminal())

    # Write final status to log file if handler has one
    if isinstance(handlers[0], PromptLearningHandler) and handlers[0]._log_file_handle:
        handlers[0]._write_log(f"Final status: {final_status.get('status', 'unknown')}")
        handlers[0]._write_log(preview_json(final_status, limit=600))

    click.echo(f"Final status: {final_status.get('status', 'unknown')}")
    click.echo(preview_json(final_status, limit=600))

    # Display final summary for GEPA/MIPRO jobs if requested
    if verbose_summary and algorithm in ("gepa", "mipro"):
        optimization_curve = None
        if isinstance(handlers[0], PromptLearningHandler):
            optimization_curve = handlers[0].optimization_curve

        from synth_ai.sdk.api.train.summary import display_prompt_learning_summary

        # Pass log_writer if handler has one
        log_writer = None
        if isinstance(handlers[0], PromptLearningHandler) and handlers[0]._log_file_handle:
            log_writer = handlers[0]._write_log
        display_prompt_learning_summary(
            job_id=job_id,
            backend_base=backend_base,
            api_key=synth_key,
            optimization_curve=optimization_curve,
            show_curve=show_curve,
            algorithm=algorithm,
            log_writer=log_writer,
        )

    # Save results file locally
    # Parse and validate results_folder from config (already done above, but ensure it's available)
    if "results_folder" not in locals():
        results_folder = parse_results_folder(cfg_path)

    # Close log file if handler has one (flush is already called by streamer, but ensure it's closed)
    if isinstance(handlers[0], PromptLearningHandler) and handlers[0]._log_file_handle:
        handlers[0].flush()

    _save_prompt_learning_results_locally(
        backend_base=backend_base,
        api_key=synth_key,
        job_id=job_id,
        config_path=cfg_path,
        results_folder=results_folder,
    )
