# Copyright (c) Microsoft. All rights reserved.

"""InsightTracker and InsightLightningStore - Insight tracking integration.

This module provides Insight tracking integration:

1. **InsightLightningStore** (recommended): Creates a real MantisRun for experiment
   tracking with automatic session management, heartbeats, and crash detection.
   ```python
   from mantisdk.store import InsightLightningStore

   store = InsightLightningStore(
       api_key="pk-lf-abc123",
       secret_key="sk-lf-xyz789",
       insight_url="https://insight.withmetis.ai",
       project_id="proj-123",
       run_name="gepa-experiment-v1",
       config={"lr": 0.01, "epochs": 100},
   )
   # Store creates a real MantisRun record in the database
   # Rollouts create real TraceSession records
   # All traces are linked via run_id
   ```

2. **InsightTracker** (DEPRECATED): Legacy event-based tracking with virtual job IDs.
   Use InsightLightningStore instead, which uses real MantisRun records.

Features:
- Real MantisRun records with lifecycle management (heartbeats, crash detection)
- Session creation for rollouts (TraceSession records linked to runs)
- OTLP trace export support
- Automatic reward-to-score conversion
"""

from __future__ import annotations

import atexit
import base64
import logging
import queue
import threading
import time
import uuid
import warnings
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from mantisdk.run import Run
    from .insight_listener import InsightRunListener
    from .memory import InMemoryLightningStore  # noqa: F401
    from .client_server import LightningStoreClient  # noqa: F401
    from .redis_stream import RedisStreamLightningStore  # noqa: F401

import httpx

from mantisdk.emitter.reward import get_rewards_from_span, is_reward_span
from mantisdk.types import (
    Attempt,
    ResourcesUpdate,
    Rollout,
    Span,
)

from .listener import StorageListener

logger = logging.getLogger(__name__)


@dataclass
class InsightEvent:
    """An event to be sent to the Insight API."""

    id: str
    type: str
    timestamp: str  # ISO 8601 datetime string
    data: Dict[str, Any]


class InsightTracker:
    """DEPRECATED: A StorageListener that streams storage events to Insight.

    .. deprecated::
        InsightTracker is deprecated. Use InsightLightningStore instead, which
        creates real MantisRun records with proper lifecycle management, heartbeats,
        and crash detection.

    This tracker implements the StorageListener protocol and can be attached
    to any LightningStore to enable Insight tracking. However, it uses virtual
    job IDs that are not persisted to the database.

    For new code, use InsightLightningStore which:
    - Creates real MantisRun records in PostgreSQL
    - Has automatic heartbeats and crash detection
    - Creates real TraceSession records for rollouts
    - Links all traces via run_id

    Args:
        api_key: Insight public API key for authentication.
        secret_key: Insight secret key for authentication.
        insight_url: Insight server URL (e.g., "http://localhost:3000").
        project_id: Project ID to associate events with.
        flush_interval: Seconds between automatic flushes (default: 1.0).
        max_buffer_size: Maximum events before forcing a flush (default: 1000).
        request_timeout: HTTP request timeout in seconds (default: 10.0).
        max_retries: Maximum retry attempts for failed requests (default: 3).
    """

    def __init__(
        self,
        *,
        api_key: str,
        secret_key: str,
        insight_url: str,
        project_id: str,
        flush_interval: float = 1.0,
        max_buffer_size: int = 1000,
        request_timeout: float = 10.0,
        max_retries: int = 3,
    ) -> None:
        warnings.warn(
            "InsightTracker is deprecated. Use InsightLightningStore instead, which "
            "creates real MantisRun records with proper lifecycle management.",
            DeprecationWarning,
            stacklevel=2,
        )
        # Store configuration
        self._api_key = api_key
        self._secret_key = secret_key
        self._insight_url = insight_url.rstrip("/")
        self._project_id = project_id
        self._flush_interval = flush_interval
        self._max_buffer_size = max_buffer_size
        self._request_timeout = request_timeout
        self._max_retries = max_retries

        # Generate a unique job ID for this tracker instance
        self._job_id = f"job-{uuid.uuid4().hex[:12]}"

        # Track if job has been completed to prevent double-complete
        self._completed = False

        # Event buffer (thread-safe queue)
        self._event_buffer: queue.Queue[InsightEvent] = queue.Queue()

        # Background sender thread control
        self._stop_event = threading.Event()
        self._sender_thread: Optional[threading.Thread] = None

        # Start the background sender thread
        self._start_sender_thread()

        # Emit job.created event immediately
        self._emit_job_created()

        # Register cleanup on exit
        atexit.register(self._cleanup)

        logger.info(
            f"InsightTracker initialized - streaming to {self._insight_url} "
            f"(project={self._project_id}, job={self._job_id})"
        )

    # ─────────────────────────────────────────────────────────────
    # StorageListener Protocol Implementation
    # ─────────────────────────────────────────────────────────────

    @property
    def capabilities(self) -> Dict[str, bool]:
        """Return the capabilities of the listener."""
        return {
            "otlp_traces": True,  # Enable OTLP trace export to Insight
        }

    @property
    def job_id(self) -> str:
        """Return the job ID for this tracker instance."""
        return self._job_id

    def otlp_traces_endpoint(self) -> Optional[str]:
        """Return the OTLP/HTTP traces endpoint."""
        endpoint = f"{self._insight_url}/api/public/otel/v1/traces"
        logger.debug(f"OTLP traces endpoint: {endpoint}")
        return endpoint

    def get_otlp_headers(self) -> Dict[str, str]:
        """Return the authentication headers for OTLP export.

        Insight's OTLP endpoint uses Basic Auth with format: public_key:secret_key
        """
        credentials = f"{self._api_key}:{self._secret_key}"
        encoded = base64.b64encode(credentials.encode()).decode()
        return {
            "Authorization": f"Basic {encoded}",
        }

    async def on_job_created(self, job_id: str, project_id: Optional[str] = None) -> None:
        """Called when a job is created. (Usually handled internally)"""
        pass  # We emit job.created in __init__

    async def on_rollout_created(self, rollout: Rollout) -> None:
        """Called when a rollout is created."""
        self._emit(
            "rollout.created",
            {
                "id": rollout.rollout_id,
                "input": rollout.input,
                "status": rollout.status,
                "resource_id": rollout.resources_id,
                "mode": rollout.mode,
                "start_time": rollout.start_time,
                "config": {
                    "max_attempts": rollout.config.max_attempts,
                    "retry_condition": rollout.config.retry_condition,
                    "timeout_seconds": rollout.config.timeout_seconds,
                    "unresponsive_seconds": rollout.config.unresponsive_seconds,
                },
                "metadata": rollout.metadata,
            },
        )

    async def on_rollout_updated(self, rollout: Rollout) -> None:
        """Called when a rollout is updated."""
        data: Dict[str, Any] = {
            "id": rollout.rollout_id,
            "status": rollout.status,
        }
        if rollout.end_time is not None:
            data["end_time"] = rollout.end_time
        self._emit("rollout.status_changed", data)

    async def on_attempt_created(self, attempt: Attempt) -> None:
        """Called when an attempt is created."""
        self._emit(
            "attempt.created",
            {
                "id": attempt.attempt_id,
                "rollout_id": attempt.rollout_id,
                "sequence_id": attempt.sequence_id,
                "status": attempt.status,
                "start_time": attempt.start_time,
                "worker_id": attempt.worker_id,
            },
        )

    async def on_attempt_updated(self, attempt: Attempt, rollout_id: str) -> None:
        """Called when an attempt is updated."""
        data: Dict[str, Any] = {
            "id": attempt.attempt_id,
            "rollout_id": rollout_id,
            "status": attempt.status,
        }
        if attempt.end_time is not None:
            data["end_time"] = attempt.end_time
        self._emit("attempt.status_changed", data)

    async def on_span_created(self, span: Span) -> None:
        """Called when a span is added.

        If the span is a reward span, also sends the reward as an Insight score.
        """
        # Emit span event to the event buffer
        self._emit(
            "span.emitted",
            {
                "trace_id": span.trace_id,
                "span_id": span.span_id,
                "parent_id": span.parent_id,
                "attempt_id": span.attempt_id,
                "rollout_id": span.rollout_id,
                "sequence_id": span.sequence_id,
                "name": span.name,
                "status": {
                    "status_code": span.status.status_code,
                    "description": span.status.description,
                },
                "attributes": span.attributes,
                "start_time": span.start_time,
                "end_time": span.end_time,
            },
        )

        # Check if this is a reward span and send as Insight score
        if is_reward_span(span):
            rewards = get_rewards_from_span(span)
            for reward in rewards:
                self._send_score(
                    name=reward.name,
                    value=reward.value,
                    trace_id=span.trace_id,
                    observation_id=span.span_id,
                    rollout_id=span.rollout_id,
                    attempt_id=span.attempt_id,
                )

    async def on_resource_registered(self, resource: ResourcesUpdate) -> None:
        """Called when a resource is registered/updated.

        Sends FULL resource content for complete experiment tracking.
        """
        # Serialize resources to JSON-compatible dicts
        # Resources are Pydantic models (PromptTemplate, LLM, etc.) which need model_dump()
        serialized_resources: Dict[str, Any] = {}
        for name, res in resource.resources.items():
            if hasattr(res, "model_dump"):
                serialized_resources[name] = res.model_dump()
            elif hasattr(res, "dict"):
                # Fallback for older Pydantic v1 models
                serialized_resources[name] = res.dict()
            else:
                # Already a dict or primitive
                serialized_resources[name] = res

        self._emit(
            "resource.registered",
            {
                "id": resource.resources_id,
                "version": resource.version,
                "create_time": resource.create_time,
                "update_time": resource.update_time,
                "resources": serialized_resources,
            },
        )

    # ─────────────────────────────────────────────────────────────
    # Event Emission Helpers
    # ─────────────────────────────────────────────────────────────

    def _emit(self, event_type: str, data: Dict[str, Any]) -> None:
        """Add an event to the buffer for later sending."""
        # Format: 2026-01-17T19:51:40.123Z (Zod datetime expects 'Z' suffix for UTC)
        now = datetime.now(timezone.utc)
        timestamp = now.strftime("%Y-%m-%dT%H:%M:%S.") + f"{now.microsecond // 1000:03d}Z"
        event = InsightEvent(
            id=f"evt-{uuid.uuid4().hex[:8]}",
            type=event_type,
            timestamp=timestamp,
            data=data,
        )
        try:
            self._event_buffer.put_nowait(event)
        except queue.Full:
            logger.warning(f"Event buffer full, dropping event: {event_type}")
            return

        # Force flush if buffer is at capacity
        if self._event_buffer.qsize() >= self._max_buffer_size:
            self._trigger_flush()

    def _emit_job_created(self) -> None:
        """Emit the job.created event when the tracker is initialized."""
        self._emit(
            "job.created",
            {
                "project_id": self._project_id,
                "type": "agent",
            },
        )
        # Immediately flush to ensure job is created before other events
        self._flush_events()

    def _send_score(
        self,
        name: str,
        value: float,
        trace_id: str,
        observation_id: Optional[str] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
    ) -> None:
        """Send a score to Insight's /api/public/scores endpoint.

        Args:
            name: Score name (e.g., "primary", "task_completion").
            value: Numeric score value.
            trace_id: OTEL trace ID to link the score to.
            observation_id: Optional span ID to link to specific observation.
            rollout_id: Optional rollout ID for metadata.
            attempt_id: Optional attempt ID for metadata.
        """
        score_payload: Dict[str, Any] = {
            "name": name,
            "value": value,
            "dataType": "NUMERIC",
            "traceId": trace_id,
        }

        if observation_id:
            score_payload["observationId"] = observation_id

        # Add rollout/attempt context as metadata
        metadata: Dict[str, Any] = {
            "source": "agent_lightning",
            "job_id": self._job_id,
        }
        if rollout_id:
            metadata["rollout_id"] = rollout_id
        if attempt_id:
            metadata["attempt_id"] = attempt_id
        score_payload["metadata"] = metadata

        # Create Basic auth header
        auth_string = f"{self._api_key}:{self._secret_key}"
        auth_bytes = base64.b64encode(auth_string.encode("utf-8")).decode("utf-8")

        # Send score in a non-blocking way (fire and forget with retries)
        def send_score_async() -> None:
            for attempt in range(self._max_retries):
                try:
                    with httpx.Client(timeout=self._request_timeout) as client:
                        response = client.post(
                            f"{self._insight_url}/api/public/scores",
                            headers={
                                "Authorization": f"Basic {auth_bytes}",
                                "Content-Type": "application/json",
                            },
                            json=score_payload,
                        )
                        response.raise_for_status()
                        logger.debug(f"Successfully sent score '{name}={value}' to Insight")
                        return
                except httpx.HTTPStatusError as e:
                    if e.response.status_code == 401:
                        logger.error("Unauthorized (401) sending score to Insight. Check API credentials.")
                        return
                    else:
                        try:
                            error_body = e.response.json()
                        except Exception:
                            error_body = e.response.text
                        logger.warning(
                            f"HTTP error sending score (attempt {attempt + 1}/{self._max_retries}): "
                            f"{e.response.status_code} - {error_body}"
                        )
                except httpx.HTTPError as e:
                    logger.warning(f"Failed to send score (attempt {attempt + 1}/{self._max_retries}): {e}")

                # Exponential backoff
                if attempt < self._max_retries - 1:
                    backoff_time = 2**attempt
                    time.sleep(backoff_time)

            logger.error(f"Failed to send score '{name}' after {self._max_retries} retries")

        # Run in a thread to avoid blocking
        threading.Thread(target=send_score_async, daemon=True, name="insight-score-sender").start()

    # ─────────────────────────────────────────────────────────────
    # Background Sender Thread
    # ─────────────────────────────────────────────────────────────

    def _start_sender_thread(self) -> None:
        """Start the background sender thread."""

        def sender_loop() -> None:
            while not self._stop_event.is_set():
                # Wait for the flush interval or until stopped
                self._stop_event.wait(timeout=self._flush_interval)
                if not self._stop_event.is_set():
                    self._flush_events()

        self._sender_thread = threading.Thread(
            target=sender_loop,
            daemon=True,
            name="insight-sender",
        )
        self._sender_thread.start()

    def _trigger_flush(self) -> None:
        """Trigger an immediate flush by interrupting the wait."""
        # The flush will happen on the next iteration since we're using a timeout
        pass

    def _flush_events(self) -> None:
        """Flush all buffered events to the Insight API."""
        events: List[InsightEvent] = []

        # Drain the queue
        while True:
            try:
                event = self._event_buffer.get_nowait()
                events.append(event)
            except queue.Empty:
                break

        if not events:
            return

        self._send_events(events)

    def _send_events(self, events: List[InsightEvent]) -> None:
        """Send events to the Insight API with retry logic."""
        if not events:
            return

        payload = {
            "job_id": self._job_id,
            # Note: project_id is not sent - it's derived from the API key auth
            "events": [
                {
                    "id": e.id,
                    "type": e.type,
                    "timestamp": e.timestamp,
                    "data": e.data,
                }
                for e in events
            ],
        }

        # Create Basic auth header from api_key:secret_key
        auth_string = f"{self._api_key}:{self._secret_key}"
        auth_bytes = base64.b64encode(auth_string.encode("utf-8")).decode("utf-8")

        for attempt in range(self._max_retries):
            try:
                with httpx.Client(timeout=self._request_timeout) as client:
                    response = client.post(
                        f"{self._insight_url}/api/public/v1/agent/ingest",
                        headers={
                            "Authorization": f"Basic {auth_bytes}",
                            "Content-Type": "application/json",
                        },
                        json=payload,
                    )
                    response.raise_for_status()
                    logger.debug(f"Successfully sent {len(events)} events to Insight")
                    return
            except httpx.HTTPStatusError as e:
                if e.response.status_code == 401:
                    # Unauthorized - bad API key, won't fix itself
                    logger.error("Unauthorized (401) sending events to Insight. Check API credentials.")
                    return
                elif e.response.status_code == 429:
                    # Rate limited - back off
                    logger.warning("Rate limited (429) sending events, retrying with backoff...")
                else:
                    # Log the full response for debugging
                    try:
                        error_body = e.response.json()
                    except Exception:
                        error_body = e.response.text
                    logger.warning(
                        f"HTTP error sending events (attempt {attempt + 1}/{self._max_retries}): "
                        f"{e.response.status_code} - {error_body}"
                    )
            except httpx.HTTPError as e:
                logger.warning(f"Failed to send events (attempt {attempt + 1}/{self._max_retries}): {e}")

            # Exponential backoff
            if attempt < self._max_retries - 1:
                backoff_time = 2**attempt
                time.sleep(backoff_time)

        logger.error(f"Failed to send {len(events)} events after {self._max_retries} retries - events dropped")

    # ─────────────────────────────────────────────────────────────
    # Lifecycle Methods
    # ─────────────────────────────────────────────────────────────

    def complete(self, summary: Optional[Dict[str, Any]] = None) -> None:
        """Mark the job as complete and flush all remaining events."""
        if self._completed:
            return
        self._completed = True
        self._emit("job.completed", {"summary": summary or {}})
        self._stop_event.set()
        self._flush_events()
        logger.info(f"InsightTracker job {self._job_id} completed")

    def fail(self, error: str) -> None:
        """Mark the job as failed and flush all remaining events."""
        self._emit("job.failed", {"error": error})
        self._stop_event.set()
        self._flush_events()
        logger.error(f"InsightTracker job {self._job_id} failed: {error}")

    def _cleanup(self) -> None:
        """Cleanup resources on exit."""
        if not self._stop_event.is_set():
            # Flush any remaining events before exiting
            self._flush_events()
            self._stop_event.set()

    def __enter__(self) -> "InsightTracker":
        return self

    def __exit__(
        self,
        exc_type: Optional[type],
        exc_val: Optional[BaseException],
        exc_tb: Optional[Any],
    ) -> None:
        if exc_type:
            self.fail(str(exc_val) if exc_val else "Unknown error")
        elif not self._completed:
            self.complete()


# ─────────────────────────────────────────────────────────────────────────────
# InsightLightningStore - Uses InsightRunListener for consolidated tracking
# ─────────────────────────────────────────────────────────────────────────────


def InsightLightningStore(
    *,
    api_key: str,
    secret_key: str,
    insight_url: str,
    project_id: str,
    # Run configuration
    run_name: Optional[str] = None,
    run_id: Optional[str] = None,
    config: Optional[Dict[str, Any]] = None,
    labels: Optional[Dict[str, Any]] = None,
    tags: Optional[List[str]] = None,
    source: str = "insight_store",
    heartbeat_interval: int = 30,
    # Store configuration
    thread_safe: bool = False,
    init_tracing: bool = False,
    # Legacy parameters (passed to listener)
    flush_interval: float = 1.0,
    max_buffer_size: int = 1000,
    request_timeout: float = 10.0,
    max_retries: int = 3,
    **kwargs: Any,
) -> "InMemoryLightningStore":
    """Create an InMemoryLightningStore with Insight tracking.

    This function creates a real MantisRun record in the database for experiment
    tracking. All Insight communication is consolidated in InsightRunListener:
    - Run creation and lifecycle management
    - Heartbeats for crash detection
    - Metrics logging
    - Session creation for rollouts
    - Score submission

    Args:
        api_key: Insight public API key for authentication.
        secret_key: Insight secret key for authentication.
        insight_url: Insight server URL (e.g., "http://localhost:3000").
        project_id: Project ID to associate the run with.
        run_name: Optional name for the run (auto-generated if not provided).
        config: Hyperparameters/config dict (logged for reproducibility).
        labels: Arbitrary key-value labels for filtering runs.
        tags: Simple string tags for categorization.
        source: Source identifier (default: "insight_store").
        heartbeat_interval: Seconds between heartbeats (default: 30).
        thread_safe: Whether the underlying store is thread-safe (default: False).
        init_tracing: If True, automatically initialize mantisdk.tracing with:
            - Insight OTLP exporter configured from the provided credentials
            - TracingContextSpanProcessor for automatic context injection
            - OpenInference instrumentors (litellm, openai, anthropic, etc.)

    Returns:
        An InMemoryLightningStore with InsightRunListener attached.

    Example:
        ```python
        # Basic usage
        store = InsightLightningStore(
            api_key="pk-lf-abc123",
            secret_key="sk-lf-xyz789",
            insight_url="https://insight.withmetis.ai",
            project_id="proj-123",
            run_name="gepa-experiment-v1",
            config={"lr": 0.01, "epochs": 100},
        )

        # Log metrics directly on the store
        store.log({"loss": 0.5, "accuracy": 0.92}, step=100)

        # With automatic tracing SDK setup (recommended)
        store = InsightLightningStore(
            api_key="pk-lf-abc123",
            secret_key="sk-lf-xyz789",
            insight_url="https://insight.withmetis.ai",
            project_id="proj-123",
            init_tracing=True,
        )

        trainer = Trainer(algorithm=GEPA(...), store=store)
        # Run is tracked with heartbeats, metrics, sessions, etc.

        # When done, call complete() or fail()
        store.complete()
        ```
    """
    import os

    from .insight_listener import InsightRunListener
    from .memory import InMemoryLightningStore as MemStore

    # Set env vars for tracing SDK and other components
    os.environ["INSIGHT_HOST"] = insight_url
    os.environ["INSIGHT_PUBLIC_KEY"] = api_key
    os.environ["INSIGHT_SECRET_KEY"] = secret_key
    os.environ["INSIGHT_PROJECT_ID"] = project_id

    logger.debug(
        f"InsightLightningStore: Set env vars - "
        f"INSIGHT_HOST={insight_url}, "
        f"INSIGHT_PUBLIC_KEY={api_key[:20]}..., "
        f"INSIGHT_PROJECT_ID={project_id}"
    )

    # Initialize tracing SDK if requested
    if init_tracing:
        import mantisdk.tracing as tracing

        tracing.init(instrument_default=True)

    # Create consolidated listener - handles ALL Insight communication
    listener = InsightRunListener(
        api_key=api_key,
        secret_key=secret_key,
        insight_url=insight_url,
        project_id=project_id,
        run_name=run_name,
        run_id=run_id,  # Reuse existing run if provided (e.g., platform runs)
        config=config,
        labels=labels,
        tags=tags,
        source=source,
        heartbeat_interval=heartbeat_interval,
        request_timeout=request_timeout,
        max_retries=max_retries,
    )

    # Start the run (creates DB record, starts heartbeat, registers handlers)
    listener.start()

    # Create store with listener attached
    store = MemStore(thread_safe=thread_safe, listeners=[listener], **kwargs)

    # Attach listener to store for access
    store._insight_listener = listener  # type: ignore[attr-defined]

    # Expose clean API methods on store instance
    def log(
        metrics,
        value=None,
        *,
        step=None,
        labels=None,
        unit="",
        metric_type="gauge",
    ):
        """Log metrics for this run."""
        listener.log(metrics, value, step=step, labels=labels, unit=unit, metric_type=metric_type)

    def complete(summary: Optional[Dict[str, Any]] = None) -> None:
        """Mark the run as completed."""
        listener.complete()

    def fail(error: Optional[str] = None) -> None:
        """Mark the run as failed."""
        listener.fail(error)

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

    def flush() -> None:
        """Manually flush buffered metrics."""
        listener.flush()

    # Monkey-patch the API methods onto the store instance
    store.log = log  # type: ignore[attr-defined]
    store.complete = complete  # type: ignore[attr-defined]
    store.fail = fail  # type: ignore[attr-defined]
    store.log_state = log_state  # type: ignore[attr-defined]
    store.flush = flush  # type: ignore[attr-defined]

    # Add run_id property for convenience
    store.run_id = listener.run_id  # type: ignore[attr-defined]
    store.job_id = listener.run_id  # type: ignore[attr-defined]  # Backward compat

    # Deprecated: _run attribute for backward compatibility
    # Creates a simple object that forwards to the listener
    class _RunCompat:
        """Backward compatibility wrapper for store._run access."""

        def __init__(self, listener):
            self._listener = listener

        @property
        def id(self):
            return self._listener.run_id

        def log(self, *args, **kwargs):
            return self._listener.log(*args, **kwargs)

        def log_state(self, state):
            return self._listener.log_state(state)

        def flush(self):
            return self._listener.flush()

        def finish(self, state="completed"):
            return self._listener.finish(state)

    store._run = _RunCompat(listener)  # type: ignore[attr-defined]

    logger.info(
        f"InsightLightningStore created with MantisRun {listener.run_id} "
        f"(project={project_id}, name={listener.run_name})"
    )

    return store


# ─────────────────────────────────────────────────────────────────────────────
# RedisStreamInsightLightningStore - For platform runs with distributed coord
# ─────────────────────────────────────────────────────────────────────────────


def RedisStreamInsightLightningStore(
    *,
    redis_url: str,
    run_id: str,
    api_key: str,
    secret_key: str,
    insight_url: str,
    project_id: str,
    run_name: Optional[str] = None,
    config: Optional[Dict[str, Any]] = None,
    labels: Optional[Dict[str, Any]] = None,
    tags: Optional[List[str]] = None,
    source: str = "platform",
    heartbeat_interval: int = 30,
    consumer_group: str = "mantis-runners",
    thread_safe: bool = False,
    init_tracing: bool = False,
) -> "RedisStreamLightningStore":
    """Create a RedisStreamLightningStore with InsightRunListener attached.

    Used by the Python BullMQ worker for platform-managed runs. Combines
    Redis Streams for distributed rollout coordination with InsightRunListener
    for run lifecycle management (heartbeats, metrics, sessions).

    Args:
        redis_url: Redis connection string.
        run_id: Existing MantisRun ID (created by the platform start endpoint).
        api_key: Insight public API key.
        secret_key: Insight secret key.
        insight_url: Insight server URL.
        project_id: Project ID.
        run_name: Optional run name.
        config: Run config dict.
        labels: Run labels.
        tags: Run tags.
        source: Run source (default: "platform").
        heartbeat_interval: Heartbeat interval in seconds.
        consumer_group: Redis consumer group name.
        thread_safe: Whether the store is thread-safe.
        init_tracing: If True, auto-initialize mantisdk.tracing.

    Returns:
        RedisStreamLightningStore with InsightRunListener attached.
    """
    import os

    from .insight_listener import InsightRunListener
    from .redis_stream import RedisStreamLightningStore as RedisStore

    # Set non-secret env vars for SDK components that read from env
    os.environ["INSIGHT_HOST"] = insight_url
    os.environ["INSIGHT_PUBLIC_KEY"] = api_key
    os.environ["INSIGHT_PROJECT_ID"] = project_id

    if init_tracing:
        # Temporarily set secret key for tracing init, then remove (F6)
        os.environ["INSIGHT_SECRET_KEY"] = secret_key
        import mantisdk.tracing as tracing

        tracing.init(instrument_default=True)
        os.environ.pop("INSIGHT_SECRET_KEY", None)

    listener = InsightRunListener(
        api_key=api_key,
        secret_key=secret_key,
        insight_url=insight_url,
        project_id=project_id,
        run_name=run_name,
        config=config,
        labels=labels,
        tags=tags,
        source=source,
        heartbeat_interval=heartbeat_interval,
    )
    listener.start()

    store = RedisStore(
        redis_url=redis_url,
        run_id=run_id,
        consumer_group=consumer_group,
        thread_safe=thread_safe,
        listeners=[listener],
    )

    store.log = listener.log  # type: ignore[attr-defined]
    store.complete = listener.complete  # type: ignore[attr-defined]
    store.fail = listener.fail  # type: ignore[attr-defined]
    store.log_state = listener.log_state  # type: ignore[attr-defined]
    store.flush = listener.flush  # type: ignore[attr-defined]

    from mantisdk.utils.redact import redact_url

    logger.info(
        f"RedisStreamInsightLightningStore created for run {run_id} "
        f"(redis={redact_url(redis_url)}, project={project_id})"
    )

    return store


def create_platform_store_client_from_env() -> "LightningStoreClient":
    """Create a store client for agent containers running on the Mantis platform.

    Only requires ``MANTIS_STORE_URL`` (set by the orchestrator). All Insight
    integration (OTLP export, sessions, scores, heartbeats) is handled
    server-side by the worker's InsightRunListener.

    Returns:
        A :class:`LightningStoreClient` connected to the worker's store server.

    Raises:
        EnvironmentError: If MANTIS_STORE_URL is not set.
    """
    import os

    from .client_server import LightningStoreClient

    store_url = os.environ.get("MANTIS_STORE_URL")
    if not store_url:
        raise EnvironmentError(
            "MANTIS_STORE_URL is not set. This variable is injected automatically "
            "by the Mantis platform orchestrator when starting agent containers."
        )

    logger.info(f"Creating store client for {store_url}")
    return LightningStoreClient(server_address=store_url)


def create_platform_store_from_env() -> "RedisStreamLightningStore":
    """DEPRECATED: Use create_platform_store_client_from_env() instead.

    This function creates a RedisStreamLightningStore which has been
    replaced by LightningStoreClient for platform runs.
    """
    import os
    import warnings

    warnings.warn(
        "create_platform_store_from_env() is deprecated. "
        "Use create_platform_store_client_from_env() instead, which uses "
        "LightningStoreClient (HTTP) rather than Redis Streams.",
        DeprecationWarning,
        stacklevel=2,
    )

    required = ["MANTIS_RUN_ID", "INSIGHT_HOST", "INSIGHT_PUBLIC_KEY", "INSIGHT_SECRET_KEY", "INSIGHT_PROJECT_ID"]
    missing = [k for k in required if not os.environ.get(k)]
    if missing:
        raise EnvironmentError(
            f"Missing required environment variables: {', '.join(missing)}. "
            "These are set automatically by the Mantis platform orchestrator."
        )

    run_id = os.environ["MANTIS_RUN_ID"]
    host = os.environ.get("REDIS_HOST", "localhost")
    port = os.environ.get("REDIS_PORT", "6379")
    auth = os.environ.get("REDIS_AUTH", "")
    redis_url = f"redis://:{auth}@{host}:{port}" if auth else f"redis://{host}:{port}"

    return RedisStreamInsightLightningStore(
        redis_url=redis_url,
        run_id=run_id,
        api_key=os.environ["INSIGHT_PUBLIC_KEY"],
        secret_key=os.environ["INSIGHT_SECRET_KEY"],
        insight_url=os.environ["INSIGHT_HOST"],
        project_id=os.environ["INSIGHT_PROJECT_ID"],
    )
