import importlib
import os
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any, cast

import click
from pydantic import ValidationError

try:
    _models_module = cast(Any, importlib.import_module("synth_ai.sdk.api.models.supported"))
    UnsupportedModelError = cast(type[Exception], _models_module.UnsupportedModelError)
    ensure_allowed_model = cast(Callable[..., None], _models_module.ensure_allowed_model)
    normalize_model_identifier = cast(
        Callable[[str], str], _models_module.normalize_model_identifier
    )
except Exception as exc:  # pragma: no cover - critical dependency
    raise RuntimeError("Unable to load supported model helpers") from exc

try:
    _sft_module = cast(Any, importlib.import_module("synth_ai.sdk.learning.sft.config"))
    prepare_sft_job_payload = cast(
        Callable[..., dict[str, Any]], _sft_module.prepare_sft_job_payload
    )
except Exception as exc:  # pragma: no cover - critical dependency
    raise RuntimeError("Unable to load SFT payload helpers") from exc

from synth_ai.core.config.resolver import ConfigResolver
from synth_ai.core.telemetry import log_info
from synth_ai.sdk.localapi.auth import ensure_localapi_auth

from .configs import PromptLearningConfig, RLConfig, SFTConfig
from .supported_algos import (
    AlgorithmValidationError,
    ensure_model_supported_for_algorithm,
    validate_algorithm_config,
)
from .utils import TrainError, ensure_api_base


@dataclass(slots=True)
class RLBuildResult:
    payload: dict[str, Any]
    task_url: str
    idempotency: str | None


@dataclass(slots=True)
class SFTBuildResult:
    payload: dict[str, Any]
    train_file: Path
    validation_file: Path | None


@dataclass(slots=True)
class PromptLearningBuildResult:
    payload: dict[str, Any]
    task_url: str


def _format_validation_error(path: Path, exc: ValidationError) -> str:
    lines: list[str] = []
    for error in exc.errors():
        loc = ".".join(str(part) for part in error.get("loc", ()))
        msg = error.get("msg", "invalid value")
        lines.append(f"{loc or '<root>'}: {msg}")
    details = "\n".join(f"  - {line}" for line in lines) or "  - Invalid configuration"
    return f"Config validation failed ({path}):\n{details}"


def build_rl_payload(
    *,
    config_path: Path,
    task_url: str,
    overrides: dict[str, Any],
    idempotency: str | None,
    allow_experimental: bool | None = None,
) -> RLBuildResult:
    ctx: dict[str, Any] = {"config_path": str(config_path), "task_url": task_url}
    log_info("build_rl_payload invoked", ctx=ctx)
    # Load and validate config with SDK-level checks
    from synth_ai.sdk.api.train.utils import load_toml
    from synth_ai.sdk.api.train.validation.validation import validate_rl_config

    try:
        raw_config = load_toml(config_path)
        validated_config = validate_rl_config(raw_config)  # Adds defaults & validates
        rl_cfg = RLConfig.from_mapping(validated_config)
    except ValidationError as exc:
        raise click.ClickException(_format_validation_error(config_path, exc)) from exc

    data = rl_cfg.to_dict()

    # Remove smoke section - it's CLI-only and should not be sent to the trainer
    if "smoke" in data:
        del data["smoke"]

    # Ensure required [reference] section for backend validators
    try:
        ref_cfg = data.get("reference") if isinstance(data, dict) else None
        if not isinstance(ref_cfg, dict):
            data["reference"] = {"placement": "none"}
        else:
            ref_cfg.setdefault("placement", "none")
    except Exception:
        # Defensive: never fail builder due to optional defaults
        data["reference"] = {"placement": "none"}
    try:
        spec = validate_algorithm_config(rl_cfg.algorithm.model_dump(), expected_family="rl")
    except AlgorithmValidationError as exc:
        raise click.ClickException(str(exc)) from exc
    services = data.get("services") if isinstance(data.get("services"), dict) else {}
    model_cfg = rl_cfg.model

    cli_task_url = overrides.get("task_url")
    env_task_url = task_url or os.environ.get("TASK_APP_URL")
    config_task_url = services.get("task_url") if isinstance(services, dict) else None
    final_task_url = ConfigResolver.resolve(
        "task_app_url",
        cli_value=cli_task_url,
        env_value=env_task_url,
        config_value=config_task_url,
        required=True,
    )
    assert final_task_url is not None  # required=True guarantees non-None

    model_source = (model_cfg.source or "").strip() if model_cfg else ""
    model_base = (model_cfg.base or "").strip() if model_cfg else ""
    override_model = (overrides.get("model") or "").strip()
    if override_model:
        model_source = override_model
        model_base = ""
    if bool(model_source) == bool(model_base):
        details = (
            f"Config: {config_path}\n[model].source={model_source!r} | [model].base={model_base!r}"
        )
        hint = (
            "Set exactly one: [model].base for a base model (e.g. 'Qwen/Qwen3-1.7B') "
            "or [model].source for a fine-tuned model id. Also remove any conflicting "
            "'[policy].model' entries."
        )
        raise click.ClickException(
            "Invalid model config: exactly one of [model].source or [model].base is required.\n"
            + details
            + "\nHint: "
            + hint
        )

    try:
        if model_source:
            model_source = normalize_model_identifier(model_source)
        if model_base:
            model_base = normalize_model_identifier(model_base)
    except UnsupportedModelError as exc:
        raise click.ClickException(str(exc)) from exc

    base_model_for_training: str | None = None
    if model_source:
        base_model_for_training = ensure_allowed_model(
            model_source,
            allow_finetuned_prefixes=True,
            allow_experimental=allow_experimental,
        )
    elif model_base:
        base_model_for_training = ensure_allowed_model(
            model_base,
            allow_finetuned_prefixes=False,
            allow_experimental=allow_experimental,
        )
    if base_model_for_training:
        try:
            ensure_model_supported_for_algorithm(base_model_for_training, spec)
        except AlgorithmValidationError as exc:
            raise click.ClickException(str(exc)) from exc

    # Force TOML services.task_url to the effective endpoint to avoid split URLs
    try:
        if isinstance(data.get("services"), dict):
            data["services"]["task_url"] = final_task_url
        else:
            data["services"] = {"task_url": final_task_url}
    except Exception:
        pass

    payload_data: dict[str, Any] = {
        "endpoint_base_url": final_task_url.rstrip("/"),
        "config": data,
    }
    payload: dict[str, Any] = {
        "job_type": "rl",
        "compute": data.get("compute", {}),
        "data": payload_data,
        "tags": {"source": "train-cli"},
    }
    if model_source:
        payload_data["model"] = model_source
    if model_base:
        payload_data["base_model"] = model_base

    backend = overrides.get("backend")
    if backend:
        metadata_default: dict[str, Any] = {}
        metadata = cast(dict[str, Any], payload.setdefault("metadata", metadata_default))
        metadata["backend_base_url"] = ensure_api_base(str(backend))

    return RLBuildResult(payload=payload, task_url=final_task_url, idempotency=idempotency)


def build_sft_payload(
    *,
    config_path: Path,
    dataset_override: Path | None,
    allow_experimental: bool | None,
) -> SFTBuildResult:
    ctx: dict[str, Any] = {
        "config_path": str(config_path),
        "dataset_override": str(dataset_override) if dataset_override else None,
    }
    log_info("build_sft_payload invoked", ctx=ctx)
    try:
        sft_cfg = SFTConfig.from_path(config_path)
    except ValidationError as exc:
        raise TrainError(_format_validation_error(config_path, exc)) from exc

    data = sft_cfg.to_dict()
    try:
        algo_mapping = sft_cfg.algorithm.model_dump() if sft_cfg.algorithm else None
        spec = validate_algorithm_config(algo_mapping, expected_family="sft")
    except AlgorithmValidationError as exc:
        raise TrainError(str(exc)) from exc
    data_cfg = data.get("data") if isinstance(data.get("data"), dict) else {}
    hp_cfg = data.get("hyperparameters") if isinstance(data.get("hyperparameters"), dict) else {}
    train_cfg = data.get("training") if isinstance(data.get("training"), dict) else {}
    compute_cfg = data.get("compute") if isinstance(data.get("compute"), dict) else {}

    raw_dataset = dataset_override or sft_cfg.job.data or sft_cfg.job.data_path
    if not raw_dataset:
        raise TrainError("Dataset not specified; pass --dataset or set [job].data")
    dataset_path = Path(raw_dataset)
    # Resolve relative paths from current working directory, not config directory
    dataset_path = (
        dataset_path if dataset_path.is_absolute() else (Path.cwd() / dataset_path)
    ).resolve()
    if not dataset_path.exists():
        raise TrainError(f"Dataset not found: {dataset_path}")

    validation_path = (
        data_cfg.get("validation_path")
        if isinstance(data_cfg, dict)
        else None
        if isinstance(data_cfg, dict) and isinstance(data_cfg.get("validation_path"), str)
        else None
    )
    validation_file = None
    if validation_path:
        vpath = Path(validation_path)
        # Resolve relative paths from current working directory, not config directory
        vpath = (vpath if vpath.is_absolute() else (Path.cwd() / vpath)).resolve()
        if not vpath.exists():
            click.echo(f"[WARN] Validation dataset {vpath} missing; continuing without validation")
        else:
            validation_file = vpath

    hp_block: dict[str, Any] = {
        "n_epochs": int(hp_cfg.get("n_epochs", 1) if isinstance(hp_cfg, dict) else 1),
    }
    for key in (
        "batch_size",
        "global_batch",
        "per_device_batch",
        "gradient_accumulation_steps",
        "sequence_length",
        "learning_rate",
        "warmup_ratio",
        "train_kind",
    ):
        if isinstance(hp_cfg, dict) and key in hp_cfg:
            hp_block[key] = hp_cfg[key]
    if isinstance(hp_cfg, dict) and isinstance(hp_cfg.get("parallelism"), dict):
        hp_block["parallelism"] = hp_cfg["parallelism"]

    compute_block = {
        k: compute_cfg[k]
        for k in ("gpu_type", "gpu_count", "nodes")
        if isinstance(compute_cfg, dict) and k in compute_cfg
    }

    effective = {
        "compute": compute_block,
        "data": {
            "topology": data_cfg.get("topology", {})
            if isinstance(data_cfg, dict) and isinstance(data_cfg.get("topology"), dict)
            else {}
        },
        "training": {
            k: v
            for k, v in (train_cfg.items() if isinstance(train_cfg, dict) else [])
            if k in ("mode", "use_qlora")
        },
    }

    validation_cfg = (
        train_cfg.get("validation")
        if isinstance(train_cfg, dict) and isinstance(train_cfg.get("validation"), dict)
        else None
    )
    if isinstance(validation_cfg, dict):
        hp_block.update(
            {
                "evaluation_strategy": validation_cfg.get("evaluation_strategy", "steps"),
                "eval_steps": int(validation_cfg.get("eval_steps", 0) or 0),
                "save_best_model_at_end": bool(validation_cfg.get("save_best_model_at_end", True)),
                "metric_for_best_model": validation_cfg.get("metric_for_best_model", "val.loss"),
                "greater_is_better": bool(validation_cfg.get("greater_is_better", False)),
            }
        )
        effective.setdefault("training", {})["validation"] = {
            "enabled": bool(validation_cfg.get("enabled", True))
        }

    raw_model = (sft_cfg.job.model or "").strip()
    if not raw_model:
        model_block = data.get("model")
        if isinstance(model_block, str):
            raw_model = model_block.strip()
    if not raw_model:
        raise TrainError("Model not specified; set [job].model or [model].base in the config")

    try:
        base_model = ensure_allowed_model(
            raw_model,
            allow_finetuned_prefixes=False,
            allow_experimental=allow_experimental,
        )
    except UnsupportedModelError as exc:
        raise TrainError(str(exc)) from exc

    if base_model:
        try:
            ensure_model_supported_for_algorithm(base_model, spec)
        except AlgorithmValidationError as exc:
            raise TrainError(str(exc)) from exc

    try:
        payload = prepare_sft_job_payload(
            model=raw_model,
            training_file=None,
            hyperparameters=hp_block,
            metadata={"effective_config": effective},
            training_type="sft_offline",
            training_file_field="training_file_id",
            require_training_file=False,
            include_training_file_when_none=True,
            allow_finetuned_prefixes=False,
        )
    except UnsupportedModelError as exc:
        raise TrainError(str(exc)) from exc
    except ValueError as exc:
        raise TrainError(str(exc)) from exc

    return SFTBuildResult(payload=payload, train_file=dataset_path, validation_file=validation_file)


def build_prompt_learning_payload(
    *,
    config_path: Path,
    task_url: str | None,
    overrides: dict[str, Any],
    allow_experimental: bool | None = None,
) -> PromptLearningBuildResult:
    """Build payload for prompt learning job (MIPRO or GEPA)."""
    ctx: dict[str, Any] = {"config_path": str(config_path), "task_url": task_url}
    log_info("build_prompt_learning_payload invoked", ctx=ctx)
    from pydantic import ValidationError

    from .configs.prompt_learning import load_toml

    # SDK-SIDE VALIDATION: Catch errors BEFORE sending to backend
    from .validators import validate_prompt_learning_config

    raw_config = load_toml(config_path)
    validate_prompt_learning_config(raw_config, config_path)

    try:
        pl_cfg = PromptLearningConfig.from_path(config_path)
    except ValidationError as exc:
        raise click.ClickException(_format_validation_error(config_path, exc)) from exc

    # Early validation: Check required fields for GEPA
    if pl_cfg.algorithm == "gepa":
        if not pl_cfg.gepa:
            raise click.ClickException(
                "GEPA config missing: [prompt_learning.gepa] section is required"
            )
        if not pl_cfg.gepa.evaluation:
            raise click.ClickException(
                "GEPA config missing: [prompt_learning.gepa.evaluation] section is required"
            )
        train_seeds = getattr(pl_cfg.gepa.evaluation, "train_seeds", None) or getattr(
            pl_cfg.gepa.evaluation, "seeds", None
        )
        if not train_seeds:
            raise click.ClickException(
                "GEPA config missing train_seeds: [prompt_learning.gepa.evaluation] must have 'train_seeds' or 'seeds' field"
            )
        val_seeds = getattr(pl_cfg.gepa.evaluation, "val_seeds", None) or getattr(
            pl_cfg.gepa.evaluation, "validation_seeds", None
        )
        if not val_seeds:
            raise click.ClickException(
                "GEPA config missing val_seeds: [prompt_learning.gepa.evaluation] must have 'val_seeds' or 'validation_seeds' field"
            )

    cli_task_url = overrides.get("task_url") or task_url
    env_task_url = os.environ.get("TASK_APP_URL")
    config_task_url = (pl_cfg.task_app_url or "").strip() or None

    # For prompt learning, prefer config value over env if config is explicitly set
    # This allows TOML files to specify task_app_url without env var interference
    # But CLI override always wins
    if cli_task_url:
        # CLI override takes precedence
        final_task_url = ConfigResolver.resolve(
            "task_app_url",
            cli_value=cli_task_url,
            env_value=None,  # Don't check env when CLI is set
            config_value=config_task_url,
            required=True,
        )
    elif config_task_url:
        # Config explicitly set - use it (ignore env var to avoid conflicts)
        final_task_url = config_task_url
    else:
        # No config, fall back to env or error
        final_task_url = ConfigResolver.resolve(
            "task_app_url",
            cli_value=None,
            env_value=env_task_url,
            config_value=None,
            required=True,
        )
    assert final_task_url is not None  # required=True guarantees non-None

    # Get task_app_api_key from config or environment
    config_api_key = (pl_cfg.task_app_api_key or "").strip() or None
    cli_api_key = overrides.get("task_app_api_key")
    env_api_key = ensure_localapi_auth()
    task_app_api_key = ConfigResolver.resolve(
        "task_app_api_key",
        cli_value=cli_api_key,
        env_value=env_api_key,
        config_value=config_api_key,
        required=True,
    )

    # Build config dict for backend
    config_dict = pl_cfg.to_dict()

    # ASSERT: MIPRO fields exist in Pydantic model
    if pl_cfg.algorithm == "mipro":
        assert pl_cfg.mipro is not None, "pl_cfg.mipro is None for MIPRO algorithm"
        bootstrap_in_model = getattr(pl_cfg.mipro, "bootstrap_train_seeds", None)
        online_in_model = getattr(pl_cfg.mipro, "online_pool", None)
        assert bootstrap_in_model is not None, (
            f"pl_cfg.mipro.bootstrap_train_seeds is None! pl_cfg.mipro keys: {dir(pl_cfg.mipro) if pl_cfg.mipro else 'N/A'}"
        )
        assert online_in_model is not None, (
            f"pl_cfg.mipro.online_pool is None! pl_cfg.mipro keys: {dir(pl_cfg.mipro) if pl_cfg.mipro else 'N/A'}"
        )

    # Ensure task_app_url and task_app_api_key are set
    pl_section = config_dict.get("prompt_learning", {})
    if isinstance(pl_section, dict):
        pl_section["task_app_url"] = final_task_url
        pl_section["task_app_api_key"] = task_app_api_key

        # GEPA: Extract train_seeds from nested structure for backwards compatibility
        # Backend checks for train_seeds at top level before parsing nested structure
        if pl_cfg.algorithm == "gepa" and pl_cfg.gepa:
            # Try to get train_seeds directly from the gepa config object first
            train_seeds = None
            if pl_cfg.gepa.evaluation:
                train_seeds = getattr(pl_cfg.gepa.evaluation, "train_seeds", None) or getattr(
                    pl_cfg.gepa.evaluation, "seeds", None
                )

            # If not found, try from serialized dict
            if not train_seeds:
                gepa_section = pl_section.get("gepa", {})
                # Handle case where gepa_section might still be a Pydantic model
                if hasattr(gepa_section, "model_dump"):
                    gepa_section = gepa_section.model_dump(mode="python")
                elif hasattr(gepa_section, "dict"):
                    gepa_section = gepa_section.dict()

                if isinstance(gepa_section, dict):
                    eval_section = gepa_section.get("evaluation", {})
                    # Handle case where eval_section might still be a Pydantic model
                    if hasattr(eval_section, "model_dump"):
                        eval_section = eval_section.model_dump(mode="python")
                    elif hasattr(eval_section, "dict"):
                        eval_section = eval_section.dict()

                    if isinstance(eval_section, dict):
                        train_seeds = eval_section.get("train_seeds") or eval_section.get("seeds")

                    # Update gepa_section back to pl_section in case we converted it
                    pl_section["gepa"] = gepa_section

            # Add train_seeds to top level for backwards compatibility
            if train_seeds and not pl_section.get("train_seeds"):
                pl_section["train_seeds"] = train_seeds
            if train_seeds and not pl_section.get("evaluation_seeds"):
                pl_section["evaluation_seeds"] = train_seeds

        # MIPRO: CRITICAL - Ensure bootstrap_train_seeds and online_pool are ALWAYS in mipro section
        # Handle Pydantic model serialization - mipro might be a model object, not a dict
        if pl_cfg.algorithm == "mipro":
            mipro_section = pl_section.get("mipro", {})

            # ASSERT: Check what we got from to_dict()
            assert pl_cfg.mipro is not None, "pl_cfg.mipro is None"
            bootstrap_before_convert = getattr(pl_cfg.mipro, "bootstrap_train_seeds", None)
            online_before_convert = getattr(pl_cfg.mipro, "online_pool", None)
            assert bootstrap_before_convert is not None, (
                f"bootstrap_train_seeds is None in model! mipro_section type: {type(mipro_section)}, keys: {list(mipro_section.keys()) if isinstance(mipro_section, dict) else 'N/A'}"
            )
            assert online_before_convert is not None, (
                f"online_pool is None in model! mipro_section type: {type(mipro_section)}, keys: {list(mipro_section.keys()) if isinstance(mipro_section, dict) else 'N/A'}"
            )

            # Convert Pydantic model to dict if needed
            if hasattr(mipro_section, "model_dump"):
                mipro_section = mipro_section.model_dump(mode="python")
            elif hasattr(mipro_section, "dict"):
                mipro_section = mipro_section.dict()

            if not isinstance(mipro_section, dict):
                mipro_section = {}

            # ASSERT: After conversion, check if fields are present
            bootstrap_after_convert = mipro_section.get("bootstrap_train_seeds")
            online_after_convert = mipro_section.get("online_pool")
            if bootstrap_after_convert is None:
                raise AssertionError(
                    f"bootstrap_train_seeds missing after conversion! mipro_section keys: {list(mipro_section.keys())}, bootstrap_before_convert: {bootstrap_before_convert}"
                )
            if online_after_convert is None:
                raise AssertionError(
                    f"online_pool missing after conversion! mipro_section keys: {list(mipro_section.keys())}, online_before_convert: {online_before_convert}"
                )

            # CRITICAL: Get fields from Pydantic model FIRST (most reliable)
            # These fields MUST be present - get them from the source of truth
            if not pl_cfg.mipro:
                raise ValueError(
                    "MIPRO config missing: pl_cfg.mipro is None. "
                    "Ensure [prompt_learning.mipro] section exists in TOML."
                )

            bootstrap_from_model = getattr(pl_cfg.mipro, "bootstrap_train_seeds", None)
            online_from_model = getattr(pl_cfg.mipro, "online_pool", None)
            test_from_model = getattr(pl_cfg.mipro, "test_pool", None)
            reference_from_model = getattr(pl_cfg.mipro, "reference_pool", None)

            # FORCE these fields into mipro_section (model is source of truth)
            # Use model values if present, otherwise keep existing dict values
            assert bootstrap_from_model is not None, (
                f"bootstrap_from_model is None! pl_cfg.mipro: {pl_cfg.mipro}"
            )
            assert online_from_model is not None, (
                f"online_from_model is None! pl_cfg.mipro: {pl_cfg.mipro}"
            )

            mipro_section["bootstrap_train_seeds"] = bootstrap_from_model
            mipro_section["online_pool"] = online_from_model

            if test_from_model is not None:
                mipro_section["test_pool"] = test_from_model
            elif not mipro_section.get("test_pool") and pl_section.get("test_pool"):
                mipro_section["test_pool"] = pl_section["test_pool"]

            if reference_from_model is not None:
                mipro_section["reference_pool"] = reference_from_model
            elif not mipro_section.get("reference_pool") and pl_section.get("reference_pool"):
                mipro_section["reference_pool"] = pl_section["reference_pool"]

            # ASSERT: Fields are now in mipro_section
            assert mipro_section.get("bootstrap_train_seeds") is not None, (
                f"bootstrap_train_seeds STILL missing after forcing! mipro_section keys: {list(mipro_section.keys())}"
            )
            assert mipro_section.get("online_pool") is not None, (
                f"online_pool STILL missing after forcing! mipro_section keys: {list(mipro_section.keys())}"
            )

            # CRITICAL: Validate fields are present BEFORE override merge
            # If they're missing here, we'll check overrides after they're extracted
            # For now, just ensure they're in mipro_section from TOML
            if not mipro_section.get("bootstrap_train_seeds") and bootstrap_from_model is None:
                raise ValueError(
                    f"MIPRO config missing bootstrap_train_seeds in TOML. "
                    f"pl_cfg.mipro.bootstrap_train_seeds={bootstrap_from_model}, "
                    f"mipro_section keys={list(mipro_section.keys())}, "
                    f"pl_section keys={list(pl_section.keys())[:10]}. "
                    f"Ensure [prompt_learning.mipro] has bootstrap_train_seeds or provide override."
                )
            if not mipro_section.get("online_pool") and online_from_model is None:
                raise ValueError(
                    f"MIPRO config missing online_pool in TOML. "
                    f"pl_cfg.mipro.online_pool={online_from_model}, "
                    f"mipro_section keys={list(mipro_section.keys())}, "
                    f"pl_section keys={list(pl_section.keys())[:10]}. "
                    f"Ensure [prompt_learning.mipro] has online_pool or provide override."
                )

            # Extract env_name from mipro section to top-level (backend expects it there)
            mipro_env_name = mipro_section.get("env_name")
            if (
                mipro_env_name
                and not pl_section.get("env_name")
                and not pl_section.get("task_app_id")
            ):
                pl_section["env_name"] = mipro_env_name

            # CRITICAL: Update mipro section back to config_dict IMMEDIATELY
            # This ensures fields are present before override merge
            pl_section["mipro"] = mipro_section
            config_dict["prompt_learning"] = pl_section

            # ASSERT: Fields are in config_dict before override merge
            assert (
                config_dict.get("prompt_learning", {}).get("mipro", {}).get("bootstrap_train_seeds")
                is not None
            ), (
                f"bootstrap_train_seeds missing from config_dict before override merge! config_dict keys: {list(config_dict.keys())}"
            )
            assert (
                config_dict.get("prompt_learning", {}).get("mipro", {}).get("online_pool")
                is not None
            ), (
                f"online_pool missing from config_dict before override merge! config_dict keys: {list(config_dict.keys())}"
            )
    else:
        config_dict["prompt_learning"] = {
            "task_app_url": final_task_url,
            "task_app_api_key": task_app_api_key,
        }

    # Build payload matching backend API format
    # Extract nested overrides if present, otherwise use flat overrides directly
    # The experiment queue passes flat overrides like {"prompt_learning.policy.model": "..."}
    # But some SDK code passes nested like {"overrides": {"prompt_learning.policy.model": "..."}}
    config_overrides = overrides.get("overrides", {}) if "overrides" in overrides else overrides
    # Remove non-override keys (backend, task_url, metadata, auto_start)
    config_overrides = {
        k: v
        for k, v in config_overrides.items()
        if k not in ("backend", "task_url", "metadata", "auto_start")
    }

    # ASSERT: Check MIPRO fields BEFORE override merge
    if pl_cfg.algorithm == "mipro":
        pre_merge_mipro = config_dict.get("prompt_learning", {}).get("mipro", {})
        assert pre_merge_mipro.get("bootstrap_train_seeds") is not None, (
            f"bootstrap_train_seeds missing BEFORE override merge! pre_merge_mipro keys: {list(pre_merge_mipro.keys()) if isinstance(pre_merge_mipro, dict) else 'NOT DICT'}"
        )
        assert pre_merge_mipro.get("online_pool") is not None, (
            f"online_pool missing BEFORE override merge! pre_merge_mipro keys: {list(pre_merge_mipro.keys()) if isinstance(pre_merge_mipro, dict) else 'NOT DICT'}"
        )

    # CRITICAL: Merge overrides into config_dict BEFORE sending to backend
    # This ensures early validation in backend sees merged values
    # Use the same _deep_update logic from experiment_queue/config_utils.py
    if config_overrides:
        from synth_ai.core.experiment_queue.config_utils import _deep_update

        _deep_update(config_dict, config_overrides)

        # ASSERT: Check MIPRO fields AFTER override merge
        if pl_cfg.algorithm == "mipro":
            post_merge_mipro = config_dict.get("prompt_learning", {}).get("mipro", {})
            assert post_merge_mipro.get("bootstrap_train_seeds") is not None, (
                f"bootstrap_train_seeds missing AFTER override merge! post_merge_mipro keys: {list(post_merge_mipro.keys()) if isinstance(post_merge_mipro, dict) else 'NOT DICT'}, overrides: {list(config_overrides.keys())[:5]}"
            )
            assert post_merge_mipro.get("online_pool") is not None, (
                f"online_pool missing AFTER override merge! post_merge_mipro keys: {list(post_merge_mipro.keys()) if isinstance(post_merge_mipro, dict) else 'NOT DICT'}, overrides: {list(config_overrides.keys())[:5]}"
            )

    # After merging overrides (or if no overrides), re-run MIPRO reorganization to ensure
    # fields like bootstrap_train_seeds and online_pool are in the right place
    # This must run AFTER override merge so we catch fields from overrides
    pl_section_in_dict = config_dict.get("prompt_learning", {})
    if pl_cfg.algorithm == "mipro" and isinstance(pl_section_in_dict, dict):
        mipro_section = pl_section_in_dict.get("mipro", {})
        if not isinstance(mipro_section, dict):
            mipro_section = {}

        # CRITICAL: After override merge, ensure fields are in mipro section
        # Check both top-level and nested locations (overrides may have added them anywhere)
        if not mipro_section.get("bootstrap_train_seeds") and pl_section_in_dict.get(
            "bootstrap_train_seeds"
        ):
            mipro_section["bootstrap_train_seeds"] = pl_section_in_dict["bootstrap_train_seeds"]

        if not mipro_section.get("online_pool") and pl_section_in_dict.get("online_pool"):
            mipro_section["online_pool"] = pl_section_in_dict["online_pool"]

        if not mipro_section.get("test_pool") and pl_section_in_dict.get("test_pool"):
            mipro_section["test_pool"] = pl_section_in_dict["test_pool"]

        if not mipro_section.get("reference_pool") and pl_section_in_dict.get("reference_pool"):
            mipro_section["reference_pool"] = pl_section_in_dict["reference_pool"]

        # ASSERT: Fields should be in mipro_section after reorganization
        assert mipro_section.get("bootstrap_train_seeds") is not None, (
            f"bootstrap_train_seeds missing after reorganization! mipro_section keys: {list(mipro_section.keys())}, pl_section keys: {list(pl_section_in_dict.keys())[:10]}"
        )
        assert mipro_section.get("online_pool") is not None, (
            f"online_pool missing after reorganization! mipro_section keys: {list(mipro_section.keys())}, pl_section keys: {list(pl_section_in_dict.keys())[:10]}"
        )

        # Update mipro section back to config_dict
        pl_section_in_dict["mipro"] = mipro_section
        config_dict["prompt_learning"] = pl_section_in_dict

        # ASSERT: Fields are in config_dict after reorganization
        assert (
            config_dict.get("prompt_learning", {}).get("mipro", {}).get("bootstrap_train_seeds")
            is not None
        ), "bootstrap_train_seeds missing from config_dict after reorganization!"
        assert (
            config_dict.get("prompt_learning", {}).get("mipro", {}).get("online_pool") is not None
        ), "online_pool missing from config_dict after reorganization!"

        # CRITICAL: Verify MIPRO required fields are present after merge
        # This is the final check before sending to backend (runs whether or not overrides were applied)
        final_mipro_section = pl_section_in_dict.get("mipro", {})

        # DEBUG: Log what we found for troubleshooting
        bootstrap_seeds = final_mipro_section.get("bootstrap_train_seeds")
        online_pool_val = final_mipro_section.get("online_pool")

        if not bootstrap_seeds or not online_pool_val:
            # Log debug info before raising error
            import logging

            logger = logging.getLogger(__name__)
            logger.error(
                "MIPRO config validation failed:\n"
                f"  bootstrap_train_seeds in mipro section: {bootstrap_seeds is not None}\n"
                f"  online_pool in mipro section: {online_pool_val is not None}\n"
                f"  mipro_section keys: {list(final_mipro_section.keys()) if isinstance(final_mipro_section, dict) else 'NOT A DICT'}\n"
                f"  pl_section keys: {list(pl_section_in_dict.keys()) if isinstance(pl_section_in_dict, dict) else 'NOT A DICT'}\n"
                f"  config_overrides keys: {list(config_overrides.keys()) if config_overrides else 'NO OVERRIDES'}"
            )

        if not bootstrap_seeds:
            raise ValueError(
                "MIPRO config missing bootstrap_train_seeds after applying overrides. "
                "Ensure it's set in TOML at [prompt_learning.mipro] or [prompt_learning] level, "
                "or provided via override 'prompt_learning.mipro.bootstrap_train_seeds'. "
                f"Current mipro_section keys: {list(final_mipro_section.keys()) if isinstance(final_mipro_section, dict) else 'NOT A DICT'}"
            )
        if not online_pool_val:
            raise ValueError(
                "MIPRO config missing online_pool after applying overrides. "
                "Ensure it's set in TOML at [prompt_learning.mipro] or [prompt_learning] level, "
                "or provided via override 'prompt_learning.mipro.online_pool'. "
                f"Current mipro_section keys: {list(final_mipro_section.keys()) if isinstance(final_mipro_section, dict) else 'NOT A DICT'}"
            )

    # CRITICAL: Final validation - ensure MIPRO fields are in config_body before sending
    pl_section_final = config_dict.get("prompt_learning", {})
    if pl_cfg.algorithm == "mipro":
        mipro_final = pl_section_final.get("mipro", {})
        if not isinstance(mipro_final, dict):
            mipro_final = {}

        # If fields are still missing, FORCE them from the Pydantic model one last time
        if pl_cfg.mipro:
            if not mipro_final.get("bootstrap_train_seeds"):
                bootstrap_val = getattr(pl_cfg.mipro, "bootstrap_train_seeds", None)
                if bootstrap_val is not None:
                    mipro_final["bootstrap_train_seeds"] = bootstrap_val

            if not mipro_final.get("online_pool"):
                online_val = getattr(pl_cfg.mipro, "online_pool", None)
                if online_val is not None:
                    mipro_final["online_pool"] = online_val

            if not mipro_final.get("test_pool"):
                test_val = getattr(pl_cfg.mipro, "test_pool", None)
                if test_val is not None:
                    mipro_final["test_pool"] = test_val

            if not mipro_final.get("reference_pool"):
                ref_val = getattr(pl_cfg.mipro, "reference_pool", None)
                if ref_val is not None:
                    mipro_final["reference_pool"] = ref_val

        # Update back to config_dict
        pl_section_final["mipro"] = mipro_final
        config_dict["prompt_learning"] = pl_section_final

        # ASSERT: Fields are in config_dict after final check
        assert (
            config_dict.get("prompt_learning", {}).get("mipro", {}).get("bootstrap_train_seeds")
            is not None
        ), (
            f"bootstrap_train_seeds missing from config_dict after final check! mipro_final keys: {list(mipro_final.keys())}"
        )
        assert (
            config_dict.get("prompt_learning", {}).get("mipro", {}).get("online_pool") is not None
        ), (
            f"online_pool missing from config_dict after final check! mipro_final keys: {list(mipro_final.keys())}"
        )

        # FINAL ASSERTION: These fields MUST be present
        assert mipro_final.get("bootstrap_train_seeds") is not None, (
            f"CRITICAL: bootstrap_train_seeds missing from config_body! mipro_final keys: {list(mipro_final.keys())}, pl_cfg.mipro.bootstrap_train_seeds: {getattr(pl_cfg.mipro, 'bootstrap_train_seeds', None) if pl_cfg.mipro else 'N/A'}"
        )
        assert mipro_final.get("online_pool") is not None, (
            f"CRITICAL: online_pool missing from config_body! mipro_final keys: {list(mipro_final.keys())}, pl_cfg.mipro.online_pool: {getattr(pl_cfg.mipro, 'online_pool', None) if pl_cfg.mipro else 'N/A'}"
        )

    # ASSERT: Verify critical overrides are reflected in config_body
    pl_section_in_dict = config_dict.get("prompt_learning", {})
    if config_overrides:
        # Check rollout budget override
        rollout_budget_key = "prompt_learning.gepa.rollout.budget"
        if rollout_budget_key in config_overrides:
            expected_budget = config_overrides[rollout_budget_key]
            gepa_section = pl_section_in_dict.get("gepa", {})
            actual_budget = (
                gepa_section.get("rollout", {}).get("budget")
                if isinstance(gepa_section, dict)
                else None
            )
            if actual_budget is not None:
                assert actual_budget == expected_budget, (
                    f"Rollout budget mismatch: config_body has {actual_budget} but override specifies {expected_budget}. "
                    f"This indicates the override wasn't applied correctly."
                )

        # Check model override
        model_key = "prompt_learning.policy.model"
        if model_key in config_overrides:
            expected_model = config_overrides[model_key]
            policy_section = pl_section_in_dict.get("policy", {})
            actual_model = policy_section.get("model") if isinstance(policy_section, dict) else None
            if actual_model is not None:
                assert actual_model == expected_model, (
                    f"Model mismatch: config_body has {actual_model} but override specifies {expected_model}. "
                    f"This indicates the override wasn't applied correctly."
                )

        # Check provider override
        provider_key = "prompt_learning.policy.provider"
        if provider_key in config_overrides:
            expected_provider = config_overrides[provider_key]
            policy_section = pl_section_in_dict.get("policy", {})
            actual_provider = (
                policy_section.get("provider") if isinstance(policy_section, dict) else None
            )
            if actual_provider is not None:
                assert actual_provider == expected_provider, (
                    f"Provider mismatch: config_body has {actual_provider} but override specifies {expected_provider}. "
                    f"This indicates the override wasn't applied correctly."
                )

    # FINAL CHECK: Ensure config_body has correct structure for backend
    # Backend expects: {"prompt_learning": {...}} (full TOML structure)
    if "prompt_learning" not in config_dict:
        raise ValueError(
            f"config_dict must have 'prompt_learning' key. Found keys: {list(config_dict.keys())}"
        )

    # CRITICAL: Final validation - check MIPRO fields are in config_body before sending
    if pl_cfg.algorithm == "mipro":
        pl_section_final = config_dict.get("prompt_learning", {})
        mipro_section_final = (
            pl_section_final.get("mipro", {}) if isinstance(pl_section_final, dict) else {}
        )

        # ASSERT: Fields MUST be present before building payload
        assert mipro_section_final.get("bootstrap_train_seeds") is not None, (
            f"CRITICAL ASSERTION FAILED: bootstrap_train_seeds missing from config_body before sending to backend! "
            f"mipro_section keys: {list(mipro_section_final.keys()) if isinstance(mipro_section_final, dict) else 'NOT A DICT'}, "
            f"pl_section keys: {list(pl_section_final.keys()) if isinstance(pl_section_final, dict) else 'NOT A DICT'}, "
            f"pl_cfg.mipro.bootstrap_train_seeds: {getattr(pl_cfg.mipro, 'bootstrap_train_seeds', None) if pl_cfg.mipro else 'N/A'}"
        )
        assert mipro_section_final.get("online_pool") is not None, (
            f"CRITICAL ASSERTION FAILED: online_pool missing from config_body before sending to backend! "
            f"mipro_section keys: {list(mipro_section_final.keys()) if isinstance(mipro_section_final, dict) else 'NOT A DICT'}, "
            f"pl_section keys: {list(pl_section_final.keys()) if isinstance(pl_section_final, dict) else 'NOT A DICT'}, "
            f"pl_cfg.mipro.online_pool: {getattr(pl_cfg.mipro, 'online_pool', None) if pl_cfg.mipro else 'N/A'}"
        )

        if not mipro_section_final.get("bootstrap_train_seeds"):
            import json

            raise ValueError(
                "CRITICAL: bootstrap_train_seeds missing from config_body before sending to backend. "
                f"mipro_section keys: {list(mipro_section_final.keys()) if isinstance(mipro_section_final, dict) else 'NOT A DICT'}, "
                f"pl_section keys: {list(pl_section_final.keys()) if isinstance(pl_section_final, dict) else 'NOT A DICT'}, "
                f"config_overrides: {json.dumps(list(config_overrides.keys()) if config_overrides else [])}"
            )
        if not mipro_section_final.get("online_pool"):
            import json

            raise ValueError(
                "CRITICAL: online_pool missing from config_body before sending to backend. "
                f"mipro_section keys: {list(mipro_section_final.keys()) if isinstance(mipro_section_final, dict) else 'NOT A DICT'}, "
                f"pl_section keys: {list(pl_section_final.keys()) if isinstance(pl_section_final, dict) else 'NOT A DICT'}, "
                f"config_overrides: {json.dumps(list(config_overrides.keys()) if config_overrides else [])}"
            )

    payload: dict[str, Any] = {
        "algorithm": pl_cfg.algorithm,
        "config_body": config_dict,
        "overrides": config_overrides,
        "metadata": overrides.get("metadata", {}),
        "auto_start": overrides.get("auto_start", True),
    }

    # CRITICAL DEBUG: Print MIPRO section structure before sending
    if pl_cfg.algorithm == "mipro":
        import json

        mipro_debug = config_dict.get("prompt_learning", {}).get("mipro", {})
        print("\n🔍 DEBUG: MIPRO section in config_body before sending:")
        print(f"  Type: {type(mipro_debug)}")
        print(
            f"  Keys: {list(mipro_debug.keys()) if isinstance(mipro_debug, dict) else 'NOT A DICT'}"
        )
        print(
            f"  bootstrap_train_seeds present: {mipro_debug.get('bootstrap_train_seeds') is not None}"
        )
        print(f"  online_pool present: {mipro_debug.get('online_pool') is not None}")
        if isinstance(mipro_debug, dict):
            print(f"  bootstrap_train_seeds value: {mipro_debug.get('bootstrap_train_seeds')}")
            print(f"  online_pool value: {mipro_debug.get('online_pool')}")
        # Print full mipro section (truncated)
        mipro_json = json.dumps(mipro_debug, indent=2, default=str)
        print(f"  Full mipro section (first 500 chars):\n{mipro_json[:500]}")
        print("🔍 END DEBUG\n")

    backend = overrides.get("backend")
    if backend:
        metadata_default: dict[str, Any] = {}
        metadata = cast(dict[str, Any], payload.setdefault("metadata", metadata_default))
        metadata["backend_base_url"] = ensure_api_base(str(backend))

    return PromptLearningBuildResult(payload=payload, task_url=final_task_url)


def build_prompt_learning_payload_from_mapping(
    *,
    raw_config: dict[str, Any],
    task_url: str | None,
    overrides: dict[str, Any],
    allow_experimental: bool | None = None,
    source_label: str = "programmatic",
) -> PromptLearningBuildResult:
    """Build payload for prompt learning job from a dictionary (programmatic use).

    This is the same as build_prompt_learning_payload but accepts a dict instead of a file path.
    Both functions route through the same PromptLearningConfig Pydantic validation.

    Args:
        raw_config: Configuration dictionary with the same structure as the TOML file.
                   Should have a 'prompt_learning' section.
        task_url: Override for task_app_url
        overrides: Config overrides (merged into config)
        allow_experimental: Allow experimental models
        source_label: Label for logging/error messages (default: "programmatic")

    Returns:
        PromptLearningBuildResult with payload and task_url

    Example:
        >>> result = build_prompt_learning_payload_from_mapping(
        ...     raw_config={
        ...         "prompt_learning": {
        ...             "algorithm": "gepa",
        ...             "task_app_url": "https://tunnel.example.com",
        ...             "policy": {"model": "gpt-4o-mini", "provider": "openai"},
        ...             "gepa": {...},
        ...         }
        ...     },
        ...     task_url=None,
        ...     overrides={},
        ... )
    """
    ctx: dict[str, Any] = {"source": source_label}
    log_info("build_prompt_learning_payload_from_mapping invoked", ctx=ctx)
    from pydantic import ValidationError

    # SDK-SIDE VALIDATION: Catch errors BEFORE sending to backend
    from .validators import validate_prompt_learning_config

    # Use a pseudo-path for error messages (validator expects Path object)
    pseudo_path = Path(f"<{source_label}>")
    validate_prompt_learning_config(raw_config, pseudo_path)

    try:
        pl_cfg = PromptLearningConfig.from_mapping(raw_config)
    except ValidationError as exc:
        # Format validation errors for dict-based config
        lines: list[str] = []
        for error in exc.errors():
            loc = ".".join(str(part) for part in error.get("loc", ()))
            msg = error.get("msg", "invalid value")
            lines.append(f"{loc or '<root>'}: {msg}")
        details = "\n".join(f"  - {line}" for line in lines) or "  - Invalid configuration"
        raise click.ClickException(
            f"Config validation failed ({source_label}):\n{details}"
        ) from exc

    # Early validation: Check required fields for GEPA
    if pl_cfg.algorithm == "gepa":
        if not pl_cfg.gepa:
            raise click.ClickException(
                "GEPA config missing: [prompt_learning.gepa] section is required"
            )
        if not pl_cfg.gepa.evaluation:
            raise click.ClickException(
                "GEPA config missing: [prompt_learning.gepa.evaluation] section is required"
            )
        train_seeds = getattr(pl_cfg.gepa.evaluation, "train_seeds", None) or getattr(
            pl_cfg.gepa.evaluation, "seeds", None
        )
        if not train_seeds:
            raise click.ClickException(
                "GEPA config missing train_seeds: [prompt_learning.gepa.evaluation] must have 'train_seeds' or 'seeds' field"
            )
        val_seeds = getattr(pl_cfg.gepa.evaluation, "val_seeds", None) or getattr(
            pl_cfg.gepa.evaluation, "validation_seeds", None
        )
        if not val_seeds:
            raise click.ClickException(
                "GEPA config missing val_seeds: [prompt_learning.gepa.evaluation] must have 'val_seeds' or 'validation_seeds' field"
            )

    cli_task_url = overrides.get("task_url") or task_url
    env_task_url = os.environ.get("TASK_APP_URL")
    config_task_url = (pl_cfg.task_app_url or "").strip() or None

    # Resolve task_app_url with same precedence as file-based builder
    if cli_task_url:
        final_task_url = ConfigResolver.resolve(
            "task_app_url",
            cli_value=cli_task_url,
            env_value=None,
            config_value=config_task_url,
            required=True,
        )
    elif config_task_url:
        final_task_url = config_task_url
    else:
        final_task_url = ConfigResolver.resolve(
            "task_app_url",
            cli_value=None,
            env_value=env_task_url,
            config_value=None,
            required=True,
        )
    assert final_task_url is not None

    # Get task_app_api_key from config or environment
    config_api_key = (pl_cfg.task_app_api_key or "").strip() or None
    cli_api_key = overrides.get("task_app_api_key")
    env_api_key = ensure_localapi_auth()
    task_app_api_key = ConfigResolver.resolve(
        "task_app_api_key",
        cli_value=cli_api_key,
        env_value=env_api_key,
        config_value=config_api_key,
        required=True,
    )

    # Build config dict for backend
    config_dict = pl_cfg.to_dict()

    # Ensure task_app_url and task_app_api_key are set
    pl_section = config_dict.get("prompt_learning", {})
    if isinstance(pl_section, dict):
        pl_section["task_app_url"] = final_task_url
        pl_section["task_app_api_key"] = task_app_api_key

        # GEPA: Extract train_seeds from nested structure
        if pl_cfg.algorithm == "gepa" and pl_cfg.gepa:
            train_seeds = None
            if pl_cfg.gepa.evaluation:
                train_seeds = getattr(pl_cfg.gepa.evaluation, "train_seeds", None) or getattr(
                    pl_cfg.gepa.evaluation, "seeds", None
                )

            if train_seeds and not pl_section.get("train_seeds"):
                pl_section["train_seeds"] = train_seeds
            if train_seeds and not pl_section.get("evaluation_seeds"):
                pl_section["evaluation_seeds"] = train_seeds
    else:
        config_dict["prompt_learning"] = {
            "task_app_url": final_task_url,
            "task_app_api_key": task_app_api_key,
        }

    # Build payload matching backend API format
    config_overrides = overrides.get("overrides", {}) if "overrides" in overrides else overrides
    config_overrides = {
        k: v
        for k, v in config_overrides.items()
        if k not in ("backend", "task_url", "metadata", "auto_start")
    }

    # Merge overrides into config_dict
    if config_overrides:
        from synth_ai.core.experiment_queue.config_utils import _deep_update

        _deep_update(config_dict, config_overrides)

    # Final validation
    if "prompt_learning" not in config_dict:
        raise ValueError(
            f"config_dict must have 'prompt_learning' key. Found keys: {list(config_dict.keys())}"
        )

    payload: dict[str, Any] = {
        "algorithm": pl_cfg.algorithm,
        "config_body": config_dict,
        "overrides": config_overrides,
        "metadata": overrides.get("metadata", {}),
        "auto_start": overrides.get("auto_start", True),
    }

    backend = overrides.get("backend")
    if backend:
        metadata_default: dict[str, Any] = {}
        metadata = cast(dict[str, Any], payload.setdefault("metadata", metadata_default))
        metadata["backend_base_url"] = ensure_api_base(str(backend))

    return PromptLearningBuildResult(payload=payload, task_url=final_task_url)


__all__ = [
    "PromptLearningBuildResult",
    "RLBuildResult",
    "SFTBuildResult",
    "build_prompt_learning_payload",
    "build_prompt_learning_payload_from_mapping",
    "build_rl_payload",
    "build_sft_payload",
]
