/*
 * Copyright (C) 2023 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package android.tools.traces.monitors

import android.tools.function.Supplier
import android.tools.io.TraceType
import android.tools.traces.io.IResultData
import android.tools.traces.io.IoUtils
import android.tools.traces.io.ResultReader
import android.tools.traces.io.ResultWriter
import android.tools.withTracing
import java.io.File
import kotlin.io.path.createTempDirectory

/**
 * Base class for monitors containing common logic to read the trace as a byte array and save the
 * trace to another location.
 */
abstract class TraceMonitor : ITransitionMonitor {
    abstract val isEnabled: Boolean
    abstract val traceType: TraceType

    protected abstract fun doStart()

    protected abstract fun doStop(): File

    protected open fun doStopTraces(): Map<TraceType, File> = mapOf(traceType to doStop())

    final override fun start() {
        withTracing("${this::class.simpleName}#start") {
            validateStart()
            doStart()
        }
    }

    open fun validateStart() {
        if (this.isEnabled) {
            throw UnsupportedOperationException(
                "${traceType.name} trace already running. " +
                    "This is likely due to chained 'withTracing' calls."
            )
        }
    }

    /** Stops monitor. */
    override fun stop(writer: ResultWriter) {
        val artifacts =
            try {
                withTracing("${this::class.simpleName}#stop") {
                    doStopTraces()
                        .map { (key, value) -> key to moveTraceFileToTmpDir(value) }
                        .toMap()
                }
            } catch (e: Throwable) {
                throw RuntimeException(
                    "Could not stop ${traceType.name} trace and save it to ${traceType.fileName}",
                    e,
                )
            }
        artifacts.forEach { (key, value) -> writer.addTraceResult(key, value) }
    }

    private fun moveTraceFileToTmpDir(sourceFile: File): File {
        val newFile = File.createTempFile(sourceFile.name, "")
        IoUtils.moveFile(sourceFile, newFile)
        require(newFile.exists()) { "Unable to save trace file $newFile" }
        return newFile
    }

    /**
     * Uses [writer] to write the trace generated by executing the commands defined by [predicate].
     *
     * @param writer Write to use to write the collected traces
     * @param predicate Commands to execute
     * @throws UnsupportedOperationException If tracing is already activated
     */
    fun withTracing(writer: ResultWriter, predicate: Runnable) {
        withTracing("${this::class.simpleName}#withTracing") {
            try {
                this.start()
                predicate.run()
            } finally {
                this.stop(writer)
            }
        }
    }

    /**
     * Acquires the trace generated when executing the commands defined in the [predicate].
     *
     * @param predicate Commands to execute
     * @param resultReaderProvider Predicate to generate new result readers
     * @param
     * @throws UnsupportedOperationException If tracing is already activated
     */
    fun withTracing(
        resultReaderProvider: Supplier<IResultData, ResultReader>,
        predicate: Runnable,
    ): ResultReader {
        val writer = createWriter()
        withTracing(writer, predicate)
        val result = writer.write()
        return resultReaderProvider.get(result)
    }

    private fun createWriter(): ResultWriter {
        val className = this::class.simpleName ?: error("Missing class name for $this")
        val tmpDir = createTempDirectory("withTracing").toFile()
        return ResultWriter().withName(className).withOutputDir(tmpDir)
    }

    companion object {
        @JvmStatic protected val TRACE_DIR = File("/data/misc/wmtrace/")
    }
}