package com.yorvana.testsupport

import androidx.test.platform.app.InstrumentationRegistry
import androidx.test.uiautomator.UiDevice
import org.junit.rules.TestRule
import org.junit.runner.Description
import org.junit.runners.model.Statement

class RetryRule(
    private val maxRetries: Int = 3,
) : TestRule {
    companion object {
        @Volatile
        private var clearedOutputDir = false
    }

    override fun apply(
        base: Statement,
        description: Description,
    ): Statement {
        return object : Statement() {
            override fun evaluate() {
                val outputDir = "/data/local/tmp/flakes"
                val testName = "${description.className}#${description.methodName}"
                val sanitizedTestName =
                    testName
                        .replace("[", "_")
                        .replace("]", "_")
                        .replace("=", "_")
                        .replace(" ", "_")
                        .replace("#", "_")
                val markerFile = "$outputDir/flake-$sanitizedTestName.txt"

                // Delete the entire output directory once per test suite execution
                if (!clearedOutputDir) {
                    synchronized(RetryRule::class.java) {
                        if (!clearedOutputDir) {
                            try {
                                runShell("rm -rf $outputDir")
                                clearedOutputDir = true
                            } catch (e: Exception) {
                                android.util.Log.e("RetryRule", "Failed to clear output dir", e)
                            }
                        }
                    }
                }

                // Delete any existing flake marker for this test before running it
                try {
                    runShell("rm -f $markerFile")
                } catch (e: Exception) {
                    android.util.Log.e("RetryRule", "Failed to clear existing flake marker file", e)
                }

                var caughtThrowable: Throwable? = null
                for (i in 0 until maxRetries) {
                    try {
                        base.evaluate()
                        return
                    } catch (t: Throwable) {
                        if (caughtThrowable == null) {
                            caughtThrowable = t
                        } else {
                            caughtThrowable.addSuppressed(t)
                        }
                        if (t is org.junit.AssumptionViolatedException) {
                            throw t
                        }

                        android.util.Log.e("RetryRule", "${description.displayName}: Run ${i + 1} failed.", t)

                        // Write flake marker file
                        try {
                            runShell("mkdir -p $outputDir")
                            runShell("touch $markerFile")
                        } catch (e: Exception) {
                            android.util.Log.e("RetryRule", "Failed to write flake marker file", e)
                        }
                    }
                }
                throw caughtThrowable ?: AssertionError("Test failed after $maxRetries attempts")
            }
        }
    }

    private fun runShell(command: String): String {
        val instrumentation = InstrumentationRegistry.getInstrumentation()
        val device = UiDevice.getInstance(instrumentation)
        return device.executeShellCommand(command)
    }
}