"""Main optimizer class for Pydantic models using DSPy."""

import inspect
import time
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any

import dspy
from dspy.teleprompt import MIPROv2, Teleprompter  # noqa: E402
from pydantic import BaseModel
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich import box

from dspydantic.evaluators.functions import default_evaluate_fn
from dspydantic.extractor import extract_field_descriptions, extract_field_types
from dspydantic.module import PydanticOptimizerModule
from dspydantic.types import Example, FieldOptimizationProgress, OptimizationResult, create_output_model
from dspydantic.utils import convert_images_to_dspy_images, format_instruction_prompt_template

# Fast mode optimization kwargs: reduce demo count and optimizer complexity
_FAST_MODE_KWARGS: dict[str, dict] = {
    "bootstrapfewshot": {"max_bootstrapped_demos": 1},
    "bootstrapfewshotwithrandomsearch": {"max_bootstrapped_demos": 1, "num_candidates": 4},
    "miprov2": {"auto": "light", "max_bootstrapped_demos": 1},
    "miprov2zeroshot": {"auto": "light"},
}


class PydanticOptimizer:
    """Optimizer that uses DSPy to optimize Pydantic model field descriptions.

    This class optimizes field descriptions in Pydantic models by using DSPy
    to iteratively improve descriptions based on example data and a custom
    evaluation function.

    Examples:
        Basic usage without evaluation function (uses default with "exact" metric)::

            from pydantic import BaseModel, Field
            from dspydantic import PydanticOptimizer

            class User(BaseModel):
                name: str = Field(description="User name")
                age: int = Field(description="User age")

            examples = [
                Example(
                    input_data={"text": "John Doe, 30 years old"},
                    expected_output={"name": "John Doe", "age": 30}
                )
            ]

            # Configure DSPy first
            import dspy
            lm = dspy.LM("openai/gpt-4o", api_key="your-key")
            dspy.configure(lm=lm)

            optimizer = PydanticOptimizer(
                model=User,
                examples=examples
            )

        Using "exact" metric for exact string matching::

            optimizer = PydanticOptimizer(
                model=User,
                examples=examples,
                evaluate_fn="exact",
                model_id="gpt-4o",
                api_key="your-key"
            )

        Using "levenshtein" metric for fuzzy matching::

            optimizer = PydanticOptimizer(
                model=User,
                examples=examples,
                evaluate_fn="levenshtein",
                model_id="gpt-4o",
                api_key="your-key"
            )

        Using a custom evaluation function::

            def evaluate(
                example,
                optimized_descriptions,
                optimized_system_prompt,
                optimized_instruction_prompt,
            ):
                # Your custom evaluation logic here
                # Return a score between 0.0 and 1.0
                return 0.85

            optimizer = PydanticOptimizer(
                model=User,
                examples=examples,
                evaluate_fn=evaluate,
                model_id="gpt-4o",
                api_key="your-key"
            )

        Configure DSPy first::

            import dspy
            lm = dspy.LM("openai/gpt-4o", api_key="your-key")
            dspy.configure(lm=lm)

            optimizer = PydanticOptimizer(
                model=User,
                examples=examples
            )

        Passing optimizer as a string::

            optimizer = PydanticOptimizer(
                model=User,
                examples=examples,
                optimizer="miprov2"
            )

        Passing a custom optimizer instance::

            from dspy.teleprompt import MIPROv2
            custom_optimizer = MIPROv2(
                metric=lambda x, y, trace=None: 0.9,
                num_threads=8,
                auto="full"
            )
            optimizer = PydanticOptimizer(
                model=User,
                examples=examples,
                optimizer=custom_optimizer
            )

        Using None expected_output with LLM judge::

            examples_without_expected = [
                Example(
                    text="John Doe, 30 years old",
                    expected_output=None
                )
            ]
            optimizer = PydanticOptimizer(
                model=User,
                examples=examples_without_expected,
                model_id="gpt-4o",
                api_key="your-key"
            )

        Using None expected_output with custom judge LM::

            import dspy
            lm = dspy.LM("openai/gpt-4o", api_key="your-key")
            dspy.configure(lm=lm)

            judge_lm = dspy.LM("openai/gpt-4", api_key="your-key")
            optimizer = PydanticOptimizer(
                model=User,
                examples=examples_without_expected,
                evaluate_fn=judge_lm
            )

        Using None expected_output with custom judge function::

            def custom_judge(example, extracted_data, optimized_descriptions,
                            optimized_system_prompt, optimized_instruction_prompt):
                # Your custom evaluation logic here
                # Return a score between 0.0 and 1.0
                return 0.85

            optimizer = PydanticOptimizer(
                model=User,
                examples=examples_without_expected,
                evaluate_fn=custom_judge,
                model_id="gpt-4o",
                api_key="your-key"
            )

        Running optimization::

            result = optimizer.optimize()
            print(result.optimized_descriptions)
    """

    def __init__(
        self,
        model: type[BaseModel] | None,
        examples: list[Example],
        evaluate_fn: Callable[[Example, dict[str, str], str | None, str | None], float]
        | Callable[[Example, dict[str, Any], dict[str, str], str | None, str | None], float]
        | dspy.LM
        | str
        | None = None,
        system_prompt: str | None = None,
        instruction_prompt: str | None = None,
        num_threads: int = 4,
        init_temperature: float = 1.0,
        verbose: bool = False,
        optimizer: str | Teleprompter | None = None,
        train_split: float = 0.8,
        optimizer_kwargs: dict[str, Any] | None = None,
        exclude_fields: list[str] | None = None,
        include_fields: list[str] | None = None,
        evaluator_config: dict[str, Any] | None = None,
        sequential: bool = False,
        parallel_fields: bool = True,
        max_val_examples: int | None = None,
        skip_score_threshold: float | None = None,
        on_progress: Callable[[FieldOptimizationProgress], None] | None = None,
    ) -> None:
        """Initialize the Pydantic optimizer.

        Args:
            model: The Pydantic model class to optimize. If None, will auto-create
                a single-field model with field "output" when examples have string outputs.
            examples: List of examples to use for optimization.
            evaluate_fn: Optional function that evaluates the quality of optimized prompts.

                When expected_output is provided:
                    - Takes (Example, optimized_descriptions dict, optimized_system_prompt,
                      optimized_instruction_prompt), returns a float score (0.0-1.0).
                    - Can also be a string: "exact" for exact matching, "levenshtein" for
                      Levenshtein distance-based matching, or None for default evaluation
                      that performs structured extraction with the same LLM used for optimization.

                When expected_output is None:
                    - Can be a dspy.LM instance to use as a judge.
                    - Can be a callable that takes (Example, extracted_data dict,
                      optimized_descriptions dict, optimized_system_prompt,
                      optimized_instruction_prompt) and returns a float score (0.0-1.0).
                    - If None, uses the default LLM judge (same LM as optimization).
            system_prompt: Optional initial system prompt to optimize.
            instruction_prompt: Optional initial instruction prompt to optimize.
            num_threads: Number of threads for optimization.
            init_temperature: Initial temperature for optimization.
            verbose: If True, print detailed progress information.
            optimizer: Optimizer specification. Can be:

                - A string (optimizer type name): e.g., "miprov2", "gepa",
                  "bootstrapfewshot", etc. If None, optimizer will be auto-selected
                  based on dataset size.
                - A Teleprompter instance: Custom optimizer instance to use directly.

                  Available optimizer types include: "miprov2", "miprov2zeroshot", "gepa",
                  "bootstrapfewshot", "bootstrapfewshotwithrandomsearch", "knnfewshot",
                  "labeledfewshot", "copro", "simba", and all other Teleprompter subclasses.
            train_split: Fraction of examples to use for training (rest for validation).
            optimizer_kwargs: Optional dictionary of additional keyword arguments
                to pass to the optimizer constructor. These will override default
                parameters. For example: {"max_bootstrapped_demos": 8, "auto": "full"}.
                Only used if `optimizer` is a string or None.
            exclude_fields: Optional list of field paths to exclude from evaluation.
                Field paths use dot notation for nested fields
                (e.g., ["address.street", "metadata"]).
                Fields matching these paths (or starting with them) will be excluded
                from scoring. Only applies when using default evaluation functions
                (not custom evaluate_fn).
            include_fields: Optional list of field paths to include in optimization
                and evaluation. When set, only these fields (and nested fields
                under them) are optimized and scored. Mutually filters with
                exclude_fields when both are set.
            evaluator_config: Optional evaluator configuration dict with "default" and
                "field_overrides" keys. If provided, uses configured evaluators instead
                of evaluate_fn/metric. Supports string names, config dicts, and custom classes.
            sequential: If False (default), use single-pass optimization (one DSPy compile
                for all fields) with reduced demo budgets for speed. If True, optimize each
                field description independently (deepest-nested first) for maximum quality.
            parallel_fields: If True (default), parallelize field optimization when using
                sequential mode. Each field runs in a thread simultaneously. Has no effect
                when sequential=False.
            max_val_examples: Optional cap on validation set size. When set, uses only the
                first N validation examples per field optimization, reducing scoring LLM calls.
            skip_score_threshold: Optional threshold (0.0-1.0). When set in sequential mode,
                skips optimizing fields that already score above this baseline threshold.
            on_progress: Optional callback invoked after each field/phase optimization.
                Receives a FieldOptimizationProgress object with field_path, score updates, etc.

        Raises:
            ValueError: If at least one example is not provided, or if optimizer string
                is not a valid Teleprompter subclass name.
            TypeError: If optimizer is not a string, Teleprompter instance, or None.
        """
        if not examples:
            raise ValueError("At least one example must be provided")

        # Detect if examples have string outputs
        has_string_outputs = any(
            isinstance(ex.expected_output, str) for ex in examples if ex.expected_output is not None
        )

        # Auto-create OutputModel if model is None and we have string outputs
        if model is None:
            if has_string_outputs:
                model = create_output_model()
            else:
                raise ValueError(
                    "model cannot be None unless examples have string expected_output values"
                )
        elif has_string_outputs:
            # If model is provided but examples have strings, we'll convert strings to dicts
            # with {"output": <string>} format during evaluation
            pass

        self.model = model
        self.examples = examples
        self.evaluate_fn = evaluate_fn
        self.exclude_fields = exclude_fields
        self.include_fields = include_fields
        self.evaluator_config = evaluator_config
        self.system_prompt = system_prompt
        self.instruction_prompt = instruction_prompt
        self.num_threads = num_threads
        self.init_temperature = init_temperature
        self.verbose = verbose
        self.train_split = train_split
        self.optimizer_kwargs = optimizer_kwargs or {}
        self.sequential = sequential
        self.parallel_fields = parallel_fields
        self.max_val_examples = max_val_examples
        self.skip_score_threshold = skip_score_threshold
        self.on_progress = on_progress

        # Add default progress callback if verbose and no callback provided
        if verbose and on_progress is None:
            _console = Console()
            def _default_progress(p: FieldOptimizationProgress):
                if p.phase == "fields":
                    status = "[green]✓[/]" if p.improved else "[yellow]–[/]"
                    _console.print(f"  [bold]{p.field_path}[/] {p.score_before:.0%} → {p.score_after:.0%} {status}")
                    if p.optimized_value:
                        _console.print(f"    [dim]→ {p.optimized_value!r}[/]")
                elif p.phase == "skipped":
                    _console.print(f"  [dim]{p.field_path}: skipped ({p.score_before:.0%} ≥ threshold)[/]")
                elif p.phase in ("system_prompt", "instruction_prompt"):
                    status = "[green]✓[/]" if p.improved else "[yellow]–[/]"
                    _console.print(f"  [bold]{p.phase}[/] {p.score_before:.0%} → {p.score_after:.0%} {status}")
                    if p.optimized_value:
                        _console.print(f"    [dim]→ {p.optimized_value!r:.100}[/]")
            self.on_progress = _default_progress

        # Handle optimizer parameter (can be string or Teleprompter instance)
        if optimizer is None:
            # Auto-select optimizer based on dataset size
            self.optimizer_type = self._auto_select_optimizer()
            self.custom_optimizer = None
        elif isinstance(optimizer, str):
            # String provided - validate and store as type
            self.optimizer_type = optimizer.lower()
            # Validate optimizer type by checking if it's a Teleprompter subclass
            teleprompter_classes = self._get_teleprompter_subclasses()
            if self.optimizer_type not in teleprompter_classes:
                valid_optimizers = sorted(teleprompter_classes.keys())
                raise ValueError(
                    f"optimizer '{optimizer}' is not a valid Teleprompter subclass. "
                    f"Valid options: {valid_optimizers}"
                )
            self.custom_optimizer = None
        elif isinstance(optimizer, Teleprompter):
            # Teleprompter instance provided
            self.custom_optimizer = optimizer
            self.optimizer_type = "custom"
        else:
            raise TypeError(
                f"optimizer must be a string, Teleprompter instance, or None, "
                f"got {type(optimizer).__name__}"
            )

        # Apply fast mode optimizer kwargs by default (when not sequential)
        # Sequential mode can still use them but doesn't require them
        if not self.sequential and self.optimizer_type != "custom":
            fast_defaults = _FAST_MODE_KWARGS.get(self.optimizer_type, {})
            # User-supplied optimizer_kwargs override fast defaults
            self.optimizer_kwargs = {**fast_defaults, **self.optimizer_kwargs}

        # Extract field descriptions from Pydantic model
        # Field descriptions are automatically set from field names if not provided
        self.field_descriptions = extract_field_descriptions(self.model)

        # Extract field types from Pydantic model
        self.field_types = extract_field_types(self.model)

        # Check that we have something to optimize
        has_field_descriptions = bool(self.field_descriptions)
        has_system_prompt = self.system_prompt is not None
        has_instruction_prompt = self.instruction_prompt is not None

        if not (has_field_descriptions or has_system_prompt or has_instruction_prompt):
            raise ValueError(
                "At least one of the following must be provided: "
                "model fields (field descriptions are automatically set from field names "
                "if not provided), system_prompt, or instruction_prompt"
            )

    @staticmethod
    def _get_teleprompter_subclasses() -> dict[str, type[Teleprompter]]:
        """Get all subclasses of Teleprompter and create a mapping by lowercase name.

        Returns:
            Dictionary mapping lowercase class names to Teleprompter subclasses.
        """

        # Get all subclasses recursively
        def get_all_subclasses(cls: type) -> set[type]:
            """Recursively get all subclasses of a class."""
            subclasses = set()
            for subclass in cls.__subclasses__():
                subclasses.add(subclass)
                subclasses.update(get_all_subclasses(subclass))
            return subclasses

        subclasses = get_all_subclasses(Teleprompter)
        # Create mapping: lowercase class name -> class
        mapping: dict[str, type[Teleprompter]] = {}
        for subclass in subclasses:
            # Skip abstract classes or classes that shouldn't be used directly
            if subclass.__name__ == "Teleprompter":
                continue
            # Map lowercase name to class
            mapping[subclass.__name__.lower()] = subclass

        # Add special case for miprov2zeroshot (which is MIPROv2 with zero-shot settings)
        if "miprov2" in mapping:
            mapping["miprov2zeroshot"] = mapping["miprov2"]

        return mapping

    def _get_effective_fields(self) -> set[str]:
        """Return field paths to optimize after applying include/exclude filters.

        Returns:
            Set of field paths. If include_fields is set, only those (and nested)
            are included. exclude_fields removes from the result.
        """
        all_fields = set(self.field_descriptions.keys())
        if self.include_fields is not None:
            included = set()
            for path in self.include_fields:
                if path in all_fields:
                    included.add(path)
                for f in all_fields:
                    if f.startswith(f"{path}."):
                        included.add(f)
            all_fields = included
        if self.exclude_fields is not None:
            excluded = set()
            for path in self.exclude_fields:
                if path in all_fields:
                    excluded.add(path)
                for f in all_fields:
                    if f.startswith(f"{path}."):
                        excluded.add(f)
            all_fields -= excluded
        return all_fields

    def _sort_fields_by_depth(self) -> list[str]:
        """Return field paths sorted deepest-first (most dots first).

        Returns:
            List of field paths ordered by nesting depth, deepest first.
        """
        effective = self._get_effective_fields()
        return sorted(
            (p for p in self.field_descriptions.keys() if p in effective),
            key=lambda p: p.count("."),
            reverse=True,
        )

    def _auto_select_optimizer(self) -> str:
        """Auto-select the best optimizer based on the number of examples.

        Selection logic:
        - Very small datasets (1-2 examples): Use MIPROv2ZeroShot (avoids BootstrapFewShot bug)
        - Small datasets (3-19 examples): Use BootstrapFewShot
        - Larger datasets (>= 20 examples): Use BootstrapFewShotWithRandomSearch

        Returns:
            String name of the recommended optimizer type.
        """
        num_examples = len(self.examples)

        if num_examples <= 2:
            # Very small dataset - use MIPROv2ZeroShot to avoid BootstrapFewShot bug
            return "miprov2zeroshot"
        elif num_examples < 20:
            # Small dataset - use BootstrapFewShot
            return "bootstrapfewshot"
        else:
            # Larger dataset - use BootstrapFewShotWithRandomSearch
            return "bootstrapfewshotwithrandomsearch"

    def _default_evaluate_fn(
        self,
        lm: dspy.LM,
        metric: str = "exact",
        judge_lm: dspy.LM | None = None,
        evaluator_config: dict[str, Any] | None = None,
    ) -> Callable[[Example, dict[str, str], str | None, str | None], float]:
        """Create a default evaluation function that uses the LLM for structured extraction.

        Args:
            lm: The DSPy language model to use for extraction.
            metric: Comparison metric to use. Options:
                - "exact": Exact string matching (default)
                - "levenshtein": Levenshtein distance-based matching
            judge_lm: Optional separate LM to use as judge when expected_output is None.

        Returns:
            An evaluation function that performs structured extraction and compares
            with expected output (or uses judge if expected_output is None).
        """
        # Check if original evaluate_fn is a custom judge callable
        custom_judge_fn = None
        if (
            callable(self.evaluate_fn)
            and not isinstance(self.evaluate_fn, str)
            and not isinstance(self.evaluate_fn, dspy.LM)
        ):
            custom_judge_fn = self.evaluate_fn

        return default_evaluate_fn(
            lm=lm,
            model=self.model,
            system_prompt=self.system_prompt,
            instruction_prompt=self.instruction_prompt,
            metric=metric,
            judge_lm=judge_lm,
            custom_judge_fn=custom_judge_fn,
            exclude_fields=self.exclude_fields,
            include_fields=self.include_fields,
            evaluator_config=evaluator_config,
        )

    def _resolve_evaluate_fn(self, lm: dspy.LM) -> Callable[..., float]:
        """Resolve evaluate_fn from raw value to a callable.

        Args:
            lm: The DSPy language model.

        Returns:
            Resolved evaluation function.

        Raises:
            ValueError: If evaluate_fn is an invalid string.
            TypeError: If evaluate_fn has unexpected type.
        """
        raw = self.evaluate_fn
        evaluator_config_to_use = self.evaluator_config
        if evaluator_config_to_use is None and isinstance(raw, str):
            lower = raw.lower()
            if lower in ("exact", "levenshtein"):
                evaluator_config_to_use = {"default": lower, "field_overrides": {}}

        if raw is None:
            return self._default_evaluate_fn(lm, evaluator_config=evaluator_config_to_use)
        if isinstance(raw, str):
            lower = raw.lower()
            if lower in ("exact", "levenshtein"):
                return self._default_evaluate_fn(
                    lm, metric=lower, evaluator_config=evaluator_config_to_use
                )
            raise ValueError(
                f"evaluate_fn must be a callable, dspy.LM, None, or "
                f'one of ("exact", "levenshtein"), got "{raw}"'
            )
        if isinstance(raw, dspy.LM):
            return self._default_evaluate_fn(
                lm, judge_lm=raw, evaluator_config=evaluator_config_to_use
            )
        if callable(raw):
            return self._default_evaluate_fn(lm, evaluator_config=evaluator_config_to_use)
        raise TypeError(f"Unexpected type for evaluate_fn: {type(raw)}")

    def _create_metric_function(self, lm: dspy.LM) -> Callable[..., float]:
        """Create a metric function for DSPy optimization.

        Args:
            lm: The DSPy language model (needed for default evaluation function).

        Returns:
            A function that evaluates prompt performance.
        """
        evaluate_fn = self._resolve_evaluate_fn(lm)

        def metric_function(
            example: dspy.Example, prediction: dspy.Prediction, trace: Any = None
        ) -> float:
            """Evaluate the quality of optimized prompts and descriptions.

            Args:
                example: The DSPy example.
                prediction: The optimized field descriptions and prompts.
                trace: Optional trace from DSPy.

            Returns:
                A score between 0.0 and 1.0.
            """
            # Extract optimized field descriptions from prediction
            optimized_field_descriptions: dict[str, str] = {}
            optimized_system_prompt: str | None = None
            optimized_instruction_prompt: str | None = None

            for key, value in prediction.items():
                if key == "optimized_system_prompt":
                    optimized_system_prompt = value
                elif key == "optimized_instruction_prompt":
                    optimized_instruction_prompt = value
                elif key.startswith("optimized_"):
                    # Extract field path (remove "optimized_" prefix)
                    field_path = key.replace("optimized_", "")
                    optimized_field_descriptions[field_path] = value

            # If no optimized values provided (baseline evaluation), use original
            if not optimized_field_descriptions:
                optimized_field_descriptions = self.field_descriptions.copy()
            if optimized_system_prompt is None:
                optimized_system_prompt = self.system_prompt
            if optimized_instruction_prompt is None:
                optimized_instruction_prompt = self.instruction_prompt

            # Convert DSPy example to our Example type
            example_obj = self._dspy_example_to_example(example)

            # Use the evaluation function
            score = evaluate_fn(
                example_obj,
                optimized_field_descriptions,
                optimized_system_prompt,
                optimized_instruction_prompt,
            )

            # Ensure score is valid (between 0.0 and 1.0)
            if not isinstance(score, (int, float)) or score < 0.0 or score > 1.0:
                if self.verbose:
                    print(f"Warning: Invalid score {score}, defaulting to 0.0")
                return 0.0

            return float(score)

        return metric_function

    def _create_single_field_metric(
        self,
        field_path: str,
        current_descriptions: dict[str, str],
        evaluate_fn: Callable[..., float],
        optimized_demos: list[dict[str, Any]],
    ) -> Callable[..., float]:
        """Create a metric for single-field optimization that merges with current_descriptions."""

        def single_field_metric(
            example: dspy.Example, prediction: dspy.Prediction, trace: Any = None
        ) -> float:
            merged = dict(current_descriptions)
            new_val = getattr(prediction, f"optimized_{field_path}", None)
            if new_val is not None:
                merged[field_path] = new_val
            example_obj = self._dspy_example_to_example(example)
            if self._evaluate_fn_accepts_optimized_demos(evaluate_fn):
                score = evaluate_fn(
                    example_obj,
                    merged,
                    self.system_prompt,
                    self.instruction_prompt,
                    optimized_demos=optimized_demos,
                )
            else:
                score = evaluate_fn(
                    example_obj,
                    merged,
                    self.system_prompt,
                    self.instruction_prompt,
                )
            if not isinstance(score, (int, float)) or score < 0.0 or score > 1.0:
                return 0.0
            return float(score)

        return single_field_metric

    def _dspy_example_to_example(self, dspy_ex: dspy.Example) -> Example:
        """Convert a DSPy example to our Example object.

        Args:
            dspy_ex: DSPy example object.

        Returns:
            Example object.
        """
        # Extract input_data and expected_output from DSPy example
        input_data = getattr(dspy_ex, "input_data", {})
        expected_output = getattr(dspy_ex, "expected_output", None)
        # Only use {} as default if expected_output attribute doesn't exist
        # If it exists but is None, keep it as None
        if not hasattr(dspy_ex, "expected_output"):
            expected_output = {}

        # Reconstruct Example from input_data dictionary
        # input_data can contain "text" and/or "images" keys
        # Images might be dspy.Image objects or base64 strings
        text = input_data.get("text") if isinstance(input_data, dict) else None
        images = input_data.get("images") if isinstance(input_data, dict) else None
        images_base64 = (
            input_data.get("images_base64")
            if isinstance(input_data, dict)
            else None
        )

        # Create Example - if we have images, use image_base64 (first image)
        # Prefer images_base64 (original base64) if available,
        # otherwise try to extract from images
        if images_base64 and isinstance(images_base64, list) and len(images_base64) > 0:
            return Example(
                image_base64=images_base64[0],
                expected_output=expected_output,
            )
        elif images and isinstance(images, list) and len(images) > 0:
            # If images are dspy.Image objects, we need to get base64 from them
            # For now, try to get base64 from the original input_data if available
            # Otherwise, create example with text
            if text:
                return Example(
                    text=text,
                    expected_output=expected_output,
                )
            else:
                # Fallback: try to extract base64 from dspy.Image if possible
                # dspy.Image objects have a url attribute that might be a data URL
                first_image = images[0]
                if hasattr(first_image, "url"):
                    # Extract base64 from data URL if it's a data URL
                    url = first_image.url
                    if url.startswith("data:image"):
                        # Extract base64 part from data URL
                        base64_part = url.split(",")[-1] if "," in url else None
                        if base64_part:
                            return Example(
                                image_base64=base64_part,
                                expected_output=expected_output,
                            )
                        else:
                            return Example(
                                text="",
                                expected_output=expected_output,
                            )
                    else:
                        return Example(
                            text="",
                            expected_output=expected_output,
                        )
                else:
                    return Example(
                        text="",
                        expected_output=expected_output,
                    )
        elif text:
            return Example(
                text=text,
                expected_output=expected_output,
            )
        else:
            # Fallback: create a minimal example
            return Example(
                text="",
                expected_output=expected_output,
            )

    def _prepare_dspy_examples(
        self,
        descriptions_override: dict[str, str] | None = None,
    ) -> list[dspy.Example]:
        """Prepare examples as DSPy examples.

        Args:
            descriptions_override: Optional dict of field descriptions to use
                instead of self.field_descriptions. Used for sequential mode
                when optimizing a subset of fields.

        Returns:
            List of dspy.Example objects.
        """
        descriptions = descriptions_override or self.field_descriptions
        trainset = []
        input_keys = list(descriptions.keys())

        # Add field type keys to input keys (only for fields in descriptions)
        for field_path in descriptions.keys():
            if f"field_type_{field_path}" not in input_keys:
                input_keys.append(f"field_type_{field_path}")

        # Add prompts to input keys if they exist
        if self.system_prompt is not None:
            input_keys.append("system_prompt")
        if self.instruction_prompt is not None:
            input_keys.append("instruction_prompt")

        for ex in self.examples:
            # Convert input_data to dict if it's a Pydantic model
            input_data = ex.input_data
            if isinstance(input_data, BaseModel):
                input_data = input_data.model_dump()

            # Convert base64 images to dspy.Image objects if present
            # This allows DSPy to properly handle images in signatures
            if isinstance(input_data, dict) and "images" in input_data:
                base64_images = input_data.get("images")
                if base64_images:
                    dspy_images = convert_images_to_dspy_images(base64_images)
                    # Replace base64 strings with dspy.Image objects
                    # Keep original base64 in a separate key for backward compatibility
                    input_data = input_data.copy()
                    input_data["images"] = dspy_images
                    input_data["images_base64"] = base64_images  # Keep original for reference

            # Convert expected_output to dict if it's a Pydantic model or string
            expected_output = ex.expected_output
            if isinstance(expected_output, BaseModel):
                expected_output = expected_output.model_dump()
            elif isinstance(expected_output, str):
                # Convert string to dict format matching OutputModel structure
                expected_output = {"output": expected_output}

            example_dict: dict[str, Any] = {
                "input_data": input_data,
                "expected_output": expected_output,
            }
            # Add field descriptions as inputs
            example_dict.update(descriptions)

            # Add field types as inputs (with field_type_ prefix to distinguish from descriptions)
            for field_path in descriptions.keys():
                field_type = self.field_types.get(field_path, "")
                example_dict[f"field_type_{field_path}"] = field_type

            # Add prompts as inputs if they exist
            if self.system_prompt is not None:
                example_dict["system_prompt"] = self.system_prompt
            if self.instruction_prompt is not None:
                # Format instruction prompt template with example's text_dict for optimization
                # The optimizer will see formatted versions in examples, but we'll preserve
                # the template structure when optimizing
                if ex.text_dict:
                    formatted_instruction = format_instruction_prompt_template(
                        self.instruction_prompt, ex.text_dict
                    )
                    example_dict["instruction_prompt"] = (
                        formatted_instruction or self.instruction_prompt
                    )
                else:
                    example_dict["instruction_prompt"] = self.instruction_prompt

            trainset.append(dspy.Example(**example_dict).with_inputs(*input_keys))

        return trainset

    def _create_teleprompter(self, metric: Callable[..., float]) -> Teleprompter:
        """Create a Teleprompter instance with the given metric."""
        if self.custom_optimizer is not None:
            return self.custom_optimizer
        if self.optimizer_type == "miprov2zeroshot":
            default_kwargs = {
                "metric": metric,
                "num_threads": self.num_threads,
                "init_temperature": self.init_temperature,
                "auto": "light",
                "max_bootstrapped_demos": 0,
                "max_labeled_demos": 0,
            }
            return MIPROv2(**{**default_kwargs, **self.optimizer_kwargs})
        if self.optimizer_type == "miprov2":
            default_kwargs = {
                "metric": metric,
                "num_threads": self.num_threads,
                "init_temperature": self.init_temperature,
                "auto": "light",
            }
            if len(self.examples) <= 10:
                default_kwargs["max_bootstrapped_demos"] = 2
                default_kwargs["max_labeled_demos"] = 2
            return MIPROv2(**{**default_kwargs, **self.optimizer_kwargs})
        teleprompter_classes = self._get_teleprompter_subclasses()
        if self.optimizer_type not in teleprompter_classes:
            valid_optimizers = sorted(teleprompter_classes.keys())
            raise ValueError(
                f"Unknown optimizer_type: {self.optimizer_type}. Valid options: {valid_optimizers}"
            )
        optimizer_class = teleprompter_classes[self.optimizer_type]
        default_kwargs = {"metric": metric}
        # Fix: pass num_threads to all optimizers that support it (not just MIPROv2)
        try:
            sig = inspect.signature(optimizer_class.__init__)
            if "num_threads" in sig.parameters:
                default_kwargs["num_threads"] = self.num_threads
        except (ValueError, TypeError):
            pass
        if self.optimizer_type == "bootstrapfewshot" and len(self.examples) < 5:
            default_kwargs["max_bootstrapped_demos"] = max(1, min(2, len(self.examples) - 1))
        return optimizer_class(**{**default_kwargs, **self.optimizer_kwargs})

    @staticmethod
    def _evaluate_fn_accepts_optimized_demos(fn: Callable[..., float]) -> bool:
        """Check if evaluate_fn accepts optimized_demos keyword argument."""
        try:
            sig = inspect.signature(fn)
            return "optimized_demos" in sig.parameters
        except (ValueError, TypeError):
            return False

    def _optimize_fields_parallel(
        self,
        sorted_fields: list[str],
        current_descriptions: dict[str, str],
        train_examples: list[dspy.Example],
        val_examples: list[dspy.Example],
        evaluate_fn: Callable[..., float],
        baseline_avg: float,
        optimized_demos: list[dict[str, Any]],
        _emit: Callable[..., None],
    ) -> None:
        """Optimize multiple fields in parallel threads.

        Updates current_descriptions in place with optimized descriptions from all fields.
        Each field optimization runs independently with a snapshot of current_descriptions.
        """
        total_fields = len(sorted_fields)

        def optimize_field_task(idx_and_path: tuple[int, str]) -> tuple[str, str, float]:
            """Optimize a single field and return (field_path, optimized_desc, score)."""
            idx, field_path = idx_and_path
            depth = field_path.count(".")
            if self.verbose:
                print(
                    f"  Field {idx}/{total_fields}: {field_path} (depth {depth}) ... ",
                    end="",
                    flush=True,
                )

            score_before = baseline_avg  # Each field optimizes from baseline independently
            new_desc, new_score = self._optimize_single_field(
                field_path=field_path,
                current_descriptions=current_descriptions,  # Snapshot at start
                train_examples=train_examples,
                val_examples=val_examples,
                evaluate_fn=evaluate_fn,
                baseline_score=baseline_avg,
                optimized_demos=optimized_demos,
            )
            improved = new_score > baseline_avg
            if self.verbose:
                if improved:
                    print(f"{baseline_avg:.2%} → {new_score:.2%} ✓")
                else:
                    print(f"{baseline_avg:.2%} (no improvement)")

            _emit("fields", score_before, new_score, field_path=field_path, field_index=idx,
                  optimized_value=new_desc)
            return field_path, new_desc, new_score

        # Run all field optimizations in parallel
        with ThreadPoolExecutor(max_workers=self.num_threads) as executor:
            futures = {
                executor.submit(optimize_field_task, (idx, fp)): fp
                for idx, fp in enumerate(sorted_fields, 1)
            }

            for future in as_completed(futures):
                field_path, new_desc, _ = future.result()
                current_descriptions[field_path] = new_desc

    def _optimize_single_field(
        self,
        field_path: str,
        current_descriptions: dict[str, str],
        train_examples: list[dspy.Example],
        val_examples: list[dspy.Example],
        evaluate_fn: Callable[..., float],
        baseline_score: float,
        optimized_demos: list[dict[str, Any]],
    ) -> tuple[str, float]:
        """Optimize a single field description, holding all others fixed.

        Args:
            field_path: Field path to optimize.
            current_descriptions: Current best descriptions for all fields.
            train_examples: DSPy training examples.
            val_examples: DSPy validation examples.
            evaluate_fn: Evaluation function for scoring.
            baseline_score: Current best score to beat.
            optimized_demos: Few-shot demos for extraction.

        Returns:
            Tuple of (best_description, best_score).
        """
        single_field_descriptions = {field_path: current_descriptions[field_path]}
        program = PydanticOptimizerModule(
            field_descriptions=single_field_descriptions,
            field_types={field_path: self.field_types.get(field_path, "")},
            has_system_prompt=False,
            has_instruction_prompt=False,
        )

        train_single = self._prepare_dspy_examples(
            descriptions_override=single_field_descriptions,
        )
        val_single = self._prepare_dspy_examples(
            descriptions_override=single_field_descriptions,
        )
        split_idx = len(train_examples)
        train_single = train_single[:split_idx]
        val_single = val_single[split_idx:] if split_idx < len(train_single) else train_single

        metric = self._create_single_field_metric(
            field_path, current_descriptions, evaluate_fn, optimized_demos
        )
        optimizer = self._create_teleprompter(metric)

        optimizers_with_valset = (
            "miprov2zeroshot",
            "miprov2",
            "gepa",
            "bootstrapfewshotwithrandomsearch",
            "copro",
            "simba",
            "custom",
        )
        if self.optimizer_type in optimizers_with_valset:
            try:
                optimized_program = optimizer.compile(
                    program,
                    trainset=train_single,
                    valset=val_single,
                )
            except TypeError:
                optimized_program = optimizer.compile(
                    program,
                    trainset=train_single,
                )
        else:
            optimized_program = optimizer.compile(
                program,
                trainset=train_single,
            )

        program_args: dict[str, Any] = {
            field_path: current_descriptions[field_path],
            f"field_type_{field_path}": self.field_types.get(field_path, ""),
        }
        test_result = optimized_program(**program_args)
        new_description = getattr(
            test_result,
            f"optimized_{field_path}",
            current_descriptions[field_path],
        )

        merged_descriptions = dict(current_descriptions)
        merged_descriptions[field_path] = new_description

        scores = []
        for val_ex in val_examples:
            example_obj = self._dspy_example_to_example(val_ex)
            if self._evaluate_fn_accepts_optimized_demos(evaluate_fn):
                score = evaluate_fn(
                    example_obj,
                    merged_descriptions,
                    self.system_prompt,
                    self.instruction_prompt,
                    optimized_demos=optimized_demos,
                )
            else:
                score = evaluate_fn(
                    example_obj,
                    merged_descriptions,
                    self.system_prompt,
                    self.instruction_prompt,
                )
            scores.append(score)

        new_score = sum(scores) / len(scores) if scores else 0.0

        if new_score >= baseline_score:
            return (new_description, new_score)
        return (current_descriptions[field_path], baseline_score)

    def _optimize_prompt(
        self,
        prompt_type: str,
        current_descriptions: dict[str, str],
        current_system_prompt: str | None,
        current_instruction_prompt: str | None,
        train_examples: list[dspy.Example],
        val_examples: list[dspy.Example],
        evaluate_fn: Callable[..., float],
        baseline_score: float,
        optimized_demos: list[dict[str, Any]],
    ) -> tuple[str | None, float]:
        """Optimize system or instruction prompt, holding field descriptions fixed.

        Args:
            prompt_type: "system" or "instruction".
            current_descriptions: Optimized field descriptions from Phase 1.
            current_system_prompt: Current system prompt.
            current_instruction_prompt: Current instruction prompt.
            train_examples: DSPy training examples.
            val_examples: DSPy validation examples.
            evaluate_fn: Evaluation function.
            baseline_score: Current best score.
            optimized_demos: Few-shot demos.

        Returns:
            Tuple of (best_prompt_value, best_score).
        """
        if prompt_type == "system":
            if current_system_prompt is None:
                return (None, baseline_score)
            program = PydanticOptimizerModule(
                field_descriptions={},
                field_types={},
                has_system_prompt=True,
                has_instruction_prompt=False,
            )
            current_prompt = current_system_prompt
            attr_name = "optimized_system_prompt"
        else:
            if current_instruction_prompt is None:
                return (None, baseline_score)
            program = PydanticOptimizerModule(
                field_descriptions={},
                field_types={},
                has_system_prompt=False,
                has_instruction_prompt=True,
            )
            current_prompt = current_instruction_prompt
            attr_name = "optimized_instruction_prompt"

        metric = self._create_metric_function(dspy.settings.lm)
        optimizer = self._create_teleprompter(metric)

        optimizers_with_valset = (
            "miprov2zeroshot",
            "miprov2",
            "gepa",
            "bootstrapfewshotwithrandomsearch",
            "copro",
            "simba",
            "custom",
        )
        if self.optimizer_type in optimizers_with_valset:
            try:
                optimized_program = optimizer.compile(
                    program,
                    trainset=train_examples,
                    valset=val_examples,
                )
            except TypeError:
                optimized_program = optimizer.compile(
                    program,
                    trainset=train_examples,
                )
        else:
            optimized_program = optimizer.compile(
                program,
                trainset=train_examples,
            )

        program_args: dict[str, Any] = {}
        if prompt_type == "system":
            program_args["system_prompt"] = current_system_prompt
        else:
            program_args["instruction_prompt"] = current_instruction_prompt

        test_result = optimized_program(**program_args)
        new_prompt = getattr(test_result, attr_name, current_prompt)

        if prompt_type == "system":
            merged_system = new_prompt
            merged_instruction = current_instruction_prompt
        else:
            merged_system = current_system_prompt
            merged_instruction = new_prompt

        scores = []
        for val_ex in val_examples:
            example_obj = self._dspy_example_to_example(val_ex)
            if self._evaluate_fn_accepts_optimized_demos(evaluate_fn):
                score = evaluate_fn(
                    example_obj,
                    current_descriptions,
                    merged_system,
                    merged_instruction,
                    optimized_demos=optimized_demos,
                )
            else:
                score = evaluate_fn(
                    example_obj,
                    current_descriptions,
                    merged_system,
                    merged_instruction,
                )
            scores.append(score)

        new_score = sum(scores) / len(scores) if scores else 0.0

        if new_score >= baseline_score:
            return (new_prompt, new_score)
        if prompt_type == "system":
            return (current_system_prompt, baseline_score)
        return (current_instruction_prompt, baseline_score)

    def _optimize_sequential(
        self,
        lm: dspy.LM,
        evaluate_fn: Callable[..., float],
        train_examples: list[dspy.Example],
        val_examples: list[dspy.Example],
        optimized_demos: list[dict[str, Any]],
    ) -> OptimizationResult:
        """Run sequential optimization: fields deepest-first, then prompts."""
        _t0 = time.perf_counter()

        def _emit(phase, score_before, score_after, field_path=None, field_index=None, optimized_value=None):
            if self.on_progress is None:
                return
            try:
                self.on_progress(FieldOptimizationProgress(
                    phase=phase, score_before=score_before, score_after=score_after,
                    improved=score_after > score_before, total_fields=total_fields,
                    field_path=field_path, field_index=field_index,
                    elapsed_seconds=time.perf_counter() - _t0,
                    optimized_value=optimized_value,
                ))
            except Exception:
                pass  # never abort optimization due to callback error

        baseline_scores = []
        for val_ex in val_examples:
            example_obj = self._dspy_example_to_example(val_ex)
            baseline_score = evaluate_fn(
                example_obj,
                self.field_descriptions,
                self.system_prompt,
                self.instruction_prompt,
            )
            baseline_scores.append(baseline_score)

        baseline_avg = sum(baseline_scores) / len(baseline_scores) if baseline_scores else 0.0
        if self.verbose:
            print(f"Baseline average score: {baseline_avg:.2%}")
            print("\nPhase 1: Optimizing field descriptions (deepest-first)...")

        current_descriptions = dict(self.field_descriptions)
        current_score = baseline_avg
        sorted_fields = self._sort_fields_by_depth()
        total_fields = len(sorted_fields)

        _emit("baseline", baseline_avg, baseline_avg)

        # Slice val_examples if max_val_examples is set
        effective_val_examples = val_examples
        if self.max_val_examples is not None:
            effective_val_examples = val_examples[:self.max_val_examples]
            if self.verbose and len(effective_val_examples) < len(val_examples):
                print(f"(Using {len(effective_val_examples)}/{len(val_examples)} validation examples)")

        if self.parallel_fields:
            # Parallel field optimization
            self._optimize_fields_parallel(
                sorted_fields=sorted_fields,
                current_descriptions=current_descriptions,
                train_examples=train_examples,
                val_examples=effective_val_examples,
                evaluate_fn=evaluate_fn,
                baseline_avg=baseline_avg,
                optimized_demos=optimized_demos,
                _emit=_emit,
            )
            # After parallel optimization, recalculate current_score
            scores = []
            for val_ex in effective_val_examples:
                example_obj = self._dspy_example_to_example(val_ex)
                score = evaluate_fn(
                    example_obj,
                    current_descriptions,
                    self.system_prompt,
                    self.instruction_prompt,
                )
                scores.append(score)
            current_score = sum(scores) / len(scores) if scores else baseline_avg
        else:
            # Sequential field optimization (original behavior)
            for idx, field_path in enumerate(sorted_fields, 1):
                depth = field_path.count(".")
                if self.verbose:
                    print(
                        f"  Field {idx}/{total_fields}: {field_path} (depth {depth}) ... ",
                        end="",
                        flush=True,
                    )

                # Check skip threshold
                if self.skip_score_threshold is not None:
                    field_baseline_scores = []
                    for val_ex in effective_val_examples:
                        example_obj = self._dspy_example_to_example(val_ex)
                        score = evaluate_fn(
                            example_obj,
                            current_descriptions,
                            self.system_prompt,
                            self.instruction_prompt,
                        )
                        field_baseline_scores.append(score)
                    field_baseline = sum(field_baseline_scores) / len(field_baseline_scores) if field_baseline_scores else 0.0

                    if field_baseline >= self.skip_score_threshold:
                        if self.verbose:
                            print(f"{field_baseline:.2%} (skipped, already above {self.skip_score_threshold:.0%})")
                        _emit("skipped", field_baseline, field_baseline, field_path=field_path, field_index=idx)
                        continue

                score_before = current_score
                new_desc, new_score = self._optimize_single_field(
                    field_path=field_path,
                    current_descriptions=current_descriptions,
                    train_examples=train_examples,
                    val_examples=effective_val_examples,
                    evaluate_fn=evaluate_fn,
                    baseline_score=current_score,
                    optimized_demos=optimized_demos,
                )
                improved = new_score > current_score
                current_descriptions[field_path] = new_desc
                if self.verbose:
                    if improved:
                        print(f"{current_score:.2%} → {new_score:.2%} ✓")
                    else:
                        print(f"{current_score:.2%} (no improvement)")
                current_score = new_score
                _emit("fields", score_before, new_score, field_path=field_path, field_index=idx,
                      optimized_value=new_desc)

        optimized_system_prompt = self.system_prompt
        optimized_instruction_prompt = self.instruction_prompt

        if self.system_prompt is not None:
            if self.verbose:
                print("\nPhase 2a: Optimizing system_prompt ... ", end="", flush=True)
            score_before_sys = current_score
            optimized_system_prompt, current_score = self._optimize_prompt(
                "system",
                current_descriptions,
                optimized_system_prompt,
                optimized_instruction_prompt,
                train_examples,
                effective_val_examples,
                evaluate_fn,
                current_score,
                optimized_demos,
            )
            if self.verbose:
                print(f"{current_score:.2%}")
            _emit("system_prompt", score_before_sys, current_score,
                  optimized_value=optimized_system_prompt)

        if self.instruction_prompt is not None:
            if self.verbose:
                print("Phase 2b: Optimizing instruction_prompt ... ", end="", flush=True)
            score_before_instr = current_score
            optimized_instruction_prompt, current_score = self._optimize_prompt(
                "instruction",
                current_descriptions,
                optimized_system_prompt,
                optimized_instruction_prompt,
                train_examples,
                effective_val_examples,
                evaluate_fn,
                current_score,
                optimized_demos,
            )
            if self.verbose:
                print(f"{current_score:.2%}")
            _emit("instruction_prompt", score_before_instr, current_score,
                  optimized_value=optimized_instruction_prompt)

        improvement = current_score - baseline_avg
        improvement_pct = (improvement / baseline_avg * 100) if baseline_avg > 0 else 0.0

        if improvement < 0:
            if self.verbose:
                print(f"\n⚠️  WARNING: Optimization decreased performance by {abs(improvement):.2%}")
                print("Keeping original prompts and descriptions.")
            current_descriptions = self.field_descriptions.copy()
            optimized_system_prompt = self.system_prompt
            optimized_instruction_prompt = self.instruction_prompt
            current_score = baseline_avg
            improvement = 0.0
            improvement_pct = 0.0

        api_calls = 0
        total_tokens = 0
        if hasattr(lm, "history") and lm.history:
            api_calls = len(lm.history)
            for call in lm.history:
                if isinstance(call, dict):
                    usage = call.get("usage", {})
                    if isinstance(usage, dict):
                        total_tokens += usage.get("total_tokens", 0)

        if self.verbose:
            _console = Console()
            table = Table(title="Optimization Complete", box=box.ROUNDED)
            table.add_column("Metric", style="bold")
            table.add_column("Value")
            table.add_row("Baseline score", f"{baseline_avg:.2%}")
            table.add_row("Final score", f"{current_score:.2%}")
            if improvement > 0:
                table.add_row("Improvement", f"{improvement:+.2%}")
            elif improvement < 0:
                table.add_row("Improvement", f"{improvement:+.2%} [red](decreased)[/]")
            else:
                table.add_row("Improvement", "No change")
            if api_calls > 0:
                table.add_row("API calls", str(api_calls))
            if total_tokens > 0:
                table.add_row("Total tokens", f"{total_tokens:,}")
            _console.print(table)

        _emit("complete", baseline_avg, current_score)

        return OptimizationResult(
            optimized_descriptions=current_descriptions,
            optimized_system_prompt=optimized_system_prompt,
            optimized_instruction_prompt=optimized_instruction_prompt,
            metrics={
                "average_score": current_score,
                "baseline_score": baseline_avg,
                "improvement": improvement,
                "improvement_percent": improvement_pct,
                "validation_size": len(val_examples),
                "training_size": len(train_examples),
            },
            baseline_score=baseline_avg,
            optimized_score=current_score,
            optimized_demos=optimized_demos,
            api_calls=api_calls,
            total_tokens=total_tokens,
            estimated_cost_usd=None,
        )

    def optimize(self) -> OptimizationResult:
        """Optimize the Pydantic model field descriptions using DSPy.

        Returns:
            OptimizationResult containing optimized descriptions and metrics.
        """
        if self.verbose:
            print(f"\n{'='*60}")
            print("Starting DSPy Pydantic optimization")
            print(f"{'='*60}")
            print(f"Model: {self.model.__name__}")
            print(f"Optimizer: {self.optimizer_type.upper()}")
            print(f"Examples: {len(self.examples)}")
            effective_count = len(self._get_effective_fields())
            print(
                f"Fields to optimize: {effective_count}"
                + (
                    f" (of {len(self.field_descriptions)} total)"
                    if effective_count != len(self.field_descriptions)
                    else ""
                )
            )
            if self.field_descriptions:
                print("\nInitial field descriptions (set during initialization):")
                for field_path, description in self.field_descriptions.items():
                    print(f"  {field_path}: {description}")
            print(f"Optimization threads: {self.num_threads}")
            print(f"{'='*60}\n")

        # Use configured DSPy LM (should be set via dspy.configure())
        if dspy.settings.lm is None:
            raise ValueError(
                "DSPy must be configured before optimization. "
                "Call dspy.configure(lm=dspy.LM(...)) first."
            )
        lm = dspy.settings.lm
        evaluate_fn = self._resolve_evaluate_fn(lm)

        effective = self._get_effective_fields()
        effective_descriptions = {
            k: v for k, v in self.field_descriptions.items() if k in effective
        }
        program = PydanticOptimizerModule(
            field_descriptions=effective_descriptions or self.field_descriptions,
            field_types=self.field_types,
            has_system_prompt=self.system_prompt is not None,
            has_instruction_prompt=self.instruction_prompt is not None,
        )

        # Prepare examples for DSPy
        trainset = self._prepare_dspy_examples()

        # Split into train and validation sets
        # Ensure at least one example in trainset (needed for optimizers like MIPROv2)
        split_idx = max(1, int(len(trainset) * self.train_split))
        train_examples = trainset[:split_idx]
        val_examples = trainset[split_idx:] if split_idx < len(trainset) else trainset

        # Few-shot demos: up to 8 training examples for the extraction prompt
        max_few_shot = min(8, split_idx)
        optimized_demos: list[dict[str, Any]] = []
        for ex in self.examples[:max_few_shot]:
            inp = ex.input_data
            out = ex.expected_output
            if isinstance(inp, BaseModel):
                inp = inp.model_dump()
            if isinstance(out, BaseModel):
                out = out.model_dump()
            optimized_demos.append({"input_data": inp, "expected_output": out})

        if self.verbose:
            print(f"Training examples: {len(train_examples)}")
            print(f"Validation examples: {len(val_examples)}")

        if self.sequential:
            return self._optimize_sequential(
                lm=lm,
                evaluate_fn=evaluate_fn,
                train_examples=train_examples,
                val_examples=val_examples,
                optimized_demos=optimized_demos,
            )

        # Single-pass optimization (sequential=False)
        metric = self._create_metric_function(lm)
        optimizer = self._create_teleprompter(metric)

        if self.verbose and self.custom_optimizer is not None:
            print(f"Using custom optimizer: {type(optimizer).__name__}")

        # Evaluate baseline (original prompts and descriptions) on validation set
        if self.verbose:
            print("\nEvaluating baseline configuration...")

        baseline_scores = []
        for val_ex in val_examples:
            # Convert DSPy example to our Example object
            example_obj = self._dspy_example_to_example(val_ex)
            # Use original prompts and descriptions (no optimization)
            baseline_score = evaluate_fn(
                example_obj,
                self.field_descriptions,
                self.system_prompt,
                self.instruction_prompt,
            )
            baseline_scores.append(baseline_score)

        baseline_avg = (
            sum(baseline_scores) / len(baseline_scores) if baseline_scores else 0.0
        )
        if self.verbose:
            print(f"Baseline average score: {baseline_avg:.2%}")

        # Optimize
        if self.verbose:
            print("\nOptimizing prompts and field descriptions...")
            if self.system_prompt:
                print("  - System prompt")
            if self.instruction_prompt:
                print("  - Instruction prompt")
            if self.field_descriptions:
                print(f"  - {len(self.field_descriptions)} field descriptions")

        # Some optimizers support valset, others don't
        # Try to use valset if supported, fall back to trainset only if not
        optimizers_with_valset = (
            "miprov2zeroshot",
            "miprov2",
            "gepa",
            "bootstrapfewshotwithrandomsearch",
            "copro",
            "simba",
            "custom",  # Custom optimizers might support valset
        )

        if self.optimizer_type in optimizers_with_valset:
            try:
                optimized_program = optimizer.compile(
                    program,
                    trainset=train_examples,
                    valset=val_examples,
                )
            except TypeError:
                # If valset is not supported, fall back to trainset only
                if self.verbose:
                    print("Warning: Optimizer doesn't support valset, using trainset only")
                optimized_program = optimizer.compile(
                    program,
                    trainset=train_examples,
                )
        else:
            optimized_program = optimizer.compile(
                program,
                trainset=train_examples,
            )

        # Build arguments for optimized program (field descriptions, types, and prompts)
        program_args: dict[str, Any] = {}
        program_args.update(self.field_descriptions)
        # Add field types with field_type_ prefix
        for field_path, field_type in self.field_types.items():
            program_args[f"field_type_{field_path}"] = field_type
        if self.system_prompt is not None:
            program_args["system_prompt"] = self.system_prompt
        if self.instruction_prompt is not None:
            program_args["instruction_prompt"] = self.instruction_prompt

        # Test the optimized program to get optimized values
        test_result = optimized_program(**program_args)

        optimized_field_descriptions = dict(self.field_descriptions)
        for field_path in effective_descriptions.keys():
            attr_name = f"optimized_{field_path}"
            if hasattr(test_result, attr_name):
                optimized_field_descriptions[field_path] = getattr(
                    test_result, attr_name
                )

        # Extract optimized prompts
        optimized_system_prompt: str | None = None
        optimized_instruction_prompt: str | None = None
        if self.system_prompt is not None:
            if hasattr(test_result, "optimized_system_prompt"):
                optimized_system_prompt = getattr(test_result, "optimized_system_prompt")
        if self.instruction_prompt is not None:
            if hasattr(test_result, "optimized_instruction_prompt"):
                optimized_instruction_prompt = getattr(
                    test_result, "optimized_instruction_prompt"
                )

        # Evaluate optimized config on validation set
        if self.verbose:
            print("\nEvaluating optimized configuration...")

        evaluation_scores = []
        for val_ex in val_examples:
            # Get optimized descriptions and prompts for this example
            val_program_args: dict[str, Any] = {}
            for field_path in self.field_descriptions.keys():
                if hasattr(val_ex, field_path):
                    val_program_args[field_path] = getattr(val_ex, field_path)
                else:
                    val_program_args[field_path] = self.field_descriptions[field_path]

            # Add field types
            for field_path in self.field_types.keys():
                field_type_key = f"field_type_{field_path}"
                if hasattr(val_ex, field_type_key):
                    val_program_args[field_type_key] = getattr(val_ex, field_type_key)
                else:
                    val_program_args[field_type_key] = self.field_types[field_path]

            if self.system_prompt is not None:
                val_program_args["system_prompt"] = self.system_prompt
            if self.instruction_prompt is not None:
                val_program_args["instruction_prompt"] = self.instruction_prompt

            prediction = optimized_program(**val_program_args)

            # Extract optimized descriptions and prompts from prediction
            pred_descriptions: dict[str, str] = {}
            pred_system_prompt: str | None = None
            pred_instruction_prompt: str | None = None

            for field_path in self.field_descriptions.keys():
                attr_name = f"optimized_{field_path}"
                if hasattr(prediction, attr_name):
                    pred_descriptions[field_path] = getattr(prediction, attr_name)

            if self.system_prompt is not None:
                if hasattr(prediction, "optimized_system_prompt"):
                    pred_system_prompt = getattr(prediction, "optimized_system_prompt")

            if self.instruction_prompt is not None:
                if hasattr(prediction, "optimized_instruction_prompt"):
                    pred_instruction_prompt = getattr(
                        prediction, "optimized_instruction_prompt"
                    )

            # Convert DSPy example to our Example object
            example_obj = self._dspy_example_to_example(val_ex)
            if self._evaluate_fn_accepts_optimized_demos(evaluate_fn):
                score = evaluate_fn(
                    example_obj,
                    pred_descriptions,
                    pred_system_prompt,
                    pred_instruction_prompt,
                    optimized_demos=optimized_demos,
                )
            else:
                score = evaluate_fn(
                    example_obj,
                    pred_descriptions,
                    pred_system_prompt,
                    pred_instruction_prompt,
                )
            evaluation_scores.append(score)

        avg_score = (
            sum(evaluation_scores) / len(evaluation_scores) if evaluation_scores else 0.0
        )

        # Compare with baseline
        improvement = avg_score - baseline_avg
        improvement_pct = (
            (improvement / baseline_avg * 100) if baseline_avg > 0 else 0.0
        )

        # Only use optimized prompts/descriptions if they improve performance
        if improvement < 0:
            if self.verbose:
                print(
                    f"\n⚠️  WARNING: Optimization decreased performance by {abs(improvement):.2%}"
                )
                print("Keeping original prompts and descriptions instead of optimized ones.")
            optimized_field_descriptions = self.field_descriptions.copy()
            optimized_system_prompt = self.system_prompt
            optimized_instruction_prompt = self.instruction_prompt
            avg_score = baseline_avg
            improvement = 0.0
            improvement_pct = 0.0

        # Track API usage from DSPy LM history
        api_calls = 0
        total_tokens = 0
        estimated_cost_usd = None

        if hasattr(lm, "history") and lm.history:
            api_calls = len(lm.history)
            for call in lm.history:
                if isinstance(call, dict):
                    usage = call.get("usage", {})
                    if isinstance(usage, dict):
                        total_tokens += usage.get("total_tokens", 0)

        # Build result
        result = OptimizationResult(
            optimized_descriptions=optimized_field_descriptions,
            optimized_system_prompt=optimized_system_prompt,
            optimized_instruction_prompt=optimized_instruction_prompt,
            metrics={
                "average_score": avg_score,
                "baseline_score": baseline_avg,
                "improvement": improvement,
                "improvement_percent": improvement_pct,
                "validation_size": len(val_examples),
                "training_size": len(train_examples),
            },
            baseline_score=baseline_avg,
            optimized_score=avg_score,
            optimized_demos=optimized_demos,
            api_calls=api_calls,
            total_tokens=total_tokens,
            estimated_cost_usd=estimated_cost_usd,
        )

        if self.verbose:
            _console = Console()
            table = Table(title="Optimization Complete", box=box.ROUNDED)
            table.add_column("Metric", style="bold")
            table.add_column("Value")
            table.add_row("Baseline score", f"{baseline_avg:.2%}")
            table.add_row("Final score", f"{avg_score:.2%}")
            if improvement > 0:
                table.add_row("Improvement", f"{improvement:+.2%}")
            elif improvement < 0:
                table.add_row("Improvement", f"{improvement:+.2%} [red](decreased)[/]")
            else:
                table.add_row("Improvement", "No change")
            if api_calls > 0:
                table.add_row("API calls", str(api_calls))
            if total_tokens > 0:
                table.add_row("Total tokens", f"{total_tokens:,}")
            _console.print(table)

        return result

