/*
 * Copyright 2013-2026, Seqera Labs
 *
 * 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 nextflow.trace

import java.nio.file.Files
import java.nio.file.Path
import java.util.regex.Pattern

import groovy.json.JsonSlurper
import groovy.json.StringEscapeUtils
import groovy.transform.CompileStatic
import groovy.transform.Memoized
import groovy.transform.PackageScope
import groovy.util.logging.Slf4j
import nextflow.cloud.types.CloudMachineInfo
import nextflow.container.resolver.ContainerMeta
import nextflow.extension.Bolts
import nextflow.processor.TaskId
import nextflow.script.ProcessDef
import nextflow.util.Duration
import nextflow.util.KryoHelper
import nextflow.util.MemoryUnit
/**
  * This object represent holds the information of a single process run,
  * its content is saved to a trace file line
  *
  * @author Paolo Di Tommaso <paolo.ditommaso@gmail.com>
  */
@Slf4j
@CompileStatic
class TraceRecord implements Serializable {

    // note: ?i stands for ignore case - ?m stands for multiline
    static public final Pattern SECRET_REGEX = ~/(?im)(^AWS[^=]*|.*TOKEN[^=]*|.*SECRET[^=]*)=(.*)$/

    TraceRecord() {
        this.store = new LinkedHashMap<>(FIELDS.size())
    }

    @PackageScope
    TraceRecord(Map store) {
        this.store = store
    }

    final private static String DEFAULT_DATE_FORMAT = "yyyy-MM-dd HH:mm:ss.SSS"

    final public static String NA = '-'

    final public static List<String> NON_PRIMITIVE_TYPES = ['date','time','perc','mem']

    final public static Map<String,String> FIELDS = [
            task_id:    'str',
            hash:       'str',
            native_id:  'str',
            process:    'str',
            module:     'str',
            container:  'str',
            tag:        'str',
            name:       'str',
            status:     'str',
            exit:       'str',
            submit:     'date',
            start:      'date',
            complete:   'date',
            duration:   'time',
            realtime:   'time',
            '%cpu':     'perc',     // -- ps field '%cpu'
            '%mem':     'perc',     // -- ps field '%mem'
            rss:        'mem',      // -- ps field 'rss'
            vmem:       'mem',      // -- ps field 'vsize'
            peak_rss:   'mem',      // -- /proc/$pid/status field 'VmHWM'  (Peak resident set size i.e. high water mark)
            peak_vmem:  'mem',      // -- /proc/$pid/status field 'VmPeak' (Peak virtual memory size)
            rchar:      'mem',      // -- /proc/$pid/io
            wchar:      'mem',      // -- /proc/$pid/io
            syscr:      'num',      // -- /proc/$pid/io
            syscw:      'num',      // -- /proc/$pid/io
            read_bytes: 'mem',      // -- /proc/$pid/io
            write_bytes:'mem',      // -- /proc/$pid/io
            attempt:    'num',
            workdir:    'str',
            script:     'str',
            scratch:    'str',
            queue:      'str',
            cpus:       'num',
            memory:     'mem',
            disk:       'mem',
            time:       'time',
            env:        'str',
            error_action:'str',
            vol_ctxt: 'num',        // -- /proc/$pid/status field 'voluntary_ctxt_switches'
            inv_ctxt: 'num',        // -- /proc/$pid/status field 'nonvoluntary_ctxt_switches'
            hostname: 'str',
            cpu_model:  'str',
            accelerator: 'num',
            accelerator_type: 'str'
    ]

    static public Map<String,Closure<String>> FORMATTER = [
            str: this.&fmtString,
            num: this.&fmtNumber,
            date: this.&fmtDate,
            time: this.&fmtTime,
            mem: this.&fmtMemory,
            perc: this.&fmtPercent
    ]

    @PackageScope
    static TimeZone TIMEZONE = null

    transient private String executorName
    transient private CloudMachineInfo machineInfo
    transient private ContainerMeta containerMeta
    transient private Integer numSpotInterruptions
    transient private String logStreamId
    transient private Map<String,Object> resourceAllocation
    transient private Map<String,Object> gpuMetrics

    /**
     * Convert the given value to a string
     *
     * @param fmt
     * @return The value as a string
     */
    @PackageScope
    static String fmtString(value, String fmt) {
        if( value instanceof Number )
            return value != Integer.MAX_VALUE ? value.toString() : NA

        value ? value.toString() : NA
    }


    /**
     * Convert the given value to a number.
     * <p>Objects of type {@link Duration} (time) are  converted to milliseconds
     * <p>Objects of type {@link MemoryUnit) are converted to number of bytes
     * <p>Objects of type {@link Date) are converted to milliseconds since Unix epoch
     *
     * @param fmt
     * @return The value as a string
     */
    @PackageScope
    static String fmtNumber(def value, String fmt) {
        if( value == null )
            return NA

       if( value instanceof Duration )
            return String.valueOf(value.toMillis())

        if( value instanceof MemoryUnit )
            return String.valueOf(value.toBytes())

        if( value instanceof Date )
            return String.valueOf(value.getTime())

        return value.toString()
    }

    /**
     * Converts the value to a date string
     *
     * @param value The value is supposed to be the number of milliseconds since Unix epoch
     * @param fmt
     * @return The formatted date string
     */
    @PackageScope
    static String fmtDate(def value, String fmt) {
        if( !value ) return NA
        if( !fmt )
            fmt = DEFAULT_DATE_FORMAT
        Bolts.format(new Date(value as long), fmt, TIMEZONE)
    }

    /**
     * Converts the value to a duration string.
     *
     * See {@link Duration}
     * @param value
     * @param fmt
     * @return
     */
    @PackageScope
    static String fmtTime(def value, String fmt) {
        if( value == null ) return NA
        new Duration(value as long).toString()
    }

    /**
     * Converts the value to a memory unit string
     * <p>
     * See {@link MemoryUnit}
     *
     * @param value
     * @param fmt
     * @return
     */
    @PackageScope
    static String fmtMemory( def value, String fmt ) {
        if( value == null ) return NA

        if( value instanceof Number )
            return new MemoryUnit(value.toLong()).toString()

        String str = value.toString()
        if( str.isLong() ) {
            return new MemoryUnit(str.toLong()).toString()
        }

        return str
    }

    /**
     * Converts the value to a percent value string
     * @param value
     * @param fmt
     * @return
     */
    @PackageScope
    static String fmtPercent( def value, String fmt ) {
        if( value == null ) return NA
        try {
            if( value instanceof Number )
                return String.format(Locale.ROOT, '%.1f%%', value.toFloat())
            else {
                def x = value.toString().toFloat()
                return String.format(Locale.ROOT, '%.1f%%', x)
            }

        }
        catch( Exception e ) {
            log.debug "Not a valid percentage value: '$value'"
            return NA
        }
    }

    @PackageScope
    Map<String,Object> store

    Map<String,Object> getStore() { store }

    @Memoized
    Set<String> keySet() {
        FIELDS.keySet()
    }

    def propertyMissing(String name, value) {
        put(name,value)
    }

    def propertyMissing(String name) {
        get(name)
    }

    def containsKey( String name ) {
        assert name
        store.containsKey(name)
    }

    def get( String name ) {
        assert keySet().contains(name), "Not a valid TraceRecord field: '$name'"
        if( name == 'env' ) {
            final ret = store.get(name)
            return ret ? secureEnvString(ret.toString()) : ret
        }
        return store.get(name)
    }

    protected String secureEnvString( String str ) {
        str.replaceAll(SECRET_REGEX, '$1=[secure]')
    }


    void put( String name, def value ) {
        if( !keySet().contains(name) ) {
            log.warn1 "Unknown trace record field: $name"
            return
        }

        // vmpeak: Peak virtual memory size
        // this is a synonym of 'max_vmem' field
        if( name == 'vmpeak' ) {
            store.put('max_vmem', value)
        }

        // Peak resident set size ("high water mark")
        // This is a synonym of 'max_rss' field
        else if( name == 'vmhwm' ) {
            store.put('max_rss', value)
        }

        else if( name == 'env' ) {
            store.put(name, value ? secureEnvString(value.toString()) : value)
        }

        else {
            store.put(name, value)
        }
    }

    void putAll( Map<String,Object> values ) {
        if( !values )
            return

        for( String key : values.keySet() ) {
            put(key, values.get(key))
        }
    }

    /**
     * Get a trace field value and apply a conversion rule to it
     *
     * @param name The field name e.g. task_id, status, etc.
     * @param converter A converter string
     * @return A string value formatted according the specified converter
     */
    String getFmtStr( String name, String converter = null ) {
        assert name
        def val = store.get(name)

        String sType=null
        String sFormat=null
        if( converter ) {
            int p = converter.indexOf(':')
            if( p == -1 ) {
                sType = converter
            }
            else {
                sType = converter.substring(0,p)
                sFormat = converter.substring(p+1)
            }
        }

        def type = sType ?: FIELDS.get(name)
        if( !type )
            throw new IllegalArgumentException("Not a valid trace field name: '$name'")


        def formatter = FORMATTER.get(type)
        if( !formatter )
            throw new IllegalArgumentException("Not a valid trace formatter for field: '$name' with type: '$type'")

        try {
            return formatter.call(val,sFormat)
        }
        catch( Throwable e ) {
            log.debug "Not a valid trace value -- field: '$name'; value: '$val'; format: '$sFormat'"
            return null
        }
    }

    TaskId getTaskId() { TaskId.of(get('task_id')) }

    String getWorkDir() { get('workdir') }

    String getProcessName() { get('process') }

    String getSimpleName() { ProcessDef.stripScope(processName) }

    /**
     * Render the specified list of fields to a single string value
     *
     * @param fields The list of fields to be rendered, each entry can optionally specify a
     *      a format string separating it from the name by a colon character e.g. name:format
     * @param separator A delimiter that separates fields entries in the final string
     * @return The final string containing all field values
     */
    String renderText( List<String> fields, List<String> formats, String separator ) {
        def result = new ArrayList(fields.size())
        for( int i=0; i<fields.size(); i++ ) {
            String name = fields[i]
            String format = i<formats?.size() ? formats[i] : null
            result << (getFmtStr(name, format) ?: NA)
        }

        return result.join(separator)
    }

    CharSequence renderJson(StringBuilder result, List<String> fields, List<String> formats) {
        final QUOTE = '"'
        if( result == null ) result = new StringBuilder()
        result << "{"
        for( int i=0; i<fields.size(); i++ ) {
            if( i ) result << ','
            String name = fields[i]
            String format = i<formats?.size() ? formats[i] : null
            String value = StringEscapeUtils.escapeJavaScript(getFmtStr(name, format) ?: NA)
            result << QUOTE << name << QUOTE << ":" << QUOTE << value << QUOTE
        }
        result << "}"
        return result
    }

    CharSequence renderJson(StringBuilder result = new StringBuilder()) {
        List<String> fields = []
        List<String> formats = []
        FIELDS.each { name, type -> fields << name; formats << type }
        renderJson(result, fields, formats)
    }

    String toString() {
        "${this.class.simpleName} ${this.store}"
    }


    /**
     * Parse the trace file
     * <p>
     * Trace example:
     * <pre>
     * pid state %cpu %mem vmem rss max_vmem max_rss rchar wchar syscr syscw read_bytes write_bytes
     *  1 0 0 0 11084 1220 11084 1220 4790 12 11 1 0 0 0
     * </pre>
     *
     */
    TraceRecord parseTraceFile( Path file ) {

        try(BufferedReader reader = Files.newBufferedReader(file)) {
            final lines = reader.readLines()
            if( !lines )
                return this
            if( lines[0] != 'nextflow.trace/v2' )
                return parseLegacy(file, lines)

            for( int i=0; i<lines.size(); i++ ) {
                final pair = lines[i].tokenize('=')
                final name = pair[0]
                final value = pair[1]
                if( value == null )
                    continue

                switch (name) {
                    case '%cpu':
                    case '%mem':
                        // fields '%cpu' and '%mem' are expressed as percent value
                        this.put(name, parseInt(value, file, name) / 10F)
                        break

                    case 'rss':
                    case 'vmem':
                    case 'peak_rss':
                    case 'peak_vmem':
                        // these fields are provided in KB, so they are normalized to bytes
                        def val = parseLong(value, file, name) * 1024
                        this.put(name, val)
                        break

                    case 'cpu_model':
                        this.put(name, value)
                        break

                    default:
                        def val = parseLong(value, file, name)
                        this.put(name, val)
                        break
                }

            }

            return this
        }
    }

    private TraceRecord parseLegacy( Path file, List<String> lines) {
        String[] header = null
        for( int count=0; count<lines.size(); count++ ) {
            String row = lines[count]

            /*
             * 1st line -- parse the header
             */
            if( count == 0 ) {
                header = row.trim().split(/\s+/)
            }

            /*
             * 2nd line -- parse values produced by 'ps'
             */
            else if( count == 1 ) {
                String[] values = row.trim().split(/\s+/)
                for( int i=0; i<values.length; i++ ) {

                    final name = header[i]
                    if( i==2 || i==3 ) {
                        // fields '%cpu' and '%mem' are expressed as percent value
                        this.put(name, parseInt(values[i], file, i) / 10F)
                    }
                    else if( i>3 ) {
                        def val = parseLong(values[i], file, i)
                        // fields from index 4 to 7 (vmem,rss,peak_vmem, peak_rss) are provided in KB, so they are normalized to bytes
                        if( i<8 ) val *= 1024
                        this.put(name, val)
                    }

                }
            }

            // third line is the cpu realtime
            else if( count == 2 ) {
                try {
                    def elapsed = row.toString().trim().toLong()
                    this.put('realtime', elapsed)
                }
                catch( Exception e ) {
                    log.debug "Not a valid trace `realtime` value: '$count'"
                }
            }

        }

        return this
    }

    /**
     * Parse the Fusion trace file (.fusion/trace.json) produced by the Fusion client.
     * The file contains task metrics in the following format:
     * <pre>
     * {
     *   "proc": {
     *     "realtime": 660541, "pct_cpu": 1045, "cpu_name": "Intel Xeon ...",
     *     "rchar": 14112539262, "wchar": 12668821375, "syscr": 1823378, "syscw": 169293,
     *     "read_bytes": 8011776, "write_bytes": 102400,
     *     "vmem": 39015152, "rss": 14826068, "peak_vmem": 39047920, "peak_rss": 15775480,
     *     "pct_mem": 56, "vol_ctxt": 413015, "inv_ctxt": 1540
     *   },
     *   "gpu": {
     *     "name": "Tesla T4", "mem": 15360, "driver": "580.126.09",
     *     "pct": 75, "peak": 100, "pct_mem": 40.1, "peak_mem": 74.1
     *   },
     *   "cgroup": {
     *     "version": "v2", "memory_current": 25469927424, "memory_peak": 41178980352,
     *     "memory_rss": 67919872, "memory_peak_rss": 14783070208,
     *     "memory_limit": 77309411328, "cpu_usage_usec": 785302059
     *   }
     * }
     * </pre>
     *
     * @param file Path to the .fusion/trace.json file
     * @return The parsed JSON as a map with keys: {@code proc}, {@code gpu}, {@code cgroup}
     */
    static Map<String,Object> parseFusionTraceFile(Path file) {
        return (Map<String,Object>) new JsonSlurper().parse(file)
    }

    static long toLong(Object value) {
        if( value instanceof Number )
            return value.longValue()
        return value != null ? value.toString().toLong() : 0L
    }

    static float toFloat(Object value) {
        if( value instanceof Number )
            return value.floatValue()
        return value != null ? value.toString().toFloat() : 0.0f
    }

    private long parseInt( String str, Path file, String row )  {
        try {
            str.toInteger()
        }
        catch( NumberFormatException e ) {
            log.debug "[WARN] Not a valid integer number `$str` -- offending row: $row in file `$file`"
            return 0
        }
    }

    private long parseLong( String str, Path file , String row )  {
        try {
            str.toLong()
        }
        catch( NumberFormatException e ) {
            log.debug "[WARN] Not a valid long number `$str` -- offending row: $row in file `$file`"
            return 0
        }
    }

    @Deprecated
    private long parseInt( String str, Path file, int index )  {
        try {
            str.toInteger()
        }
        catch( NumberFormatException e ) {
            log.debug "[WARN] Not a valid integer number `$str` -- offending column: $index in file `$file`"
            return 0
        }
    }

    @Deprecated
    private long parseLong( String str, Path file, int index )  {
        try {
            str.toLong()
        }
        catch( NumberFormatException e ) {
            log.debug "[WARN] Not a valid long number `$str` -- offending column: $index in file `$file`"
            return 0
        }
    }

    @Override
    boolean equals( Object that ) {
        if ( this.is(that) ) return true
        if ( !(that instanceof TraceRecord) ) return false
        this.store.equals(((TraceRecord)that).store)
    }

    @Override
    int hashCode() {
        store.hashCode()
    }

    byte[] serialize() {
        KryoHelper.serialize(store)
    }

    static TraceRecord deserialize(byte[] buffer) {
        Map map = (Map)KryoHelper.deserialize(buffer)
        new TraceRecord(map)
    }

    TraceRecord setCached(boolean value) {
        if( value ) {
            store.status = 'CACHED'
        }
        return this
    }

    boolean isCached() {
        store.status == 'CACHED'
    }

    boolean isCompleted() {
        store.status == 'COMPLETED'
    }

    String getExecutorName() {
        return executorName
    }

    void setExecutorName(String value ) {
        this.executorName = value
    }

    CloudMachineInfo getMachineInfo() {
        return machineInfo
    }

    void setMachineInfo(CloudMachineInfo value) {
        this.machineInfo = value
    }

    Integer getNumSpotInterruptions() {
        return numSpotInterruptions
    }

    void setNumSpotInterruptions(Integer numSpotInterruptions) {
        this.numSpotInterruptions = numSpotInterruptions
    }

    String getLogStreamId() {
        return logStreamId
    }

    void setLogStreamId(String logStreamId) {
        this.logStreamId = logStreamId
    }

    ContainerMeta getContainerMeta() {
        return containerMeta
    }

    void setContainerMeta(ContainerMeta meta) {
        this.containerMeta = meta
    }

    Map<String,Object> getResourceAllocation() {
        return resourceAllocation
    }

    void setResourceAllocation(Map<String,Object> value) {
        this.resourceAllocation = value
    }

    Map<String,Object> getGpuMetrics() {
        return gpuMetrics
    }

    void setGpuMetrics(Map<String,Object> value) {
        this.gpuMetrics = value
    }
}
