# Copyright (c) Metis. All rights reserved.

"""InsightRunListener - Consolidated Insight backend communication via Run composition.

This module provides a single point of communication with the Insight backend,
composing the Run class for core lifecycle management and adding:

- Session creation for rollouts
- Score submission (rewards to scores)
- OTLP trace export configuration
- StorageListener protocol implementation
"""

from __future__ import annotations

import base64
import logging
import threading
import time
from typing import Any, Dict, List, Literal, Optional, Union

import httpx

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

from .listener import StorageListener

logger = logging.getLogger(__name__)

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


class InsightRunListener(StorageListener):
    """Consolidated listener that handles Insight backend communication.

    This listener composes the Run class for core lifecycle management
    (run creation, heartbeats, metrics, finish) and adds Insight-specific
    functionality:

    - Session creation for rollouts
    - Score submission (rewards to scores)
    - OTLP trace export configuration
    - StorageListener protocol for store integration

    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 for this run.
        run_name: Optional human-readable name for the run.
        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).
        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,
        # Run configuration
        run_name: 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,
        # Request configuration
        request_timeout: float = 10.0,
        max_retries: int = 3,
        # Legacy: if run_id is provided, skip run creation (backward compat)
        run_id: Optional[str] = None,
    ) -> None:
        self._api_key = api_key
        self._secret_key = secret_key
        self._insight_url = insight_url.rstrip("/")
        self._project_id = project_id
        self._request_timeout = request_timeout
        self._max_retries = max_retries

        # Create the Run instance (core lifecycle management)
        # Pass self as listener so Run can notify us of events
        self._run = Run(
            name=run_name,
            project_id=project_id,
            api_url=insight_url,
            api_key=secret_key,  # Run expects secret_key as api_key
            config=config,
            labels=labels,
            tags=tags,
            source=source,
            heartbeat_interval=heartbeat_interval,
            listener=self,  # Allow Run to notify us
        )

        # If run_id is pre-set (backward compat), inject it
        if run_id:
            self._run._run_id = run_id

        # Session mapping: rollout_id -> session_id
        self._session_map: Dict[str, str] = {}
        self._session_map_lock = threading.Lock()

        # Track if we've started
        self._started = False

    # ─────────────────────────────────────────────────────────────
    # Properties (delegated to Run)
    # ─────────────────────────────────────────────────────────────

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

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

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

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

    @property
    def is_started(self) -> bool:
        """Check if the run has been started."""
        return self._started

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

    @property
    def capabilities(self) -> Dict[str, bool]:
        """Return the capabilities of the listener."""
        return {
            "otlp_traces": True,
        }

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

    def get_otlp_headers(self) -> Dict[str, str]:
        """Return the authentication headers for OTLP export."""
        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. We handle run creation in start()."""
        pass

    async def on_rollout_created(self, rollout: Rollout) -> None:
        """Called when a rollout is created. Creates a session via API."""
        session_id = self._create_session_for_rollout(rollout)
        if session_id:
            with self._session_map_lock:
                self._session_map[rollout.rollout_id] = session_id
            logger.debug(f"Created session {session_id} for rollout {rollout.rollout_id}")

    async def on_rollout_updated(self, rollout: Rollout) -> None:
        """Called when a rollout is updated."""
        pass

    async def on_attempt_created(self, attempt: Attempt) -> None:
        """Called when an attempt is created."""
        pass

    async def on_attempt_updated(self, attempt: Attempt, rollout_id: str) -> None:
        """Called when an attempt is updated."""
        pass

    async def on_span_created(self, span: Span) -> None:
        """Called when a span is added. Converts rewards to scores."""
        if is_reward_span(span):
            rewards = get_rewards_from_span(span)
            # Look up session ID from rollout_id for session-level scoring
            session_id = self.get_session_id(span.rollout_id) if span.rollout_id else None
            logger.info(
                f"Reward span detected: rollout={span.rollout_id}, "
                f"session={session_id}, trace={span.trace_id}, "
                f"rewards={[(r.name, r.value) for r in rewards]}"
            )
            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,
                    session_id=session_id,
                )

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

    # ─────────────────────────────────────────────────────────────
    # Run Lifecycle (delegated to Run)
    # ─────────────────────────────────────────────────────────────

    def start(self) -> "InsightRunListener":
        """Start the run: create on backend, start heartbeat, register handlers.

        Returns:
            self for chaining
        """
        if self._started:
            logger.warning("Run already started")
            return self

        # Delegate to Run._start() which handles:
        # - Run creation on backend
        # - Heartbeat loop
        # - Signal handlers
        # - atexit handler
        self._run._start()
        self._started = True

        logger.info(f"InsightRunListener started with run {self._run.id} " f"(project={self._project_id})")

        return self

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

        Args:
            state: Final state ("completed", "failed", "cancelled").
        """
        self._run.finish(state=state)

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

        Args:
            summary: Optional summary dict (logged but not sent to API currently).
        """
        if summary:
            logger.debug(f"Run completed with summary: {list(summary.keys())}")
        self._run.finish(state="completed")

    def fail(self, error: Optional[str] = None):
        """Mark the run as failed.

        Args:
            error: Optional error message (logged but not sent to API).
        """
        if error:
            logger.error(f"Run failed: {error}")
        self._run.finish(state="failed")

    # ─────────────────────────────────────────────────────────────
    # Metrics Logging (delegated to Run)
    # ─────────────────────────────────────────────────────────────

    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.

        Delegates to Run.log() which handles buffering and batch sending.

        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.
            metric_type: Type hint for visualization.
        """
        self._run.log(
            metrics,
            value,
            step=step,
            labels=labels,
            unit=unit,
            metric_type=metric_type,
        )

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

        Delegates to Run.log_state() for Tables-like functionality.
        See Run.log_state() for full documentation.

        Args:
            state: Dictionary containing structured state data.
        """
        self._run.log_state(state)

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

        Delegates to Run.flush().
        """
        self._run.flush()

    # ─────────────────────────────────────────────────────────────
    # Session Management (Insight-specific)
    # ─────────────────────────────────────────────────────────────

    def _get_headers(self) -> Dict[str, str]:
        """Get HTTP headers for API requests."""
        headers = {"Content-Type": "application/json"}
        credentials = base64.b64encode(f"{self._api_key}:{self._secret_key}".encode()).decode()
        headers["Authorization"] = f"Basic {credentials}"
        return headers

    def _create_session_for_rollout(self, rollout: Rollout) -> Optional[str]:
        """Create a session for the rollout via the rollouts API."""
        if not self._run.id or self._run.id.startswith("local-"):
            logger.debug(f"Skipping session creation for local run")
            return None

        try:
            # Extract a human-readable name from the rollout input.
            # Handles both direct task dicts ({"name": "..."}) and
            # GEPA-wrapped inputs ({"input": {"name": "..."}, "id": "..."}).
            rollout_name = None
            if isinstance(rollout.input, dict):
                rollout_name = rollout.input.get("name") or rollout.input.get("task_name")
                if not rollout_name and isinstance(rollout.input.get("input"), dict):
                    inner = rollout.input["input"]
                    rollout_name = inner.get("name") or inner.get("task_name")
            if not rollout_name:
                rollout_name = str(rollout.input)[:100] if rollout.input else None
            session_id = rollout.session_id or rollout.rollout_id

            with httpx.Client(timeout=self._request_timeout) as client:
                response = client.post(
                    f"{self._insight_url}/api/public/runs/{self._run.id}/rollouts",
                    headers=self._get_headers(),
                    json={
                        "name": rollout_name,
                        "sessionId": session_id,
                    },
                )
                response.raise_for_status()
                data = response.json()
                return data.get("sessionId")
        except httpx.HTTPError as e:
            logger.warning(f"Failed to create session for rollout {rollout.rollout_id}: {e}")
            return None

    def get_session_id(self, rollout_id: str) -> Optional[str]:
        """Get the session ID for a rollout."""
        with self._session_map_lock:
            return self._session_map.get(rollout_id)

    # ─────────────────────────────────────────────────────────────
    # Score Submission (Insight-specific)
    # ─────────────────────────────────────────────────────────────

    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,
        session_id: Optional[str] = None,
    ) -> None:
        """Send a score to Insight's /api/public/scores endpoint.

        Scores are linked to sessions (rollouts) when session_id is provided,
        enabling session-level grading in the run detail view.
        """
        if not self._run.id or self._run.id.startswith("local-"):
            logger.debug(f"Skipping score submission for local run")
            return

        score_payload: Dict[str, Any] = {
            "name": name,
            "value": value,
            "dataType": "NUMERIC",
        }

        # Link to session if available, otherwise fall back to trace
        if session_id:
            score_payload["sessionId"] = session_id
        elif trace_id:
            score_payload["traceId"] = trace_id
            if observation_id:
                score_payload["observationId"] = observation_id

        metadata: Dict[str, Any] = {
            "source": "insight_run_listener",
            "run_id": self._run.id,
        }
        if rollout_id:
            metadata["rollout_id"] = rollout_id
        if attempt_id:
            metadata["attempt_id"] = attempt_id
        score_payload["metadata"] = metadata

        logger.info(f"Submitting score: name={name}, value={value}, " f"trace={trace_id}, session={session_id}")

        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=self._get_headers(),
                            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.")
                        return
                    logger.warning(
                        f"HTTP error sending score (attempt {attempt + 1}/{self._max_retries}): "
                        f"{e.response.status_code}"
                    )
                except httpx.HTTPError as e:
                    logger.warning(f"Failed to send score (attempt {attempt + 1}/{self._max_retries}): {e}")

                if attempt < self._max_retries - 1:
                    time.sleep(2**attempt)

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

        threading.Thread(target=send_score_async, daemon=True, name="insight-score-sender").start()

    # ─────────────────────────────────────────────────────────────
    # Cleanup
    # ─────────────────────────────────────────────────────────────

    def cleanup(self) -> None:
        """Cleanup resources. Deprecated - use finish() instead."""
        if not self._run._finished:
            self.finish()
