// SPDX-FileCopyrightText: © 2024 LEDGER SAS
// SPDX-License-Identifier: Apache-2.0

import type { CurrencyConfig } from '../../config'
import type { CoinModuleApi, Memo, MemoNotSupported, TxData, TxDataNotSupported } from '../types'
import { instrumentApi } from './instrumenter'

export interface Span {
  end(): void
  setError(error: unknown): void
}

export interface Tracer {
  startSpan(name: string): Span
}

/**
 * Wraps a {@link CoinModuleApi} to open a tracing span around each method invocation.
 *
 * The span is ended on success, or after `setError` on throw.
 * The Context is forwarded unchanged to the underlying method.
 * Handles both synchronous methods (e.g. `combine`) and async methods uniformly.
 */
export function withTracing<
  C extends CurrencyConfig,
  M extends Memo = MemoNotSupported,
  T extends TxData = TxDataNotSupported,
  Extra extends object = Record<never, never>,
>(api: CoinModuleApi<C, M, T> & Extra, tracer: Tracer): CoinModuleApi<C, M, T> & Extra {
  return instrumentApi(api, (name, fn, args) => {
    const span = tracer.startSpan(name)
    let result: unknown
    try {
      result = fn(...args)
    } catch (e) {
      span.setError(e)
      span.end()
      throw e
    }
    if (result instanceof Promise) {
      return result.then(
        (v) => {
          span.end()
          return v
        },
        (e) => {
          span.setError(e)
          span.end()
          throw e
        }
      )
    }
    span.end()
    return result
  })
}