"""Experiment model with ActiveRecord-style API."""

from __future__ import annotations

import asyncio
import inspect
import json
import logging
import time
import warnings
from datetime import datetime
from typing import Any, Callable

import httpx
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from lumenova_beacon.datasets.types import PaginatedResponse
from lumenova_beacon.exceptions import (
    ExperimentError,
    ExperimentNotFoundError,
    ExperimentValidationError,
)
from lumenova_beacon.experiments.types import (
    EnrichedRecord,
    ExperimentRun,
    ExperimentRunStats,
    ExperimentRunStatus,
    ExperimentStatus,
    ExperimentStep,
    ExperimentStepType,
)
from lumenova_beacon.utils.client_helpers import get_base_url, get_transport
from lumenova_beacon.utils.datetime import parse_iso_datetime
from lumenova_beacon.utils.http_errors import HTTPErrorHandler

logger = logging.getLogger(__name__)


# Centralized error handler for experiments
_error_handler = HTTPErrorHandler(
    not_found_exc=ExperimentNotFoundError,
    validation_exc=ExperimentValidationError,
    base_exc=ExperimentError,
)

# Retry decorator for async functions
_retry_async = retry(
    retry=retry_if_exception_type(
        (httpx.NetworkError, httpx.TimeoutException, httpx.HTTPStatusError)
    ),
    wait=wait_exponential(multiplier=1, min=1, max=10),
    stop=stop_after_attempt(3),
    reraise=True,
)

# Terminal statuses for polling loop
_TERMINAL_STATUSES = {
    ExperimentStatus.COMPLETED,
    ExperimentStatus.FAILED,
    ExperimentStatus.STOPPED,
    ExperimentStatus.COMPLETED_WITH_ERRORS,
}


class Experiment:
    """Experiment model with ActiveRecord-style API.

    Experiments are pipeline-based evaluation runs that execute a sequence of
    steps (prompt, LLM, evaluation, external agent) against dataset records.

    Examples:
        Create an experiment with pipeline steps:
            >>> from lumenova_beacon.experiments import (
            ...     Experiment, ExperimentStep, ExperimentStepType,
            ... )
            >>> experiment = Experiment.create(
            ...     name="My Pipeline",
            ...     dataset_id="dataset-uuid",
            ...     steps=[
            ...         ExperimentStep(
            ...             step_type=ExperimentStepType.PROMPT,
            ...             output_column="rendered",
            ...             config={"prompt_content": "Summarize: {{input}}",
            ...                     "variable_mappings": {"input": "text"}},
            ...         ),
            ...         ExperimentStep(
            ...             step_type=ExperimentStepType.LLM,
            ...             output_column="response",
            ...             config={"model_config_id": "config-uuid",
            ...                     "variable_mappings": {"input": "rendered"}},
            ...         ),
            ...     ],
            ... )

        Fire-and-forget start:
            >>> experiment.start()

        Start with external agents (blocks until agent steps done):
            >>> experiment.start(external_agents={
            ...     "My Agent": my_agent_fn,
            ... })
    """

    def __init__(
        self,
        id: str,
        name: str,
        dataset_id: str,
        status: int,
        created_at: datetime,
        description: str | None = None,
        result_dataset_id: str | None = None,
        execution_type: str | None = None,
        started_at: datetime | None = None,
        completed_at: datetime | None = None,
        updated_at: datetime | None = None,
        dataset_name: str | None = None,
        record_count: int = 0,
        step_count: int = 0,
        steps: list[ExperimentStep] | None = None,
        total_runs: int = 0,
        completed_runs: int = 0,
        failed_runs: int = 0,
        pending_runs: int = 0,
        completed_records: int = 0,
        failed_records: int = 0,
        project_id: str | None = None,
    ):
        self.id = id
        self.name = name
        self.description = description
        self.dataset_id = dataset_id
        self.result_dataset_id = result_dataset_id
        self.execution_type = execution_type
        self.status = ExperimentStatus(status)
        self.started_at = started_at
        self.completed_at = completed_at
        self.created_at = created_at
        self.updated_at = updated_at
        self.dataset_name = dataset_name
        self.record_count = record_count
        self.step_count = step_count
        self.steps = steps
        self.total_runs = total_runs
        self.completed_runs = completed_runs
        self.failed_runs = failed_runs
        self.pending_runs = pending_runs
        self.completed_records = completed_records
        self.failed_records = failed_records
        self.project_id = project_id

    # =========================================================================
    # Async CRUD Methods
    # =========================================================================

    @classmethod
    @_retry_async
    async def aget(cls, experiment_id: str) -> Experiment:
        """Get a single experiment by ID (async).

        Returns the experiment with its steps included.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{experiment_id}"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.get(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                return cls._from_dict(response.json())
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to fetch experiment: {e}")

    @classmethod
    @_retry_async
    async def acreate(
        cls,
        name: str,
        dataset_id: str,
        steps: list[ExperimentStep | dict[str, Any]],
        *,
        execution_type: str = "server",
        clone_dataset: bool = True,
        output_dataset_name: str | None = None,
        description: str | None = None,
    ) -> Experiment:
        """Create a new experiment with pipeline steps (async).

        The experiment is created in DRAFT status. Use astart() to begin execution.

        Args:
            name: Name of the experiment.
            dataset_id: Source dataset UUID.
            steps: Pipeline steps in execution order. Each can be an
                ExperimentStep object or a raw dict with step_type,
                output_column, and config.
            execution_type: "server" (default) or "local" for SDK-managed.
            clone_dataset: Whether to clone the source dataset for results.
            output_dataset_name: Custom name for the cloned result dataset.
            description: Optional description.

        Returns:
            Created Experiment instance in DRAFT status.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments"

        step_dicts = [
            s.to_dict() if isinstance(s, ExperimentStep) else s for s in steps
        ]

        payload: dict[str, Any] = {
            "name": name,
            "dataset_id": dataset_id,
            "steps": step_dicts,
            "execution_type": execution_type,
            "clone_dataset": clone_dataset,
        }
        if description is not None:
            payload["description"] = description
        if output_dataset_name is not None:
            payload["output_dataset_name"] = output_dataset_name

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.post(
                    url,
                    json=payload,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                data = response.json()
                logger.debug("Created experiment: %s", data["id"])
                return cls._from_dict(data)
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to create experiment: {e}")

    @classmethod
    @_retry_async
    async def alist(
        cls,
        page: int = 1,
        page_size: int = 20,
        search: str | None = None,
        status: ExperimentStatus | int | None = None,
    ) -> tuple[list[Experiment], PaginatedResponse]:
        """List experiments with pagination and filtering (async)."""
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments"

        params: dict[str, Any] = {"page": page, "page_size": page_size}
        if search:
            params["search"] = search
        if status is not None:
            params["status"] = (
                int(status) if isinstance(status, ExperimentStatus) else status
            )

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.get(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                    params=params,
                )
                response.raise_for_status()
                data = response.json()
                experiments = [cls._from_dict(item) for item in data["items"]]
                pagination = PaginatedResponse(
                    total=data["total"],
                    page=data["page"],
                    page_size=data["page_size"],
                    total_pages=data["total_pages"],
                )
                return experiments, pagination
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to list experiments: {e}")

    @_retry_async
    async def aupdate(
        self,
        *,
        name: str | None = None,
        description: str | None = None,
    ) -> None:
        """Update experiment metadata (name/description only)."""
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}"

        payload: dict[str, Any] = {}
        if name is not None:
            payload["name"] = name
        if description is not None:
            payload["description"] = description
        if not payload:
            return

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.patch(
                    url,
                    json=payload,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                if name is not None:
                    self.name = name
                if description is not None:
                    self.description = description
                logger.debug("Updated experiment: %s", self.id)
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to update experiment: {e}")

    @_retry_async
    async def aupdate_draft(
        self,
        *,
        name: str | None = None,
        description: str | None = None,
        dataset_id: str | None = None,
        steps: list[ExperimentStep | dict[str, Any]] | None = None,
        clone_dataset: bool | None = None,
        output_dataset_name: str | None = None,
    ) -> Experiment:
        """Fully update a draft experiment (async).

        Unlike aupdate(), this can modify steps and dataset.
        Only works when experiment is in DRAFT status.

        Returns:
            Updated Experiment instance.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/draft"

        payload: dict[str, Any] = {}
        if name is not None:
            payload["name"] = name
        if description is not None:
            payload["description"] = description
        if dataset_id is not None:
            payload["dataset_id"] = dataset_id
        if steps is not None:
            payload["steps"] = [
                s.to_dict() if isinstance(s, ExperimentStep) else s for s in steps
            ]
        if clone_dataset is not None:
            payload["clone_dataset"] = clone_dataset
        if output_dataset_name is not None:
            payload["output_dataset_name"] = output_dataset_name

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.put(
                    url,
                    json=payload,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                data = response.json()
                self._update_from_response(data)
                logger.debug("Updated draft experiment: %s", self.id)
                return self
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to update draft experiment: {e}")

    @_retry_async
    async def adelete(self) -> None:
        """Delete this experiment (async)."""
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.delete(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                logger.debug("Deleted experiment: %s", self.id)
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to delete experiment: {e}")

    # =========================================================================
    # Async Execution Control
    # =========================================================================

    async def astart(
        self,
        external_agents: dict[str, Callable] | None = None,
        *,
        poll_interval: float = 2.0,
        trace_calls: bool = True,
        show_progress: bool = True,
        max_concurrency: int = 1,
    ) -> Experiment:
        """Start or continue an experiment (async).

        With no args: fire-and-forget (starts the experiment, returns immediately).
        With external_agents: starts/continues + executes matching external agent
        steps, blocking until all matching agent steps are done.

        Args:
            external_agents: Mapping of agent_name -> callable (sync or async,
                or LangChain Runnable) for external_agent steps. The agent_name
                must match the ``agent_name`` field in the step's config.
            poll_interval: Seconds between polling cycles when external_agents
                is provided.
            trace_calls: Whether to create tracing spans for agent invocations.

        Returns:
            Self with updated status.
        """
        # Start experiment if still in DRAFT
        if self.status == ExperimentStatus.DRAFT:
            await self._post_start()

        # Fire-and-forget if no agents provided
        if not external_agents:
            self._warn_if_has_external_steps()
            return self

        # Ensure steps are loaded (needed for agent_name lookup)
        if not self.steps:
            updated = await Experiment.aget(self.id)
            self._update_from_response_obj(updated)

        # Build lookup: step_id -> ExperimentStep (external_agent steps only)
        external_step_map = {
            step.id: step
            for step in (self.steps or [])
            if step.step_type
            in (ExperimentStepType.EXTERNAL_AGENT, "external_agent")
        }

        # Initialize progress bar
        pbar = None
        if show_progress:
            from lumenova_beacon.experiments._progress import _ProgressBar

            pbar = _ProgressBar(total=None, desc=self.name or self.id)

        logger.debug(
            "Polling for experiment %s (status=%s, external_steps=%d)",
            self.id,
            self.status.name,
            len(external_step_map),
        )

        # Polling loop
        try:
            while True:
                updated = await Experiment.aget(self.id)
                self._update_from_response_obj(updated)

                logger.debug(
                    "Poll: status=%s records=%d total_runs=%d completed=%d failed=%d",
                    self.status.name,
                    self.record_count,
                    self.total_runs,
                    self.completed_runs,
                    self.failed_runs,
                )

                if self.status == ExperimentStatus.DRAFT:
                    # Experiment was reset externally — start it again
                    await self._post_start()
                    continue

                if self.status in _TERMINAL_STATUSES:
                    logger.debug("Experiment reached terminal status: %s", self.status.name)
                    return self

                # Fetch all pending runs
                pending_runs = await self._fetch_all_pending_runs()
                logger.debug(
                    "Fetched %d pending run(s) total, %d for external-agent step(s)",
                    len(pending_runs),
                    sum(1 for r in pending_runs if r.step_id in external_step_map),
                )

                # Filter to external_agent step runs only
                external_pending = [
                    r for r in pending_runs if r.step_id in external_step_map
                ]

                # Lock in total on the first iteration where we see external runs:
                # pending count + already completed = full external-agent run count
                if pbar is not None and pbar._total is None and external_pending:
                    pbar.set_total(len(external_pending) + pbar._n)

                if not external_pending:
                    # No external runs yet — server still processing other steps
                    logger.debug("No external-agent runs pending yet, sleeping %.1fs", poll_interval)
                    await asyncio.sleep(poll_interval)
                    continue

                # Split into runs we can handle vs. runs for other agents
                matching: list[tuple[ExperimentRun, str, ExperimentStep]] = []
                for run in external_pending:
                    step = external_step_map[run.step_id]
                    agent_name = step.config.get("agent_name", "")
                    if agent_name in external_agents:
                        matching.append((run, agent_name, step))

                logger.debug(
                    "%d external-agent run(s) pending, %d match provided agents %s",
                    len(external_pending),
                    len(matching),
                    list(external_agents.keys()),
                )

                if matching:
                    # Execute all matching runs
                    # Build record cache only for runs without server-resolved input
                    needs_record = [
                        run for run, _, _ in matching
                        if not (run.result or {}).get("resolved_input")
                    ]
                    record_cache = (
                        await self._build_record_cache(
                            {r.record_id for r in needs_record}
                        )
                        if needs_record
                        else {}
                    )

                    sem = asyncio.Semaphore(max_concurrency)

                    async def _run_one(run: ExperimentRun, agent_name: str, step: ExperimentStep) -> None:
                        agent_fn = external_agents[agent_name]
                        resolved = (run.result or {}).get("resolved_input")
                        if resolved:
                            input_value = resolved.get("input", "")
                        else:
                            record_data = record_cache.get(run.record_id, {})
                            input_value = _resolve_variable_mappings(
                                step.config, record_data
                            ).get("input", "")
                        logger.debug("Executing run %s (agent=%s)", run.id, agent_name)
                        async with sem:
                            await self._execute_run(run, agent_fn, input_value, trace_calls)
                        logger.debug("Completed run %s", run.id)
                        if pbar is not None:
                            pbar.update(1)

                    await asyncio.gather(*[
                        _run_one(run, agent_name, step)
                        for run, agent_name, step in matching
                    ])
                    continue  # Check for more work immediately

                # Pending external runs exist but none match our agents — done
                logger.debug(
                    "External runs exist but none match provided agents %s — stopping",
                    list(external_agents.keys()),
                )
                return self
        finally:
            if pbar is not None:
                pbar.close()

    @_retry_async
    async def _post_start(self) -> None:
        """POST /start to transition DRAFT -> QUEUED."""
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/start"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.post(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                data = response.json()
                self.status = ExperimentStatus(data["status"])
                if data.get("started_at"):
                    self.started_at = parse_iso_datetime(data["started_at"])
                if data.get("updated_at"):
                    self.updated_at = parse_iso_datetime(data["updated_at"])
                logger.debug("Started experiment: %s", self.id)
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to start experiment: {e}")

    @_retry_async
    async def acancel(self) -> Experiment:
        """Stop a running or queued experiment (async).

        The experiment transitions to STOPPED status. Use continue_experiment()
        to resume execution from where it left off.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/cancel"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.post(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                data = response.json()
                self.status = ExperimentStatus(data["status"])
                if data.get("completed_at"):
                    self.completed_at = parse_iso_datetime(data["completed_at"])
                if data.get("updated_at"):
                    self.updated_at = parse_iso_datetime(data["updated_at"])
                logger.debug("Cancelled experiment: %s", self.id)
                return self
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to cancel experiment: {e}")

    @_retry_async
    async def aclone(self) -> Experiment:
        """Clone this experiment (async). Returns a new Experiment in DRAFT."""
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/clone"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.post(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                data = response.json()
                logger.debug("Cloned experiment %s -> %s", self.id, data["id"])
                return Experiment._from_dict(data)
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to clone experiment: {e}")

    @_retry_async
    async def areset(self) -> Experiment:
        """Reset this experiment back to DRAFT status (async).

        Clears all execution timestamps, progress counters, and deletes all runs.
        Only works for experiments in terminal states: COMPLETED, FAILED,
        COMPLETED_WITH_ERRORS, STOPPED.

        Returns:
            Self with updated state (status=DRAFT, counters cleared).
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/reset"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.post(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                data = response.json()
                self._update_from_response(data)
                logger.debug("Reset experiment to DRAFT: %s", self.id)
                return self
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to reset experiment: {e}")

    async def acontinue_experiment(
        self,
        external_agents: dict[str, Callable] | None = None,
        *,
        poll_interval: float = 2.0,
        trace_calls: bool = True,
        show_progress: bool = True,
        max_concurrency: int = 1,
    ) -> Experiment:
        """Continue a stopped or waiting experiment from where it left off (async).

        For STOPPED experiments: re-queues execution for records that didn't
        complete all pipeline steps. Records that fully completed are preserved.

        For WAITING_FOR_EXTERNAL experiments: resumes the external agent polling
        loop, equivalent to calling astart() with external_agents.

        Args:
            external_agents: Mapping of agent_name -> callable for external_agent
                steps. Required when the experiment has external agent steps.
            poll_interval: Seconds between polling cycles when external_agents
                is provided.
            trace_calls: Whether to create tracing spans for agent invocations.
            max_concurrency: Maximum number of runs to execute in parallel.

        Returns:
            Self with updated state.
        """
        if self.status == ExperimentStatus.WAITING_FOR_EXTERNAL:
            # Already waiting for agents — go straight into the polling loop
            return await self.astart(
                external_agents=external_agents,
                poll_interval=poll_interval,
                trace_calls=trace_calls,
                show_progress=show_progress,
                max_concurrency=max_concurrency,
            )

        # STOPPED -> POST /continue to re-queue, then optionally poll
        await self._post_continue()

        if external_agents:
            return await self.astart(
                external_agents=external_agents,
                poll_interval=poll_interval,
                trace_calls=trace_calls,
                show_progress=show_progress,
                max_concurrency=max_concurrency,
            )

        self._warn_if_has_external_steps()
        return self

    @_retry_async
    async def _post_continue(self) -> None:
        """POST /continue to transition STOPPED -> QUEUED."""
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/continue"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.post(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                data = response.json()
                self._update_from_response(data)
                logger.debug("Continuing experiment: %s", self.id)
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to continue experiment: {e}")

    @_retry_async
    async def aretry_failed(self) -> dict[str, Any]:
        """Retry all failed runs (async).

        Returns:
            Dict with retried_count, experiment_id, and retried_run_ids.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/retry-failed"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.post(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                logger.debug("Retried failed runs for experiment: %s", self.id)
                return response.json()
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to retry failed runs: {e}")

    @_retry_async
    async def abulk_rerun(self, run_ids: list[str]) -> dict[str, Any]:
        """Rerun specific runs by ID (async).

        Args:
            run_ids: List of run UUIDs to rerun.

        Returns:
            Dict with retried_count, experiment_id, and retried_run_ids.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/runs/bulk-rerun"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.post(
                    url,
                    json={"run_ids": run_ids},
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                logger.debug("Bulk rerun %d runs for experiment: %s", len(run_ids), self.id)
                return response.json()
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to bulk rerun runs: {e}")

    # =========================================================================
    # Async Progress & Results
    # =========================================================================

    @_retry_async
    async def aget_progress(self) -> dict[str, Any]:
        """Get real-time execution progress (async).

        Returns:
            Dict with experiment_id, total_records, total_steps, total_runs,
            completed_runs, failed_runs, completed_records, failed_records,
            progress_percent.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/progress"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.get(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                return response.json()
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to get experiment progress: {e}")

    @_retry_async
    async def alist_records(
        self,
        page: int = 1,
        page_size: int = 20,
        status: str | None = None,
    ) -> tuple[list[EnrichedRecord], PaginatedResponse]:
        """List enriched records with per-step execution status (async).

        Args:
            page: Page number (starting from 1).
            page_size: Items per page.
            status: Filter by record status (pending, running, success, failed).

        Returns:
            Tuple of (enriched records, pagination metadata).
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/records"

        params: dict[str, Any] = {"page": page, "page_size": page_size}
        if status is not None:
            params["status"] = status

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.get(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                    params=params,
                )
                response.raise_for_status()
                data = response.json()
                records = [EnrichedRecord.from_dict(item) for item in data["items"]]
                pagination = PaginatedResponse(
                    total=data["total"],
                    page=data["page"],
                    page_size=data["page_size"],
                    total_pages=data["total_pages"],
                )
                return records, pagination
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to list experiment records: {e}")

    @_retry_async
    async def alist_runs(
        self,
        page: int = 1,
        page_size: int = 20,
        status: ExperimentRunStatus | str | None = None,
        step_id: str | None = None,
    ) -> tuple[list[ExperimentRun], PaginatedResponse]:
        """List individual experiment runs (async).

        Args:
            page: Page number (starting from 1).
            page_size: Items per page.
            status: Filter by run status.
            step_id: Filter by step UUID.

        Returns:
            Tuple of (runs, pagination metadata).
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/runs"

        params: dict[str, Any] = {"page": page, "page_size": page_size}
        if status is not None:
            params["status"] = (
                status.value if isinstance(status, ExperimentRunStatus) else status
            )
        if step_id is not None:
            params["step_id"] = step_id

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.get(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                    params=params,
                )
                response.raise_for_status()
                data = response.json()
                runs = [ExperimentRun.from_dict(item) for item in data["items"]]
                pagination = PaginatedResponse(
                    total=data["total"],
                    page=data["page"],
                    page_size=data["page_size"],
                    total_pages=data["total_pages"],
                )
                return runs, pagination
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to list experiment runs: {e}")

    @_retry_async
    async def aget_run(self, run_id: str) -> ExperimentRun:
        """Get a specific run by ID (async)."""
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/runs/{run_id}"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.get(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                return ExperimentRun.from_dict(response.json())
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to get experiment run: {e}")

    @_retry_async
    async def aget_run_stats(self) -> ExperimentRunStats:
        """Get aggregated run statistics (async)."""
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/run-stats"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.get(
                    url,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                return ExperimentRunStats.from_dict(response.json())
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to get run stats: {e}")

    @classmethod
    @_retry_async
    async def apreview_output(
        cls,
        dataset_id: str,
        configurations: list[dict[str, Any]],
        record_ids: list[str],
        variable_mappings: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Preview experiment output without creating an experiment (async).

        Executes configurations against sample records to preview results.
        Limited to 3 records maximum.

        Args:
            dataset_id: Source dataset UUID.
            configurations: List of configuration dicts, each with:
                label, prompt_content, model_config_id, model_parameters.
            record_ids: Dataset record UUIDs to preview (max 3).
            variable_mappings: Variable-to-column mapping.

        Returns:
            Dict with ``results`` (list of per-run preview results) and
            ``total_execution_time_ms``.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/preview-output"

        payload: dict[str, Any] = {
            "dataset_id": dataset_id,
            "configurations": configurations,
            "record_ids": record_ids,
        }
        if variable_mappings is not None:
            payload["variable_mappings"] = variable_mappings

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.post(
                    url,
                    json=payload,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                return response.json()
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to preview experiment output: {e}")

    # =========================================================================
    # Async SDK-managed (local) experiment methods
    # =========================================================================

    @_retry_async
    async def aupdate_status(self, status: ExperimentStatus | int) -> Experiment:
        """Update experiment status for SDK-managed experiments (async).

        Only valid for experiments with execution_type='local'.

        Args:
            status: New status (RUNNING, COMPLETED, FAILED, COMPLETED_WITH_ERRORS).
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/status"

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.patch(
                    url,
                    json={"status": int(status)},
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                data = response.json()
                self._update_from_response(data)
                logger.debug("Updated experiment status to %s: %s", status, self.id)
                return self
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to update experiment status: {e}")

    @_retry_async
    async def aupdate_progress(
        self,
        *,
        total_runs: int | None = None,
        completed_runs: int | None = None,
        failed_runs: int | None = None,
    ) -> Experiment:
        """Update progress counters for SDK-managed experiments (async).

        Only valid for experiments with execution_type='local'.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/progress"

        payload: dict[str, Any] = {}
        if total_runs is not None:
            payload["total_runs"] = total_runs
        if completed_runs is not None:
            payload["completed_runs"] = completed_runs
        if failed_runs is not None:
            payload["failed_runs"] = failed_runs

        if not payload:
            return self

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.patch(
                    url,
                    json=payload,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                data = response.json()
                self._update_from_response(data)
                return self
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to update experiment progress: {e}")

    @_retry_async
    async def aupdate_run(
        self,
        run_id: str,
        *,
        output: str | None = None,
        status: str = "completed",
        error_message: str | None = None,
        execution_time_ms: int | None = None,
        execution_metadata: dict[str, Any] | None = None,
    ) -> ExperimentRun:
        """Write back results for an external_agent run (async).

        Args:
            run_id: Run UUID to update.
            output: Output text from the agent.
            status: "completed" or "failed".
            error_message: Error details if failed.
            execution_time_ms: Execution time in milliseconds.
            execution_metadata: Additional metadata (model, tokens, etc.).

        Returns:
            Updated ExperimentRun.
        """
        base_url = get_base_url()
        transport = get_transport("Experiment operations")
        url = f"{base_url}/api/v1/experiments/{self.id}/runs/{run_id}"

        payload: dict[str, Any] = {"status": status}
        if output is not None:
            payload["output"] = output
        if error_message is not None:
            payload["error_message"] = error_message
        if execution_time_ms is not None:
            payload["execution_time_ms"] = execution_time_ms
        if execution_metadata is not None:
            payload["execution_metadata"] = execution_metadata

        try:
            async with httpx.AsyncClient(verify=transport.verify) as client:
                response = await client.patch(
                    url,
                    json=payload,
                    headers=transport.headers,
                    timeout=transport.timeout,
                )
                response.raise_for_status()
                return ExperimentRun.from_dict(response.json())
        except httpx.HTTPStatusError as e:
            _error_handler.handle(e)
        except httpx.HTTPError as e:
            raise ExperimentError(f"Failed to update experiment run: {e}")

    async def _fetch_all_pending_runs(self) -> list[ExperimentRun]:
        """Fetch all pending runs across all pages."""
        all_runs: list[ExperimentRun] = []
        page = 1

        while True:
            runs, pagination = await self.alist_runs(
                page=page, page_size=100, status=ExperimentRunStatus.PENDING
            )
            all_runs.extend(runs)
            if page >= pagination.total_pages:
                break
            page += 1

        return all_runs

    async def _build_record_cache(
        self, record_ids: set[str]
    ) -> dict[str, dict[str, Any]]:
        """Build a mapping of record_id -> record data for lookups."""
        cache: dict[str, dict[str, Any]] = {}
        page = 1

        while True:
            records, pagination = await self.alist_records(
                page=page, page_size=100
            )
            for record in records:
                if record.record_id in record_ids:
                    cache[record.record_id] = record.data
            # Stop if we've found all needed records or exhausted pages
            if len(cache) >= len(record_ids) or page >= pagination.total_pages:
                break
            page += 1

        return cache

    async def _execute_run(
        self,
        run: ExperimentRun,
        agent_fn: Callable,
        input_value: str,
        trace_calls: bool,
    ) -> None:
        """Execute a single pending run locally and write back the result."""
        start_time = time.monotonic()
        is_runnable = _is_langchain_runnable(agent_fn)
        trace_id = None

        try:
            if is_runnable:
                raw_output, trace_id = await _invoke_runnable(
                    agent_fn, input_value, trace_calls, self.id, run.record_id
                )
                output_value = _extract_runnable_output(raw_output)
            else:
                raw_output, trace_id = await _invoke_callable(
                    agent_fn, input_value, trace_calls, self.id, run.record_id
                )
                output_value = _extract_output(raw_output)

            elapsed_ms = int((time.monotonic() - start_time) * 1000)
            output_str = (
                output_value
                if isinstance(output_value, str)
                else json.dumps(output_value, default=str)
            )
            await self.aupdate_run(
                run.id,
                output=output_str,
                status="completed",
                execution_time_ms=elapsed_ms,
                execution_metadata={"trace_id": trace_id} if trace_id else None,
            )

        except Exception as e:
            elapsed_ms = int((time.monotonic() - start_time) * 1000)
            logger.warning("Error executing run %s: %s", run.id, e)
            await self.aupdate_run(
                run.id,
                status="failed",
                error_message=str(e),
                execution_time_ms=elapsed_ms,
                execution_metadata={"trace_id": trace_id} if trace_id else None,
            )

    # =========================================================================
    # Async Dataset accessors
    # =========================================================================

    async def aget_dataset(self):
        """Get the source Dataset instance (async)."""
        from lumenova_beacon.datasets import Dataset

        try:
            return await Dataset.aget(self.dataset_id)
        except Exception as e:
            raise ExperimentError(f"Failed to get source dataset: {e}") from e

    async def aget_result_dataset(self):
        """Get the result Dataset instance (async). Returns None if not yet complete."""
        if self.result_dataset_id is None:
            return None

        from lumenova_beacon.datasets import Dataset

        try:
            return await Dataset.aget(self.result_dataset_id)
        except Exception as e:
            raise ExperimentError(f"Failed to get result dataset: {e}") from e

    # =========================================================================
    # Sync wrappers
    # =========================================================================

    @classmethod
    def get(cls, experiment_id: str) -> Experiment:
        """Get a single experiment by ID."""
        return asyncio.run(cls.aget(experiment_id))

    @classmethod
    def create(
        cls,
        name: str,
        dataset_id: str,
        steps: list[ExperimentStep | dict[str, Any]],
        *,
        execution_type: str = "server",
        clone_dataset: bool = True,
        output_dataset_name: str | None = None,
        description: str | None = None,
    ) -> Experiment:
        """Create a new experiment with pipeline steps."""
        return asyncio.run(
            cls.acreate(
                name=name,
                dataset_id=dataset_id,
                steps=steps,
                execution_type=execution_type,
                clone_dataset=clone_dataset,
                output_dataset_name=output_dataset_name,
                description=description,
            )
        )

    @classmethod
    def list(
        cls,
        page: int = 1,
        page_size: int = 20,
        search: str | None = None,
        status: ExperimentStatus | int | None = None,
    ) -> tuple[list[Experiment], PaginatedResponse]:
        """List experiments with pagination and filtering."""
        return asyncio.run(
            cls.alist(page=page, page_size=page_size, search=search, status=status)
        )

    def update(
        self,
        *,
        name: str | None = None,
        description: str | None = None,
    ) -> None:
        """Update experiment metadata (name/description only)."""
        asyncio.run(self.aupdate(name=name, description=description))

    def update_draft(
        self,
        *,
        name: str | None = None,
        description: str | None = None,
        dataset_id: str | None = None,
        steps: list[ExperimentStep | dict[str, Any]] | None = None,
        clone_dataset: bool | None = None,
        output_dataset_name: str | None = None,
    ) -> Experiment:
        """Fully update a draft experiment."""
        return asyncio.run(
            self.aupdate_draft(
                name=name,
                description=description,
                dataset_id=dataset_id,
                steps=steps,
                clone_dataset=clone_dataset,
                output_dataset_name=output_dataset_name,
            )
        )

    def delete(self) -> None:
        """Delete this experiment."""
        asyncio.run(self.adelete())

    def start(
        self,
        external_agents: dict[str, Callable] | None = None,
        *,
        poll_interval: float = 2.0,
        trace_calls: bool = True,
        show_progress: bool = True,
        max_concurrency: int = 1,
    ) -> Experiment:
        """Start or continue an experiment.

        With no args: fire-and-forget (starts the experiment).
        With external_agents: starts + executes matching external agent steps.
        """
        return asyncio.run(
            self.astart(
                external_agents=external_agents,
                poll_interval=poll_interval,
                trace_calls=trace_calls,
                show_progress=show_progress,
                max_concurrency=max_concurrency,
            )
        )

    def cancel(self) -> Experiment:
        """Stop a running or queued experiment."""
        return asyncio.run(self.acancel())

    def clone(self) -> Experiment:
        """Clone this experiment. Returns a new Experiment in DRAFT."""
        return asyncio.run(self.aclone())

    def reset(self) -> Experiment:
        """Reset this experiment back to DRAFT status."""
        return asyncio.run(self.areset())

    def continue_experiment(
        self,
        external_agents: dict[str, Callable] | None = None,
        *,
        poll_interval: float = 2.0,
        trace_calls: bool = True,
        show_progress: bool = True,
        max_concurrency: int = 1,
    ) -> Experiment:
        """Continue a stopped or waiting experiment from where it left off."""
        return asyncio.run(
            self.acontinue_experiment(
                external_agents=external_agents,
                poll_interval=poll_interval,
                trace_calls=trace_calls,
                show_progress=show_progress,
                max_concurrency=max_concurrency,
            )
        )

    def retry_failed(self) -> dict[str, Any]:
        """Retry all failed runs."""
        return asyncio.run(self.aretry_failed())

    def bulk_rerun(self, run_ids: list[str]) -> dict[str, Any]:
        """Rerun specific runs by ID."""
        return asyncio.run(self.abulk_rerun(run_ids))

    def get_progress(self) -> dict[str, Any]:
        """Get real-time execution progress."""
        return asyncio.run(self.aget_progress())

    def list_records(
        self,
        page: int = 1,
        page_size: int = 20,
        status: str | None = None,
    ) -> tuple[list[EnrichedRecord], PaginatedResponse]:
        """List enriched records with per-step execution status."""
        return asyncio.run(
            self.alist_records(page=page, page_size=page_size, status=status)
        )

    def list_runs(
        self,
        page: int = 1,
        page_size: int = 20,
        status: ExperimentRunStatus | str | None = None,
        step_id: str | None = None,
    ) -> tuple[list[ExperimentRun], PaginatedResponse]:
        """List individual experiment runs."""
        return asyncio.run(
            self.alist_runs(
                page=page, page_size=page_size, status=status, step_id=step_id
            )
        )

    def get_run(self, run_id: str) -> ExperimentRun:
        """Get a specific run by ID."""
        return asyncio.run(self.aget_run(run_id))

    def get_run_stats(self) -> ExperimentRunStats:
        """Get aggregated run statistics."""
        return asyncio.run(self.aget_run_stats())

    @classmethod
    def preview_output(
        cls,
        dataset_id: str,
        configurations: list[dict[str, Any]],
        record_ids: list[str],
        variable_mappings: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Preview experiment output without creating an experiment."""
        return asyncio.run(
            cls.apreview_output(
                dataset_id=dataset_id,
                configurations=configurations,
                record_ids=record_ids,
                variable_mappings=variable_mappings,
            )
        )

    def update_status(self, status: ExperimentStatus | int) -> Experiment:
        """Update experiment status for SDK-managed experiments."""
        return asyncio.run(self.aupdate_status(status))

    def update_progress(
        self,
        *,
        total_runs: int | None = None,
        completed_runs: int | None = None,
        failed_runs: int | None = None,
    ) -> Experiment:
        """Update progress counters for SDK-managed experiments."""
        return asyncio.run(
            self.aupdate_progress(
                total_runs=total_runs,
                completed_runs=completed_runs,
                failed_runs=failed_runs,
            )
        )

    def update_run(
        self,
        run_id: str,
        *,
        output: str | None = None,
        status: str = "completed",
        error_message: str | None = None,
        execution_time_ms: int | None = None,
        execution_metadata: dict[str, Any] | None = None,
    ) -> ExperimentRun:
        """Write back results for an external_agent run."""
        return asyncio.run(
            self.aupdate_run(
                run_id,
                output=output,
                status=status,
                error_message=error_message,
                execution_time_ms=execution_time_ms,
                execution_metadata=execution_metadata,
            )
        )

    def get_dataset(self):
        """Get the source Dataset instance."""
        return asyncio.run(self.aget_dataset())

    def get_result_dataset(self):
        """Get the result Dataset instance. Returns None if not yet complete."""
        return asyncio.run(self.aget_result_dataset())

    # =========================================================================
    # Internal helpers
    # =========================================================================

    def _warn_if_has_external_steps(self) -> None:
        """Warn if the experiment has external_agent steps but no agents were provided."""
        if not self.steps:
            return
        external_steps = [
            s for s in self.steps
            if s.step_type in (ExperimentStepType.EXTERNAL_AGENT, "external_agent")
        ]
        if external_steps:
            agent_names = [
                s.config.get("agent_name", "<unnamed>") for s in external_steps
            ]
            warnings.warn(
                f"Experiment '{self.id}' has external_agent steps "
                f"({', '.join(agent_names)}) but no external_agents were "
                f"provided. These steps will remain pending until agents "
                f"are supplied via start() or continue_experiment().",
                UserWarning,
                stacklevel=3,
            )

    @classmethod
    def _from_dict(cls, data: dict[str, Any]) -> Experiment:
        """Create an Experiment instance from API response data."""
        created_at = parse_iso_datetime(data["created_at"])

        # Parse steps if present (detail response)
        steps = None
        if data.get("steps"):
            steps = [ExperimentStep.from_dict(s) for s in data["steps"]]

        return cls(
            id=data["id"],
            name=data["name"],
            description=data.get("description"),
            dataset_id=data["dataset_id"],
            result_dataset_id=data.get("result_dataset_id"),
            execution_type=data.get("execution_type"),
            status=data["status"],
            started_at=parse_iso_datetime(data.get("started_at")),
            completed_at=parse_iso_datetime(data.get("completed_at")),
            created_at=created_at,
            updated_at=parse_iso_datetime(data.get("updated_at")),
            dataset_name=data.get("dataset_name"),
            record_count=data.get("record_count", 0),
            step_count=data.get("step_count", 0),
            steps=steps,
            total_runs=data.get("total_runs", 0),
            completed_runs=data.get("completed_runs", 0),
            failed_runs=data.get("failed_runs", 0),
            pending_runs=data.get("pending_runs", 0),
            completed_records=data.get("completed_records", 0),
            failed_records=data.get("failed_records", 0),
            project_id=data.get("project_id"),
        )

    def _update_from_response(self, data: dict[str, Any]) -> None:
        """Update this instance from an API response dict."""
        self.name = data.get("name", self.name)
        self.description = data.get("description", self.description)
        self.dataset_id = data.get("dataset_id", self.dataset_id)
        self.result_dataset_id = data.get("result_dataset_id", self.result_dataset_id)
        self.execution_type = data.get("execution_type", self.execution_type)
        if "status" in data:
            self.status = ExperimentStatus(data["status"])
        if data.get("started_at"):
            self.started_at = parse_iso_datetime(data["started_at"])
        if data.get("completed_at"):
            self.completed_at = parse_iso_datetime(data["completed_at"])
        if data.get("updated_at"):
            self.updated_at = parse_iso_datetime(data["updated_at"])
        self.dataset_name = data.get("dataset_name", self.dataset_name)
        self.record_count = data.get("record_count", self.record_count)
        self.step_count = data.get("step_count", self.step_count)
        if data.get("steps"):
            self.steps = [ExperimentStep.from_dict(s) for s in data["steps"]]
        self.total_runs = data.get("total_runs", self.total_runs)
        self.completed_runs = data.get("completed_runs", self.completed_runs)
        self.failed_runs = data.get("failed_runs", self.failed_runs)
        self.pending_runs = data.get("pending_runs", self.pending_runs)
        self.completed_records = data.get("completed_records", self.completed_records)
        self.failed_records = data.get("failed_records", self.failed_records)
        self.project_id = data.get("project_id", self.project_id)

    def _update_from_response_obj(self, other: Experiment) -> None:
        """Update this instance from another Experiment instance."""
        self.name = other.name
        self.description = other.description
        self.dataset_id = other.dataset_id
        self.result_dataset_id = other.result_dataset_id
        self.execution_type = other.execution_type
        self.status = other.status
        self.started_at = other.started_at
        self.completed_at = other.completed_at
        self.updated_at = other.updated_at
        self.dataset_name = other.dataset_name
        self.record_count = other.record_count
        self.step_count = other.step_count
        self.steps = other.steps
        self.total_runs = other.total_runs
        self.completed_runs = other.completed_runs
        self.failed_runs = other.failed_runs
        self.pending_runs = other.pending_runs
        self.completed_records = other.completed_records
        self.failed_records = other.failed_records
        self.project_id = other.project_id

    def __repr__(self) -> str:
        return (
            f"Experiment(id={self.id!r}, name={self.name!r}, "
            f"status={self.status.name}, step_count={self.step_count}, "
            f"record_count={self.record_count})"
        )


# =============================================================================
# Module-level helpers for agent execution (used by arun)
# =============================================================================


def _resolve_variable_mappings(
    config: dict[str, Any],
    record_data: dict[str, Any],
) -> dict[str, str]:
    """Resolve variable mappings from step config using record data.

    Mirrors server-side logic in step_executors._resolve_variable_mappings().

    Each mapping value can be:
    - A string: treated as a column name to look up from record_data
    - A dict with is_fixed=True: use fixed_value directly
    """
    variable_mappings = config.get("variable_mappings", {})
    resolved: dict[str, str] = {}
    for var_name, var_config in variable_mappings.items():
        if isinstance(var_config, dict) and var_config.get("is_fixed"):
            resolved[var_name] = var_config.get("fixed_value", "")
        else:
            source_column = var_config if isinstance(var_config, str) else ""
            value = record_data.get(source_column)
            resolved[var_name] = str(value) if value is not None else ""
    return resolved


def _is_langchain_runnable(fn: Any) -> bool:
    """Check if fn is a LangChain Runnable (lazy import)."""
    try:
        from langchain_core.runnables import Runnable

        return isinstance(fn, Runnable)
    except ImportError:
        return False


def _extract_output(raw_output: Any) -> Any:
    """Extract a serializable output value from callable result."""
    try:
        from langchain_core.messages import AIMessage

        if isinstance(raw_output, AIMessage):
            return raw_output.content
    except ImportError:
        pass

    if isinstance(raw_output, (str, dict, list, int, float, bool)):
        return raw_output

    return str(raw_output)


def _extract_runnable_output(raw_output: Any) -> Any:
    """Pull the final assistant message content out of a LangGraph state dict."""
    if isinstance(raw_output, dict) and "messages" in raw_output:
        messages = raw_output["messages"]
        if messages:
            last = messages[-1]
            if hasattr(last, "content"):
                return last.content
            if isinstance(last, dict) and "content" in last:
                return last["content"]
    return _extract_output(raw_output)


def _wrap_input_for_runnable(input_value: Any) -> Any:
    """Wrap a plain string/primitive as a LangGraph messages dict."""
    if isinstance(input_value, dict):
        return input_value

    try:
        from langchain_core.messages import HumanMessage

        return {"messages": [HumanMessage(content=str(input_value))]}
    except ImportError:
        return {"messages": [{"role": "user", "content": str(input_value)}]}


async def _invoke_runnable(
    fn: Any,
    input_value: Any,
    trace_calls: bool,
    experiment_id: str,
    record_id: str,
) -> tuple[Any, str | None]:
    """Invoke a LangChain Runnable with optional BeaconLangGraphHandler.

    Returns:
        Tuple of (output, trace_id). trace_id is None when trace_calls is False.
    """
    config: dict[str, Any] = {
        "configurable": {"thread_id": f"experiment:{experiment_id}:record:{record_id}"},
    }

    handler = None
    if trace_calls:
        from lumenova_beacon.tracing.integrations.langchain import (
            BeaconLangGraphHandler,
        )

        handler = BeaconLangGraphHandler(
            metadata={
                "experiment_id": experiment_id,
                "dataset_record_id": record_id,
                "source": "experiment_run",
            },
        )
        config["callbacks"] = [handler]

    wrapped_input = _wrap_input_for_runnable(input_value)
    result = await fn.ainvoke(wrapped_input, config=config)
    trace_id = handler._current_trace_id if handler else None
    return result, trace_id


async def _invoke_callable(
    fn: Any,
    mapped_input: Any,
    trace_calls: bool,
    experiment_id: str,
    record_id: str,
) -> tuple[Any, str | None]:
    """Invoke a plain callable with optional trace span.

    Returns:
        Tuple of (output, trace_id). trace_id is None when trace_calls is False.
    """
    is_async = inspect.iscoroutinefunction(fn)

    if trace_calls:
        from lumenova_beacon.core.client import get_client
        from lumenova_beacon.types import SpanType

        client = get_client()

        trace_name = getattr(fn, "__name__", None) or type(fn).__name__
        async with client.trace(
            trace_name,
            span_type=SpanType.FUNCTION,
        ) as span:
            span.set_metadata("experiment_id", experiment_id)
            span.set_metadata("dataset_record_id", record_id)
            span.set_metadata("source", "experiment_run")
            span.set_input(mapped_input)

            if is_async:
                result = await fn(mapped_input)
            else:
                result = await asyncio.to_thread(fn, mapped_input)

            span.set_output(result)
            return result, span.trace_id
    else:
        if is_async:
            return await fn(mapped_input), None
        else:
            return await asyncio.to_thread(fn, mapped_input), None
