package com.suno.android.common_networking.extensions

import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.retryWhen
import kotlin.random.Random

private const val DEFAULT_JITTER_FACTOR = 0.25

/**
 * Retries this Flow only for network-related errors (offline, timeouts, 5xx/429, etc.).
 * - Honors server-provided backoff headers via `Throwable.serverBackoffMillis()`
 * - Exponential backoff with jitter when no server hint is present
 * - Never retries cancellations
 *
 * You can plug your own predicate; by default it tries `isRetryableNetworkIssue()` first,
 * and falls back to your existing `isNoNetworkError()`.
 */
fun <T> Flow<T>.retryOnNetworkErrors(
    maxAttempts: Long = 3,
    initialDelayMs: Long = 300,
    maxDelayMs: Long = 15_000,
    jitterRatio: Double = 0.25,
    respectServerBackoffCap: Boolean = true,
    isRetryableError: (Throwable) -> Boolean = { t ->
        // Prefer the finer-grained classifier if you added it; fall back to your original.
        runCatching { t.isRetryableNetworkError }.getOrElse { false }
    },
): Flow<T> = this.retryWhen { cause, attempt ->
    // Only retry for network-ish problems (this already checks for cancellation)
    if (!isRetryableError(cause)) return@retryWhen false
    // attempts = 0,1,2,... (number of retries already done)
    if (attempt >= maxAttempts) return@retryWhen false

    // Prefer server-suggested backoff if present
    val serverMs = cause.serverBackoffMillis()
    val baseMs = if (serverMs != null) {
        if (respectServerBackoffCap) serverMs.coerceAtMost(maxDelayMs) else serverMs
    } else {
        // exp backoff: initial * 2^attempt, capped to prevent overflow
        val safeAttempt = attempt.coerceAtMost(30).toInt()
        val multiplier = 1L shl safeAttempt
        val exp = if (initialDelayMs > Long.MAX_VALUE / multiplier) {
            Long.MAX_VALUE
        } else {
            initialDelayMs * multiplier
        }
        exp.coerceAtMost(maxDelayMs)
    }

    delay(withJitter(baseMs, jitterRatio))
    true
}

private fun withJitter(
    baseMs: Long,
    ratio: Double,
): Long {
    if (ratio <= 0.0) return baseMs
    val wiggle = (baseMs * ratio).toLong().coerceAtLeast(1L)
    val min = (baseMs - wiggle).coerceAtLeast(0L)
    val max = baseMs + wiggle
    return Random.nextLong(min, max + 1)
}

@Suppress("TooGenericExceptionCaught")
suspend fun <T> retryOnNetworkErrors(
    maxAttempts: Int = 3,
    baseDelayMs: Long = 300,
    block: suspend () -> T,
): T {
    var attempt = 0
    while (attempt < maxAttempts) {
        try {
            return block()
        } catch (e: Exception) {
            if (e.isRetryableNetworkError) {
                val backoff = baseDelayMs * (1L shl attempt)
                val jitter = (backoff * DEFAULT_JITTER_FACTOR).toLong()
                delay(backoff + (0..jitter.toInt()).random())
                attempt++
            } else {
                throw e
            }
        }
    }
    error("Invalid retryOnNetworkErrors configuration: maxAttempts = $maxAttempts")
}