# Copyright (c) Metis. All rights reserved.

"""Mantis Runs - Flexible metric tracking for experiments.

This module provides wandb-style experiment tracking with automatic crash detection
via heartbeats, signal handling, and context managers.

Example usage::

    import mantisdk

    # Recommended: Context manager (handles crashes automatically)
    with mantisdk.run("experiment-v1", config={"lr": 0.01}) as run:
        for epoch in range(100):
            loss = train()
            run.log({"loss": loss, "accuracy": acc}, step=epoch)

    # Alternative: Manual lifecycle (must call finish!)
    run = mantisdk.init_run("experiment-v1")
    run.log({"metric": value})
    run.finish()
"""

from __future__ import annotations

import asyncio
import atexit
import functools
import hashlib
import json
import logging
import os
import signal
import subprocess
import sys
import threading
import time
import uuid
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, TypeVar, Union

if TYPE_CHECKING:
    from mantisdk.store.listener import StorageListener

import requests

logger = logging.getLogger(__name__)

# Type variable for decorators
F = TypeVar("F", bound=Callable[..., Any])


# =============================================================================
# Reproducibility Helpers
# =============================================================================


def _compute_config_hash(config: Optional[Dict[str, Any]]) -> Optional[str]:
    """SHA256 hash of config for detecting identical configs.

    Returns first 16 characters of the hash for brevity.
    """
    if not config:
        return None
    try:
        config_str = json.dumps(config, sort_keys=True, default=str)
        return hashlib.sha256(config_str.encode()).hexdigest()[:16]
    except Exception:
        return None


def _get_git_info() -> Dict[str, Any]:
    """Capture git commit, branch, and dirty state.

    Returns empty dict if not in a git repo or git is not available.
    """
    try:
        # Get commit SHA (first 12 chars)
        sha = (
            subprocess.check_output(
                ["git", "rev-parse", "HEAD"],
                stderr=subprocess.DEVNULL,
                timeout=5,
            )
            .decode()
            .strip()[:12]
        )

        # Get branch name
        branch = (
            subprocess.check_output(
                ["git", "rev-parse", "--abbrev-ref", "HEAD"],
                stderr=subprocess.DEVNULL,
                timeout=5,
            )
            .decode()
            .strip()
        )

        # Check if working directory is dirty
        dirty = (
            subprocess.call(
                ["git", "diff", "--quiet"],
                stderr=subprocess.DEVNULL,
                timeout=5,
            )
            != 0
        )

        return {"gitSha": sha, "gitBranch": branch, "gitDirty": dirty}
    except Exception:
        return {}


# =============================================================================
# Types
# =============================================================================

RunState = Literal["pending", "running", "completed", "failed", "cancelled", "crashed"]

# =============================================================================
# Constants
# =============================================================================

MAX_RUN_NAME_LENGTH = 256
MAX_METRIC_BUFFER_SIZE = 10000


@dataclass
class Metric:
    """A self-describing metric with metadata."""

    name: str
    value: float
    step: Optional[int] = None
    unit: str = ""
    metric_type: str = "gauge"
    labels: Dict[str, str] = field(default_factory=dict)


# =============================================================================
# Run and Rollout Context Management
# =============================================================================

_current_run: ContextVar[Optional["Run"]] = ContextVar("mantis_current_run", default=None)
_current_rollout: ContextVar[Optional["Rollout"]] = ContextVar("mantis_current_rollout", default=None)


def get_current_run() -> Optional["Run"]:
    """Get the currently active run, if any."""
    return _current_run.get()


def _set_current_run(run: Optional["Run"]) -> None:
    """Set the current active run."""
    _current_run.set(run)


def get_current_rollout() -> Optional["Rollout"]:
    """Get the currently active rollout (episode), if any."""
    return _current_rollout.get()


def _set_current_rollout(rollout: Optional["Rollout"]) -> None:
    """Set the current active rollout."""
    _current_rollout.set(rollout)


# =============================================================================
# Rollout Class
# =============================================================================


class Rollout:
    """A Rollout/Episode within a Mantis Run.

    Represents a single episode (multi-turn interaction) within a larger
    experiment (Run). Each rollout gets its own session_id that is injected
    into traces created within its context.

    The trace_id is captured automatically by TracingContextSpanProcessor
    when the first span starts within this rollout. This allows score()
    to link scores to the correct trace even after instrumented spans end.

    Usage::

        with mantisdk.run("gepa-v1") as run:
            for candidate in population:
                with run.rollout(f"candidate-{candidate.id}") as ep:
                    result = agent.execute(task)
                    ep.score("rubric", grade)
                    # Traces get: run_id="run-xxx", session_id="ep-yyy"
    """

    def __init__(
        self,
        run: "Run",
        name: Optional[str] = None,
    ):
        """Initialize a Rollout.

        Args:
            run: The parent Run instance.
            name: Optional human-readable name for this rollout.
        """
        self._run = run
        self._name = name
        self._session_id = f"ep-{uuid.uuid4().hex[:8]}"
        self._trace_id: Optional[str] = None
        self._started = False
        self._finished = False

    @property
    def session_id(self) -> str:
        """Get the session ID for this rollout."""
        return self._session_id

    @property
    def trace_id(self) -> Optional[str]:
        """Get the trace ID captured from the first span in this rollout."""
        return self._trace_id

    @property
    def name(self) -> Optional[str]:
        """Get the rollout name."""
        return self._name

    @property
    def run(self) -> "Run":
        """Get the parent run."""
        return self._run

    def _capture_trace_id(self, trace_id: str) -> None:
        """Store the trace_id from the first span in this rollout.

        Called by TracingContextSpanProcessor.on_start() when a span
        starts within this rollout's context.
        """
        if self._trace_id is None:
            self._trace_id = trace_id

    def __enter__(self) -> "Rollout":
        """Context manager entry - set this as the current rollout."""
        self._started = True
        _set_current_rollout(self)
        logger.debug(f"Started rollout: {self._name or self._session_id}")
        return self

    def score(
        self,
        name: str,
        value: Union[float, str],
        **kwargs,
    ) -> None:
        """Post a score linked to this rollout's session and trace.

        Convenience shorthand that auto-fills session_id and trace_id
        from this rollout. See :meth:`Run.score` for full documentation.

        Args:
            name: Score name (e.g., "rubric", "accuracy").
            value: Score value.
            **kwargs: Additional arguments passed to Run.score().

        Example::

            with run.rollout("task-1") as rollout:
                result = await run_agent(task)
                rollout.score("rubric", 0.85)
        """
        # Use the trace_id captured from spans within this rollout
        if "trace_id" not in kwargs and self._trace_id:
            kwargs["trace_id"] = self._trace_id
        self._run.score(name, value, session_id=self._session_id, **kwargs)

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - clear the current rollout."""
        _set_current_rollout(None)
        self._finished = True
        if exc_type is not None:
            logger.debug(f"Rollout {self._session_id} exited with exception: {exc_val}")
        else:
            logger.debug(f"Rollout {self._session_id} completed")
        return False  # Don't suppress exceptions


# =============================================================================
# Run Class
# =============================================================================


class Run:
    """A Mantis Run for tracking metrics with automatic heartbeat and crash detection.

    Features:
    - Automatic heartbeat to detect crashed runs
    - Signal handlers for SIGINT/SIGTERM (Ctrl+C)
    - Context manager support with exception handling
    - atexit handler for scripts that exit without calling finish()
    - Automatic trace linkage via session_id
    """

    def __init__(
        self,
        name: Optional[str] = None,
        *,
        project_id: Optional[str] = None,
        heartbeat_interval: int = 30,
        source: str = "sdk",
        labels: Optional[Dict[str, Any]] = None,
        config: Optional[Dict[str, Any]] = None,
        tags: Optional[List[str]] = None,
        api_url: Optional[str] = None,
        api_key: Optional[str] = None,
        listener: Optional["StorageListener"] = None,
    ):
        """Initialize a Run.

        Args:
            name: Human-readable name for the run.
            project_id: Project ID (from env INSIGHT_PROJECT_ID if not provided).
            heartbeat_interval: Seconds between heartbeats (default: 30).
            source: Source identifier (e.g., "gepa", "trainer", "notebook").
            labels: Arbitrary key-value labels for filtering.
            config: Hyperparameters/config dict (logged for reproducibility).
            tags: Simple string tags for categorization.
            api_url: API URL (from env INSIGHT_HOST if not provided).
            api_key: API key (from env INSIGHT_SECRET_KEY if not provided).
            listener: Optional StorageListener for GEPA/optimization callbacks.

        Note:
            Runs can have MANY sessions (episodes/rollouts). Use run.rollout()
            to create individual episodes within a run.
        """
        # Validate and truncate name if too long
        if name and len(name) > MAX_RUN_NAME_LENGTH:
            logger.warning(f"Run name truncated from {len(name)} to {MAX_RUN_NAME_LENGTH} characters")
            name = name[:MAX_RUN_NAME_LENGTH]
        self._name = name or f"run-{uuid.uuid4().hex[:8]}"
        self._project_id = project_id or os.environ.get("INSIGHT_PROJECT_ID")
        self._heartbeat_interval = heartbeat_interval
        self._source = source
        self._labels = labels or {}
        self._config = config
        self._tags = tags or []

        # API configuration
        self._api_url = api_url or os.environ.get("INSIGHT_HOST", "http://localhost:3000")
        self._api_key = api_key or os.environ.get("INSIGHT_SECRET_KEY")
        # Store public key for auth (read from env since it's set by InsightLightningStore)
        self._public_key = os.environ.get("INSIGHT_PUBLIC_KEY", "")

        # Internal state
        self._run_id: Optional[str] = None
        self._run_number: int = 0
        self._state: RunState = "pending"
        self._heartbeat_thread: Optional[threading.Thread] = None
        self._stop_heartbeat = threading.Event()
        self._flush_ticker_thread: Optional[threading.Thread] = None
        self._stop_flush_ticker = threading.Event()
        self._finished = False
        self._lock = threading.Lock()
        self._score_threads: List[threading.Thread] = []

        # Signal handling state
        self._original_sigint: Optional[Callable] = None
        self._original_sigterm: Optional[Callable] = None

        # Metric buffer for batching
        self._metric_buffer: List[Dict[str, Any]] = []
        self._metric_buffer_lock = threading.Lock()
        self._last_flush_time = time.time()
        self._flush_interval = 1.0  # Flush every 1 second for live updates

        # Optional listener for GEPA/optimization callbacks
        self._listener: Optional["StorageListener"] = listener

        # Cancellation callback -- set by Trainer/algorithm to propagate
        # remote cancellation (from UI "Cancel Run") to the algorithm.
        self._on_cancel: Optional[Callable[[], None]] = None

    @property
    def id(self) -> Optional[str]:
        """Get the run ID."""
        return self._run_id

    @property
    def run_number(self) -> int:
        """Get the auto-incrementing run number for this project."""
        return self._run_number

    @property
    def name(self) -> str:
        """Get the run name."""
        return self._name

    @property
    def state(self) -> RunState:
        """Get the current run state."""
        return self._state

    @property
    def project_id(self) -> Optional[str]:
        """Get the project ID."""
        return self._project_id

    @property
    def listener(self) -> Optional["StorageListener"]:
        """Get the attached listener, if any."""
        return self._listener

    def _notify_listener(self, method_name: str, *args: Any, **kwargs: Any) -> None:
        """Notify the listener of an event (synchronous wrapper for async methods).

        This method safely calls listener methods, handling both sync and async cases.
        Exceptions are logged but not raised to avoid disrupting the main flow.
        """
        if self._listener is None:
            return
        try:
            method = getattr(self._listener, method_name, None)
            if method is not None:
                import asyncio

                if asyncio.iscoroutinefunction(method):
                    # Run async method in a new event loop if needed
                    try:
                        loop = asyncio.get_running_loop()
                        # Already in an async context, schedule it
                        asyncio.create_task(method(*args, **kwargs))
                    except RuntimeError:
                        # No running loop, run synchronously
                        asyncio.run(method(*args, **kwargs))
                else:
                    method(*args, **kwargs)
        except Exception as e:
            logger.warning(f"Listener {method_name} failed: {e}")

    # =========================================================================
    # Rollout/Episode Methods
    # =========================================================================

    def rollout(self, name: Optional[str] = None) -> Rollout:
        """Create a new rollout (episode) within this run.

        A rollout represents a single episode or interaction sequence.
        Traces created within the rollout context get both run_id and
        session_id attributes for proper grouping.

        Usage::

            with mantisdk.run("gepa-v1") as run:
                for candidate in population:
                    with run.rollout(f"candidate-{candidate.id}") as ep:
                        result = agent.execute(task)
                        # Traces get: run_id="run-xxx", session_id="ep-yyy"

        Args:
            name: Optional human-readable name for this rollout.

        Returns:
            A Rollout context manager.
        """
        return Rollout(run=self, name=name)

    # =========================================================================
    # Lifecycle Methods
    # =========================================================================

    def _auto_init_tracing(self) -> None:
        """Auto-initialize tracing if not already done and Insight env vars are set.

        This enables the "Runs standalone" experience: when a user creates a Run,
        tracing is automatically set up so all LLM calls within the run context
        are traced and exported to Insight, linked to this run via run_id.

        Only initializes if:
        - Tracing is not already initialized
        - INSIGHT_HOST or INSIGHT_PUBLIC_KEY env vars are configured
        """
        try:
            from mantisdk.tracing.init import _initialized as tracing_initialized
            from mantisdk.tracing.init import init as tracing_init

            if tracing_initialized:
                logger.debug("Tracing already initialized, skipping auto-init")
                return

            # Only auto-init if Insight env vars are present (somewhere to send spans)
            has_insight_host = bool(os.environ.get("INSIGHT_HOST"))
            has_insight_keys = bool(os.environ.get("INSIGHT_PUBLIC_KEY") and os.environ.get("INSIGHT_SECRET_KEY"))

            if has_insight_host or has_insight_keys:
                logger.info("Auto-initializing tracing from msk.run() (Insight env vars detected)")
                tracing_init(service_name=self._source or "mantisdk")
            else:
                logger.debug("Skipping tracing auto-init: no Insight env vars configured")

        except Exception as e:
            logger.warning(f"Failed to auto-init tracing: {e}")

    def _start(self) -> "Run":
        """Initialize the run on the backend and start heartbeat.

        Automatically initializes tracing if not already done and Insight
        environment variables are configured. This ensures that traces created
        within this run are exported to Insight and linked via run_id.
        """
        if not self._project_id:
            raise ValueError(
                "project_id is required. Set INSIGHT_PROJECT_ID environment variable "
                "or pass project_id to init_run()."
            )

        # Auto-init tracing if Insight env vars are configured and tracing not yet set up
        self._auto_init_tracing()

        # Create run via API (skip if run_id is pre-set, e.g., platform runs)
        if self._run_id is None:
            self._run_id, self._run_number = self._create_run_on_backend()
        else:
            logger.info(f"Reusing existing run {self._run_id} (pre-set)")
        self._state = "running"

        # Set as current run (for trace context)
        _set_current_run(self)

        # Start heartbeat thread
        self._start_heartbeat_loop()

        # Start flush ticker for live metric updates (every 1s)
        self._start_flush_ticker()

        # Register signal handlers
        self._register_signal_handlers()

        # Register atexit handler
        atexit.register(self._atexit_handler)

        # Print run URL for convenience
        run_url = f"{self._api_url}/project/{self._project_id}/runs/{self._run_id}"
        logger.info(f"Mantis Run #{self._run_number} started: {run_url}")

        return self

    def _create_run_on_backend(self) -> tuple[str, int]:
        """Create run record on backend, return (run_id, run_number)."""
        # Validate credentials before making request
        if not self._public_key:
            logger.error(
                "Cannot create run: INSIGHT_PUBLIC_KEY not set. "
                "Set it via environment variable or pass api_key to InsightLightningStore."
            )
            return f"local-{uuid.uuid4().hex}", 0

        if not self._api_key:
            logger.error(
                "Cannot create run: INSIGHT_SECRET_KEY not set. "
                "Set it via environment variable or pass secret_key to InsightLightningStore."
            )
            return f"local-{uuid.uuid4().hex}", 0

        try:
            # Ensure config is always a dict (API rejects null)
            config = self._config if self._config is not None else {}

            # Auto-capture reproducibility info
            config_hash = _compute_config_hash(config) or ""
            git_info = _get_git_info()

            url = f"{self._api_url}/api/public/runs"
            headers = self._get_headers()
            payload = {
                "name": self._name,
                "heartbeatInterval": self._heartbeat_interval,
                "source": self._source,
                "labels": self._labels or {},
                "config": config,
                "tags": self._tags or [],
                # Reproducibility tracking (API requires strings, not null)
                "configHash": config_hash,
                "gitSha": git_info.get("gitSha") or "",
                "gitBranch": git_info.get("gitBranch") or "",
                "gitDirty": git_info.get("gitDirty", False),
                # SDK metadata
                "sdkVersion": self._get_sdk_version(),
                "clientId": str(uuid.uuid4()),
                "hostname": self._get_hostname(),
            }

            logger.debug(f"Creating run at {url}")
            response = requests.post(url, json=payload, headers=headers, timeout=30)

            if not response.ok:
                # Log detailed error for debugging
                try:
                    error_body = response.json()
                except Exception:
                    error_body = response.text
                logger.error(f"Failed to create run: HTTP {response.status_code} - {error_body}")
                return f"local-{uuid.uuid4().hex}", 0

            data = response.json()
            run_id = data.get("runId")
            run_number = data.get("runNumber", 0)
            logger.info(f"Created run {run_id} (#{run_number}) on backend")
            return run_id, run_number

        except requests.exceptions.ConnectionError as e:
            logger.error(f"Connection error creating run (is {self._api_url} reachable?): {e}")
            return f"local-{uuid.uuid4().hex}", 0
        except requests.exceptions.Timeout as e:
            logger.error(f"Timeout creating run at {self._api_url}: {e}")
            return f"local-{uuid.uuid4().hex}", 0
        except Exception as e:
            logger.error(f"Unexpected error creating run: {type(e).__name__}: {e}")
            return f"local-{uuid.uuid4().hex}", 0

    def _get_headers(self) -> Dict[str, str]:
        """Get HTTP headers for API requests."""
        import base64

        headers = {"Content-Type": "application/json"}
        if self._api_key and self._public_key:
            # Use Basic auth with public_key:secret_key
            credentials = base64.b64encode(f"{self._public_key}:{self._api_key}".encode()).decode()
            headers["Authorization"] = f"Basic {credentials}"
        elif self._api_key:
            # Fallback: try reading public key from env (in case it was set after __init__)
            public_key = os.environ.get("INSIGHT_PUBLIC_KEY", "")
            if public_key:
                credentials = base64.b64encode(f"{public_key}:{self._api_key}".encode()).decode()
                headers["Authorization"] = f"Basic {credentials}"
            else:
                logger.warning("INSIGHT_PUBLIC_KEY not set - API requests will fail auth")
        return headers

    def _get_sdk_version(self) -> str:
        """Get the mantisdk version."""
        try:
            import mantisdk

            return getattr(mantisdk, "__version__", "unknown")
        except Exception:
            return "unknown"

    def _get_hostname(self) -> str:
        """Get the hostname."""
        import socket

        try:
            return socket.gethostname()
        except Exception:
            return "unknown"

    # =========================================================================
    # Heartbeat
    # =========================================================================

    def _start_heartbeat_loop(self):
        """Start background heartbeat thread."""
        self._stop_heartbeat.clear()

        def heartbeat_loop():
            while not self._stop_heartbeat.wait(timeout=self._heartbeat_interval):
                try:
                    self._send_heartbeat()
                    # Also flush metrics periodically
                    self._maybe_flush_metrics()
                except Exception as e:
                    logger.warning(f"Heartbeat failed: {e}")

        self._heartbeat_thread = threading.Thread(
            target=heartbeat_loop,
            name=f"mantis-heartbeat-{self._run_id}",
            daemon=True,
        )
        self._heartbeat_thread.start()

    def _send_heartbeat(self):
        """Send heartbeat to backend and check for remote cancellation."""
        if not self._run_id or self._run_id.startswith("local-"):
            return

        try:
            response = requests.post(
                f"{self._api_url}/api/public/runs/{self._run_id}/heartbeat",
                headers=self._get_headers(),
                timeout=10,
            )
            response.raise_for_status()

            # Check if the run was cancelled via the UI.
            try:
                data = response.json()
                if data.get("state") == "cancelled":
                    logger.info(f"Run {self._run_id} cancelled via UI")
                    self._fire_cancel_callback()
                    # Force-exit the entire process.  The UI already set
                    # state=cancelled in Postgres so no finish request needed.
                    # os._exit() skips Python cleanup but that's intentional --
                    # we need to kill forked subprocesses (shm strategy) and
                    # can't use SIGINT from the heartbeat thread (deadlock).
                    import os as _os

                    _os._exit(130)
            except (ValueError, KeyError):
                pass  # Response wasn't JSON or missing fields -- ignore
        except Exception as e:
            logger.warning(f"Heartbeat request failed: {e}")

    # =========================================================================
    # Remote cancellation
    # =========================================================================

    def set_on_cancel(self, callback: Callable[[], None]) -> None:
        """Register a callback invoked when the run is cancelled remotely.

        Typically set to ``algorithm.request_cancel`` so that a "Cancel Run"
        click in the UI propagates to the running algorithm.
        """
        self._on_cancel = callback

    def _fire_cancel_callback(self) -> None:
        """Invoke the on_cancel callback if registered."""
        cb = self._on_cancel
        if cb is not None:
            try:
                cb()
            except Exception as e:
                logger.warning(f"on_cancel callback failed: {e}")

    # =========================================================================
    # Flush Ticker (for live metric updates)
    # =========================================================================

    def _start_flush_ticker(self):
        """Start background flush ticker thread for live metric updates."""
        self._stop_flush_ticker.clear()

        def flush_ticker_loop():
            while not self._stop_flush_ticker.wait(timeout=self._flush_interval):
                try:
                    self._maybe_flush_metrics()
                except Exception as e:
                    logger.debug(f"Flush ticker failed: {e}")

        self._flush_ticker_thread = threading.Thread(
            target=flush_ticker_loop,
            name=f"mantis-flush-ticker-{self._run_id}",
            daemon=True,
        )
        self._flush_ticker_thread.start()

    # =========================================================================
    # Signal Handling
    # =========================================================================

    def _register_signal_handlers(self):
        """Register handlers for SIGINT/SIGTERM to cleanup on Ctrl+C."""

        def signal_handler(signum, frame):
            logger.info(f"Received signal {signum}, finishing run as cancelled")
            self.finish(state="cancelled")

            # Call original handler
            if signum == signal.SIGINT and callable(self._original_sigint):
                self._original_sigint(signum, frame)
            elif signum == signal.SIGTERM and callable(self._original_sigterm):
                self._original_sigterm(signum, frame)

            sys.exit(128 + signum)

        try:
            self._original_sigint = signal.signal(signal.SIGINT, signal_handler)
            self._original_sigterm = signal.signal(signal.SIGTERM, signal_handler)
        except Exception as e:
            # Signal handling may not work in all environments (e.g., threads)
            logger.debug(f"Could not register signal handlers: {e}")

    def _restore_signal_handlers(self):
        """Restore original signal handlers."""
        try:
            if self._original_sigint is not None:
                signal.signal(signal.SIGINT, self._original_sigint)
            if self._original_sigterm is not None:
                signal.signal(signal.SIGTERM, self._original_sigterm)
        except Exception as e:
            logger.debug(f"Could not restore signal handlers: {e}")

    def _atexit_handler(self):
        """Called on Python exit - mark as failed if not finished."""
        if not self._finished:
            logger.warning("Run not finished, marking as failed")
            try:
                self.finish(state="failed")
            except Exception:
                pass  # Best effort

    # =========================================================================
    # Logging Methods
    # =========================================================================

    def log(
        self,
        metrics: Union[Dict[str, Any], str],
        value: Optional[float] = None,
        *,
        step: Optional[int] = None,
        labels: Optional[Dict[str, str]] = None,
        unit: str = "",
        metric_type: str = "gauge",
    ):
        """Log metrics for this run.

        Can be called with a dict of metrics::

            run.log({"loss": 0.5, "accuracy": 0.92}, step=100)

        Or with a single metric name and value::

            run.log("loss", 0.5, step=100, unit="", metric_type="gauge")

        Args:
            metrics: Dict of {metric_name: value} or single metric name.
            value: Metric value (only when metrics is a string).
            step: Training step/iteration number.
            labels: Per-metric labels for filtering.
            unit: Unit of measurement (e.g., "ms", "tokens", "%").
            metric_type: Type hint for visualization ("gauge", "counter", "histogram").
        """
        if self._finished:
            logger.warning("Cannot log to finished run")
            return

        if isinstance(metrics, str):
            # Single metric
            if value is None:
                raise ValueError("value is required when logging a single metric")
            metric_list = [
                {
                    "name": metrics,
                    "value": float(value),
                    "step": step,
                    "labels": labels or {},
                    "unit": unit,
                    "metricType": metric_type,
                }
            ]
        else:
            # Dict of metrics
            metric_list = [
                {
                    "name": name,
                    "value": float(val),
                    "step": step,
                    "labels": labels or {},
                    "unit": unit,
                    "metricType": metric_type,
                }
                for name, val in metrics.items()
            ]

        # Buffer metrics
        with self._metric_buffer_lock:
            self._metric_buffer.extend(metric_list)

        # Flush if buffer is large enough
        if len(self._metric_buffer) >= 100:
            self._flush_metrics()

    def log_metrics(self, metrics: List[Metric]):
        """Log a batch of Metric objects.

        Args:
            metrics: List of Metric objects with full metadata.
        """
        for m in metrics:
            self.log(
                m.name,
                m.value,
                step=m.step,
                labels=m.labels,
                unit=m.unit,
                metric_type=m.metric_type,
            )

    def _maybe_flush_metrics(self):
        """Flush metrics if enough time has passed."""
        if time.time() - self._last_flush_time > self._flush_interval:
            self._flush_metrics()

    def _flush_metrics(self):
        """Send buffered metrics to backend."""
        with self._metric_buffer_lock:
            if not self._metric_buffer:
                return
            metrics_to_send = self._metric_buffer.copy()
            self._metric_buffer.clear()
            self._last_flush_time = time.time()

        if not self._run_id or self._run_id.startswith("local-"):
            logger.debug(f"Skipping metric flush (local run): {len(metrics_to_send)} metrics")
            return

        try:
            response = requests.post(
                f"{self._api_url}/api/public/runs/{self._run_id}/metrics",
                json={
                    "metrics": metrics_to_send,
                },
                headers=self._get_headers(),
                timeout=30,
            )
            response.raise_for_status()
            logger.debug(f"Flushed {len(metrics_to_send)} metrics")
        except Exception as e:
            if not getattr(self, "_metric_flush_warned", False):
                logger.warning(f"Failed to flush metrics (suppressing further warnings): {e}")
                self._metric_flush_warned = True
            # Put metrics back in buffer (best effort) with max size limit
            with self._metric_buffer_lock:
                combined = metrics_to_send + self._metric_buffer
                if len(combined) > MAX_METRIC_BUFFER_SIZE:
                    dropped = len(combined) - MAX_METRIC_BUFFER_SIZE
                    logger.warning(
                        f"Dropping {dropped} oldest metrics to prevent buffer overflow "
                        f"(max: {MAX_METRIC_BUFFER_SIZE})"
                    )
                    # Keep the most recent metrics (drop oldest)
                    combined = combined[-MAX_METRIC_BUFFER_SIZE:]
                self._metric_buffer = combined

    def flush(self) -> None:
        """Manually flush buffered metrics to the backend.

        Call this when you want immediate visibility of logged metrics
        without waiting for the automatic flush interval.

        Example:
            run.log({"accuracy": 0.95})
            run.flush()  # Metrics visible immediately in UI
        """
        self._flush_metrics()

    # =========================================================================
    # Live State (Tables-like feature for rich structured data)
    # =========================================================================

    def log_state(self, state: Dict[str, Any]) -> None:
        """Log rich structured state for live visualization.

        This is similar to W&B Tables - allows streaming structured data
        (candidates, trees, pareto fronts) that can be visualized in real-time.

        Unlike log() which is for numeric metrics, log_state() is for
        complex objects like:
        - Candidate prompts with full text
        - Evolution tree structure (parent-child relationships)
        - Pareto front membership
        - Per-instance scores

        Args:
            state: Dictionary containing structured state data.
                   Common keys for GEPA:
                   - candidates: List of candidate dicts with idx, prompt, score, parent
                   - parents: Parent indices for tree structure
                   - scores: Validation scores per candidate
                   - pareto_front: Dict mapping instance to best candidate indices
                   - best_idx: Index of best overall candidate
                   - generation: Current generation number

        Example:
            run.log_state({
                "candidates": [
                    {"idx": 0, "prompt": "You are...", "score": 0.72, "parent": None},
                    {"idx": 1, "prompt": "As an expert...", "score": 0.78, "parent": 0},
                ],
                "pareto_front": {"instance_1": [1], "instance_2": [0, 1]},
                "best_idx": 1,
                "generation": 3,
            })
        """
        if not self._run_id or self._run_id.startswith("local-"):
            logger.info(f"[log_state] Skipping (local run): {self._run_id}")
            return

        logger.info(f"[log_state] Sending to {self._api_url}/api/public/runs/{self._run_id}/state")
        try:
            response = requests.patch(
                f"{self._api_url}/api/public/runs/{self._run_id}/state",
                json={"state": state},
                headers=self._get_headers(),
                timeout=30,
            )
            response.raise_for_status()
            logger.info(f"[log_state] Success: {response.status_code}")
        except Exception as e:
            logger.warning(f"[log_state] Failed: {e}")

    # =========================================================================
    # Scores (Insight Scores linked to traces)
    # =========================================================================

    def score(
        self,
        name: str,
        value: Union[float, str],
        *,
        trace_id: Optional[str] = None,
        session_id: Optional[str] = None,
        observation_id: Optional[str] = None,
        data_type: str = "NUMERIC",
        comment: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Post a score to Insight, linked to a trace.

        Scores are evaluation results attached to traces (visible on the Scores
        page and in trace detail views). Unlike run.log() which logs run-level
        metrics, score() creates Insight Score entities linked to specific traces.

        Auto-detection:
            - If trace_id is not provided, attempts to read it from the current
              OTel span context.
            - If session_id is not provided, reads it from the active Rollout.

        Args:
            name: Score name (e.g., "rubric", "accuracy", "task_completion").
            value: Score value. Numeric for NUMERIC type, string for CATEGORICAL.
            trace_id: Trace ID to link the score to. Auto-detected from OTel
                context if not provided.
            session_id: Session ID to link the score to. Auto-detected from
                the active Rollout if not provided.
            observation_id: Optional observation/span ID for span-level scores.
            data_type: Score data type: "NUMERIC", "CATEGORICAL", or "BOOLEAN".
            comment: Optional comment explaining the score.
            metadata: Optional metadata dict.

        Example::

            with msk.run("experiment") as r:
                with r.rollout("task-1") as rollout:
                    result = await run_agent(task)
                    grade = evaluate(result)
                    # Score linked to trace + session automatically
                    r.score("rubric", grade)

                    # Or with explicit IDs
                    r.score("accuracy", 0.95, trace_id="abc123")
        """
        if self._finished:
            logger.warning("Cannot score on a finished run")
            return

        # Auto-detect trace_id from OTel context
        if trace_id is None:
            trace_id = self._get_current_trace_id()

        # Auto-detect session_id from active Rollout
        if session_id is None:
            current_rollout = get_current_rollout()
            if current_rollout is not None:
                session_id = current_rollout.session_id

        if trace_id is None:
            logger.warning(
                "Cannot post score: no trace_id provided and no active OTel span. "
                "Call score() inside a traced block or pass trace_id explicitly."
            )
            return

        # Build score payload (matches Insight POST /api/public/scores schema).
        # The API requires EXACTLY ONE of: traceId, sessionId, or datasetRunId.
        # Prefer sessionId (scores typically grade a whole session/rollout);
        # fall back to traceId when no session context exists.
        score_payload: Dict[str, Any] = {
            "name": name,
            "value": value,
            "dataType": data_type,
        }

        if session_id:
            score_payload["sessionId"] = session_id
        elif trace_id:
            score_payload["traceId"] = trace_id
            if observation_id:
                score_payload["observationId"] = observation_id
        else:
            logger.warning(f"Cannot post score '{name}': no sessionId or traceId available.")
            return
        if comment:
            score_payload["comment"] = comment

        # Merge metadata with run context
        score_metadata: Dict[str, Any] = {
            "source": "mantisdk",
            "run_id": self._run_id,
            "run_name": self._name,
        }
        if metadata:
            score_metadata.update(metadata)
        score_payload["metadata"] = score_metadata

        # Send in a background thread (non-blocking)
        self._send_score_async(score_payload)

    def _get_current_trace_id(self) -> Optional[str]:
        """Get the trace ID from the current OTel span context.

        Returns the hex trace ID if there is an active span, None otherwise.
        """
        try:
            from opentelemetry import trace as otel_trace

            current_span = otel_trace.get_current_span()
            span_context = current_span.get_span_context()
            if span_context and span_context.is_valid:
                return format(span_context.trace_id, "032x")
        except Exception as e:
            logger.debug(f"Could not get trace_id from OTel context: {e}")
        return None

    def _send_score_async(self, score_payload: Dict[str, Any]) -> None:
        """Send a score to Insight in a background thread.

        Threads are tracked and joined in finish() to ensure delivery.
        """
        if not self._run_id or self._run_id.startswith("local-"):
            logger.debug(f"Skipping score (local run): {score_payload.get('name')}")
            return

        headers = self._get_headers()
        url = f"{self._api_url}/api/public/scores"

        def _send():
            try:
                response = requests.post(
                    url,
                    json=score_payload,
                    headers=headers,
                    timeout=30,
                )
                if not response.ok:
                    try:
                        error_body = response.json()
                    except Exception:
                        error_body = response.text
                    logger.warning(
                        f"Failed to post score '{score_payload['name']}': "
                        f"HTTP {response.status_code} - {error_body}"
                    )
                    return
                logger.info(
                    f"Score posted: {score_payload['name']}={score_payload['value']} "
                    f"(trace={score_payload.get('traceId', 'N/A')[:12]}...)"
                )
            except Exception as e:
                logger.warning(f"Failed to post score: {e}")

        thread = threading.Thread(
            target=_send,
            name=f"mantis-score-{score_payload.get('name', 'unknown')}",
            daemon=True,
        )
        thread.start()
        self._score_threads.append(thread)

    # =========================================================================
    # Finish
    # =========================================================================

    def finish(self, state: RunState = "completed"):
        """Mark the run as finished.

        Args:
            state: Final state ("completed", "failed", "cancelled").
        """
        with self._lock:
            if self._finished:
                return
            self._finished = True

        # Flush remaining metrics
        self._flush_metrics()

        # Wait for pending score submissions to complete
        for t in self._score_threads:
            t.join(timeout=10)
        self._score_threads.clear()

        # Stop heartbeat
        self._stop_heartbeat.set()
        if self._heartbeat_thread:
            self._heartbeat_thread.join(timeout=5)

        # Stop flush ticker
        self._stop_flush_ticker.set()
        if self._flush_ticker_thread:
            self._flush_ticker_thread.join(timeout=2)

        # Restore signal handlers
        self._restore_signal_handlers()

        # Unregister atexit
        try:
            atexit.unregister(self._atexit_handler)
        except Exception:
            pass

        # Clear current run
        _set_current_run(None)

        # Send finish to backend
        self._state = state
        self._send_finish(state)

        logger.info(f"Mantis Run finished: {self._name} ({state})")

    def _send_finish(self, state: RunState):
        """Send finish to backend."""
        if not self._run_id or self._run_id.startswith("local-"):
            return

        try:
            response = requests.post(
                f"{self._api_url}/api/public/runs/{self._run_id}/finish",
                json={
                    "state": state,
                },
                headers=self._get_headers(),
                timeout=10,
            )
            response.raise_for_status()
        except Exception as e:
            logger.warning(f"Failed to send finish: {e}")

    # =========================================================================
    # Context Manager
    # =========================================================================

    def __enter__(self) -> "Run":
        """Context manager entry - start the run."""
        return self._start()

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - finish with appropriate state."""
        if exc_type is not None:
            # Exception occurred - mark as failed
            logger.error(f"Run failed with exception: {exc_val}")
            self.finish(state="failed")
        elif not self._finished:
            # No exception, not manually finished - complete
            self.finish(state="completed")

        return False  # Don't suppress exceptions


# =============================================================================
# Convenience Functions
# =============================================================================


def init_run(
    name: Optional[str] = None,
    *,
    project_id: Optional[str] = None,
    heartbeat_interval: int = 30,
    source: str = "sdk",
    labels: Optional[Dict[str, Any]] = None,
    config: Optional[Dict[str, Any]] = None,
    tags: Optional[List[str]] = None,
    **kwargs,
) -> Run:
    """Initialize and start a new Mantis Run.

    This is the main entry point for creating runs. For the safest usage,
    prefer the context manager form::

        with mantisdk.run("my-experiment") as run:
            run.log({"metric": value})

    Args:
        name: Human-readable name for the run.
        project_id: Project ID (from env INSIGHT_PROJECT_ID if not provided).
        heartbeat_interval: Seconds between heartbeats (default: 30).
        source: Source identifier (e.g., "gepa", "trainer").
        labels: Arbitrary key-value labels for filtering.
        config: Hyperparameters/config dict.
        tags: Simple string tags for categorization.

    Returns:
        An initialized and started Run instance.

    Note:
        Use run.rollout() to create individual episodes/rollouts within a run.
        Each rollout gets its own session_id for trace grouping.
    """
    run = Run(
        name=name,
        project_id=project_id,
        heartbeat_interval=heartbeat_interval,
        source=source,
        labels=labels,
        config=config,
        tags=tags,
        **kwargs,
    )
    return run._start()


@contextmanager
def run(
    name: Optional[str] = None,
    **kwargs,
):
    """Context manager for a Mantis Run.

    This is the recommended way to use runs as it automatically handles
    crashes, signals, and exceptions::

        with mantisdk.run("my-experiment", config={"lr": 0.01}) as r:
            for epoch in range(100):
                loss = train()
                r.log({"loss": loss}, step=epoch)

    Args:
        name: Human-readable name for the run.
        **kwargs: Additional arguments passed to Run constructor.

    Yields:
        A Run instance.
    """
    r = Run(name=name, **kwargs)
    with r:
        yield r


def log(metrics: Dict[str, Any], step: Optional[int] = None, **kwargs):
    """Log metrics to the current run.

    Convenience function for logging to the currently active run.

    Args:
        metrics: Dict of {metric_name: value}.
        step: Training step/iteration number.
        **kwargs: Additional arguments passed to Run.log().

    Raises:
        RuntimeError: If no run is currently active.
    """
    current = get_current_run()
    if current is None:
        raise RuntimeError("No active run. Call mantisdk.init_run() or use 'with mantisdk.run(...)'")
    current.log(metrics, step=step, **kwargs)


def score(
    name: str,
    value: Union[float, str],
    **kwargs,
) -> None:
    """Post a score to Insight from the current run.

    Convenience function for posting scores to the currently active run.
    Auto-detects trace_id from OTel context and session_id from the active
    Rollout.

    Args:
        name: Score name (e.g., "rubric", "accuracy").
        value: Score value.
        **kwargs: Additional arguments passed to Run.score().

    Raises:
        RuntimeError: If no run is currently active.
    """
    current = get_current_run()
    if current is None:
        raise RuntimeError("No active run. Call mantisdk.init_run() or use 'with mantisdk.run(...)'")
    current.score(name, value, **kwargs)


def finish(state: RunState = "completed"):
    """Finish the current run.

    Convenience function for finishing the currently active run.

    Args:
        state: Final state ("completed", "failed", "cancelled").
    """
    current = get_current_run()
    if current is None:
        return
    current.finish(state=state)


# =============================================================================
# Decorator-based API
# =============================================================================


def experiment(
    name: Optional[str] = None,
    *,
    inject_run: bool = True,
    **run_kwargs,
) -> Callable[[F], F]:
    """Decorator to wrap a function in a Mantis Run.

    This provides an alternative to context managers for experiment tracking.
    The run is automatically started before the function executes and finished
    after it completes (or fails).

    Example::

        @mantisdk.experiment("gepa-v1", config={"lr": 0.01})
        def train(run):
            run.log({"started": True})
            for candidate in population:
                evaluate(candidate)

        @mantisdk.experiment(inject_run=False)
        def simple_experiment():
            # Access run via get_current_run()
            mantisdk.log({"metric": value})

    Args:
        name: Human-readable name for the run. Defaults to function name.
        inject_run: If True (default), passes the Run instance as the first
            argument to the decorated function. If False, access via
            get_current_run().
        **run_kwargs: Additional arguments passed to Run constructor
            (config, labels, tags, source, etc).

    Returns:
        A decorator function.
    """

    def decorator(fn: F) -> F:
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            run_name = name or fn.__name__
            r = Run(name=run_name, **run_kwargs)
            with r:
                if inject_run:
                    return fn(r, *args, **kwargs)
                return fn(*args, **kwargs)

        @functools.wraps(fn)
        async def async_wrapper(*args, **kwargs):
            run_name = name or fn.__name__
            r = Run(name=run_name, **run_kwargs)
            with r:
                if inject_run:
                    return await fn(r, *args, **kwargs)
                return await fn(*args, **kwargs)

        if asyncio.iscoroutinefunction(fn):
            return async_wrapper  # type: ignore
        return wrapper  # type: ignore

    return decorator


def episode(
    name: Optional[str] = None,
    *,
    name_fn: Optional[Callable[..., str]] = None,
    inject_rollout: bool = False,
) -> Callable[[F], F]:
    """Decorator to wrap a function in a Rollout (episode).

    This provides an alternative to context managers for episode tracking.
    The rollout is automatically started before the function executes and
    ended after it completes.

    Requires an active run (created via @experiment decorator or context manager).

    Example::

        @mantisdk.experiment("gepa", config={"lr": 0.01})
        def train(run):
            for candidate in population:
                evaluate(candidate)

        @mantisdk.episode()
        def evaluate(candidate):
            result = agent.execute(task)
            # Traces automatically get run_id + session_id

        # With dynamic name based on arguments
        @mantisdk.episode(name_fn=lambda c: f"candidate-{c.id}")
        def evaluate_named(candidate):
            result = agent.execute(task)

    Args:
        name: Static name for the rollout. Defaults to function name.
        name_fn: Optional function that takes the same arguments as the
            decorated function and returns a rollout name.
        inject_rollout: If True, passes the Rollout instance as the first
            argument to the decorated function. Default False.

    Returns:
        A decorator function.

    Raises:
        RuntimeError: If no active run exists when the function is called.
    """

    def decorator(fn: F) -> F:
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            current_run = get_current_run()
            if current_run is None:
                raise RuntimeError(
                    "@episode requires an active run. " "Use @experiment decorator or 'with mantisdk.run(...)' first."
                )

            # Determine rollout name
            if name_fn is not None:
                rollout_name = name_fn(*args, **kwargs)
            else:
                rollout_name = name or fn.__name__

            rollout = current_run.rollout(rollout_name)
            with rollout:
                if inject_rollout:
                    return fn(rollout, *args, **kwargs)
                return fn(*args, **kwargs)

        @functools.wraps(fn)
        async def async_wrapper(*args, **kwargs):
            current_run = get_current_run()
            if current_run is None:
                raise RuntimeError(
                    "@episode requires an active run. " "Use @experiment decorator or 'with mantisdk.run(...)' first."
                )

            # Determine rollout name
            if name_fn is not None:
                rollout_name = name_fn(*args, **kwargs)
            else:
                rollout_name = name or fn.__name__

            rollout = current_run.rollout(rollout_name)
            with rollout:
                if inject_rollout:
                    return await fn(rollout, *args, **kwargs)
                return await fn(*args, **kwargs)

        if asyncio.iscoroutinefunction(fn):
            return async_wrapper  # type: ignore
        return wrapper  # type: ignore

    return decorator
