#### Benben
#### Copyright (C) 2023-2024 Remilia Scarlet <remilia@posteo.jp>
####
#### This program is free software: you can redistribute it and/or
#### modify it under the terms of the GNU Affero General Public
#### License as published by the Free Software Foundation, either
#### version 3 of the License, or (at your option) any later version.
####
#### This program is distributed in the hope that it will be useful,
#### but WITHOUT ANY WARRANTY; without even the implied warranty of
#### MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#### Affero General Public License for more details.
####
#### You should have received a copy of the GNU Affero General Public License
#### along with this program.  If not, see <https://www.gnu.org/licenses/>.
require "./outputstream"
require "./job"

module Benben::Rendering
  class RenderingError < Exception
    getter outFilename : Path?

    def initialize(@outFilename, msg : String, cause : Exception)
      super(msg, cause)
    end
  end

  class NormalizationError < RenderingError
  end

  alias JobChannel = Channel(Bool|RenderingError|NormalizationError)
  alias ProgressChan = Channel(UInt32|Bool)
  alias ProgressDoneChan = Channel(BigFloat)
  alias ReceiverChan = Channel(Bool)

  # The `Renderer` class acts as a manager for rendering files in parallel.
  class Renderer
    @errors = [] of RenderingError|NormalizationError
    @baseOutputPath : String?

    def initialize
      unless Benben.args["quiet"].called?
        puts Banners.getBanner
      end
    end

    # Renders all of the files that were specified on the command line (or
    # indirectly via playlists).
    def render : Nil
      if Benben.fileHandler.size == 0
        RemiLib.log.fatal("No supported audio files found to render")
      end

      # Create various channels we'll use
      receiverDoneChan : ReceiverChan = ReceiverChan.new(1)
      pbarDoneChan : ProgressDoneChan = ProgressDoneChan.new(1)

      # Keep track of how long this takes
      start : Time::Span = Time.monotonic

      # Loop, creating jobs for all the audio files we have.  Note that this
      # will load all files into RAM.
      jobChan : JobChannel = JobChannel.new(Benben.args["jobs"].str.to_i32 + 3)

      unless Benben.args["quiet"].called?
        puts "Loading files into RAM..."
      end
      jobs : Array(Job) = buildJobs(jobChan)
      RemiLib.log.fatal("Nothing to do :(") if jobs.empty?

      # Start the receiver
      startReceiver(jobChan, receiverDoneChan)

      # Start the progress bar.  We use our own progress bar, not the one in the
      # Printer.
      pbarChan : ProgressChan = startProgressBar(jobs, pbarDoneChan)

      # Start the pool
      unless Benben.args["quiet"].called?
        format("Files to render: ~:d~%", jobs.size)
        puts "EQ disabled by default for rendering" if Benben.config.noEQDuringRendering?
        puts "Starting rendering jobs..."
      end
      pool : RemiLib::JobPool(Job) = RemiLib::JobPool(Job).new(Benben.args["jobs"].str.to_i32)
      totalFileSize : UInt64 = 0u64
      pool.sleepCheckInterval = 5.milliseconds

      renderStart : Time::Span = Time.monotonic
      pool.run(jobs) do |job|
        job.render(pbarChan)
        if File.exists?(job.outFilename) && !job.errored?
          totalFileSize += File.size(job.outFilename)
          job.normalizeFile(pbarChan) if Benben.args["normalize"].called?
        end
        pbarChan.send(true)
      end
      renderEnd : Time::Span = Time.monotonic

      # Tell the receiver fiber and progress bar fiber to stop.
      jobChan.send(false)
      pbarChan.send(false)
      receiverDoneChan.receive
      avgSamples : BigFloat = pbarDoneChan.receive

      Benben.args.withCalledStringArg("cue") do |arg|
        if File.exists?(arg.str) && !Benben.args["overwrite"].called?
          RemiLib.log.error("Cannot generate CUE: file exists")
        else
          generateCue(jobs, arg.str)
        end
      end

      # Print any errors
      printErrors(pool)

      # Print how much time was taken
      finish : Time::Span = Time.monotonic
      delta : Time::Span = finish - start

      unless Benben.args["quiet"].called?
        format("Number of files rendered: ~:d~%", Math.max(0, jobs.size - pool.errors.size - @errors.size))
        format("Total size of rendered files: ~a~%", totalFileSize.prettySize)
        format("Time taken: ~a~%", delta)

        # Report the time spent rendering and normalizing separate
        delta = renderEnd - renderStart
        format("  Rendering: ~a~%", delta)
        format("Average samples per second: ~:d~%", avgSamples.to_u64!)
      end
    end

    {% begin %}
      {% for info in [:Au, :Wav, :Qoa] %}
        private def getJobOutputFilename{{info.id}}(origFilename : Path) : Path
          ext : String = origFilename.extension
          if ext.empty?
            Path["#{origFilename}.#{ {{info.id.stringify.downcase}} }"]
          else
            Path[origFilename.to_s.gsub(ext, ".#{ {{info.id.stringify.downcase}} }")]
          end
        end
      {% end %}
    {% end %}

    private def getJobOutputFilename(origFilename : Path, outputPath : String) : Path
      # Construct a filename based on the desired output format.
      filename : Path = case
                        when Benben.args["au"].called?
                          getJobOutputFilenameAu(origFilename)
                        when Benben.args["qoa"].called?
                          getJobOutputFilenameQoa(origFilename)
                        else
                          getJobOutputFilenameWav(origFilename)
                        end
      Path[outputPath, filename]
    end

    # Creates an Array of Job instances based on the requested files.
    private def buildJobs(jobChan : JobChannel) : Array(Job)
      # The array we'll store Jobs in
      jobs : Array(Job) = [] of Job

      # Used to check for duplicate output filenames.
      outputPaths : Hash(Path, Bool) = {} of Path => Bool

      # Get each file
      loop do
        file = Benben.fileHandler.next
        break if file.nil?
        if file.is_a?(PcmFile)
          RemiLib.log.warn("File is already a PCM file, ignoring: " \
                           "#{Benben.fileHandler.currentFilename}")
          next
        end

        # Get the output path
        outputPath : String = getOutputPath
        outputPath = Job.adjustOutputPathForFile(outputPath, file)

        # Check if the output path exists, and if it doesn't, attempt to create
        # it.
        unless Dir.exists?(outputPath)
          begin
            Dir.mkdir_p(outputPath)
          rescue err : Exception
            RemiLib.log.fatal("Could not create output path: #{err}")
            next
          end
        end

        # Construct the filename.
        filename : Path = Path[Path[Benben.fileHandler.currentFilename].basename]
        filename = getJobOutputFilename(filename, outputPath)

        # Check for duplicate destination filenames.
        if outputPaths.has_key?(filename)
          RemiLib.log.error("Two files have the same output name, rendering only once: #{filename}")
          next
        else
          outputPaths[filename] = true
        end

        # Check for an existing file.
        if File.exists?(filename)
          unless Benben.args["overwrite"].called?
            RemiLib.log.error("File exists: #{filename}")
            next
          end
        end

        # All good, store the new job
        case file
        in VgmFile
          jobs << VgmJob.new(file.as(VgmFile), Path[Benben.fileHandler.currentFilename], filename, jobChan)
        in ModuleFile
          jobs << ModuleJob.new(file.as(ModuleFile), Path[Benben.fileHandler.currentFilename], filename, jobChan)
        in FlacFile
          jobs << FlacJob.new(file.as(FlacFile), Path[Benben.fileHandler.currentFilename], filename, jobChan)
        in OpusFile
          jobs << OpusJob.new(file.as(OpusFile), Path[Benben.fileHandler.currentFilename], filename, jobChan)
        in VorbisFile
          jobs << VorbisJob.new(file.as(VorbisFile), Path[Benben.fileHandler.currentFilename], filename, jobChan)
        in Mpeg1File
          jobs << Mpeg1Job.new(file.as(Mpeg1File), Path[Benben.fileHandler.currentFilename], filename, jobChan)
        in MidiFile
          jobs << MidiJob.new(file.as(MidiFile), Path[Benben.fileHandler.currentFilename], filename, jobChan)
        in QoaFile
          jobs << QoaJob.new(file.as(QoaFile), Path[Benben.fileHandler.currentFilename], filename, jobChan)
        in PcmFile
          raise "Should have been checked earlier"
        in WavPackFile
          jobs << WavPackJob.new(file.as(WavPackFile), Path[Benben.fileHandler.currentFilename], filename, jobChan)
        in PlayableFile
          raise "Unexpectedly received nil PlayableFile"
        end
        Fiber.yield
      end

      jobs
    end

    # Gets the base output directory where rendered files will be stored.
    private def getOutputPath : String
      if @baseOutputPath.nil?
        # Memoize
        @baseOutputPath = if Benben.args["outdir"].called?
                            Benben.args["outdir"].str
                          else
                            Path["."].expand.to_s
                          end
      end

      @baseOutputPath || raise "Failed to memoize baseOutputPath"
    end

    # Starts receiving messages from Job instances over `inChan`.
    private def startReceiver(inChan : JobChannel, receiverDoneChan : ReceiverChan) : Nil
      # Spawn a fiber.
      spawn(name: "Job Notif. Receiver") do
        loop do
          msg = inChan.receive
          if msg.is_a?(Bool)
            if msg
              # Song was finished
            else
              # Time to exit
              receiverDoneChan.send(true)
              break
            end
          elsif msg.is_a?(RenderingError) || msg.is_a?(NormalizationError)
            @errors << msg
          else
            RemiLib.log.error("Received weird message of type #{typeof(msg)} in receiver")
          end
          Fiber.yield
        end # Loop do
      end # Spawn
    end

    # Prints out any errors received by the JobPool.
    private def printErrors(pool : RemiLib::JobPool) : Nil
      unless @errors.empty? && pool.errors.empty?
        RemiLib.log.warn("Errors encountered during rendering!")

        @errors.each do |err|
          RemiLib.log.error("#{err.outFilename}: #{err.message}")
        end

        pool.errors.each do |err|
          {% if flag?(:benben_debug) %}
            # This is rather ugly output-wise, but since it's only during
            # debugging, it's not too concerning.
            RemiLib.log.error("#{err} (#{err.backtrace})")
          {% else %}
            RemiLib.log.error("#{err}")
          {% end %}
        end
      end
    end

    # Calculates the total number of samples, across all jobs, that we'll be
    # rendering.
    private def getTotalSamples(jobs : Array(Job)) : UInt64
      jobs.sum(&.calcTotalSampleSize.to_u64!)
    end

    # Starts the progress bar, then returns a channel that can be used to
    # communicate with the progress bar.  When the progress bar exits, a message
    # will be sent over `pbarDoneChan`.
    private def startProgressBar(jobs : Array(Job), pbarDoneChan : ProgressDoneChan) : ProgressChan
      # Figure out the max for the progress bar.
      totalSize : UInt64 = getTotalSamples(jobs)

      # Get the number of jobs.  We display this on the progress bar.
      totalFiles : Int32 = jobs.size

      # Create the progress bar.
      progressBar = if Benben.args["quiet"].called?
                      nil
                    else
                      RemiLib::Console::ProgressBar.new("Rendering", totalSize)
                    end

      # Create a channel that jobs can use to communicate to the progress bar.
      chan : ProgressChan = ProgressChan.new(Benben.args["jobs"].str.to_i32 + 3)

      # Setup some options.
      progressBar.try do |pb|
        pb.noAutoRefresh = true
        pb.postLabelWidth = 30u64
      end

      # Start a fiber where we'll update the progress bar.
      spawn(name: "Progress Bar") do
        numDone : UInt32 = 0
        samplesDone : UInt64 = 0
        lastUpdate : Time::Span = Time.monotonic
        avgSamples : Int64 = 0
        samplesLastPeriod : UInt64 = 0
        allAvgSamples : BigFloat = BigFloat.new(0)
        periods : UInt64 = 0
        msg : UInt32|Bool = 0
        {% if flag?(:linux) %}
          cpuUsage : Float64 = Utils.getCpu
        {% end %}
        lastCPU = Time.monotonic

        # Start receiving messages.
        loop do
          msg = chan.receive
          case msg
          in UInt32
            # Update the number of samples processed.
            samplesDone += msg
            samplesLastPeriod += msg

          in Bool
            if msg
              # A job has finished
              numDone += 1
            else
              # Completely finished, clean up and then break.
              progressBar.try do |pb|
                pb.label = "Rend: #{numDone}/#{totalFiles}"
                pb.step = pb.max
                pb.refresh
              end
              pbarDoneChan.send(allAvgSamples / (periods == 0 ? 1 : periods))
              break
            end
          end

          # Update the speed once every half second
          if (Time.monotonic - lastUpdate).total_milliseconds >= 500
            avgSamples = (samplesLastPeriod /
                          ((Time.monotonic - lastUpdate).total_nanoseconds / 500_000_000)).to_i64!
            samplesLastPeriod = 0
            allAvgSamples += avgSamples
            periods += 1
            lastUpdate = Time.monotonic
          end

          # Update the progress bar visually.
          progressBar.try do |pb|
            {% if flag?(:linux) %}
              # See if we need to update the CPU usage.  We only do this if we
              # have a progress bar.
              if (Time.monotonic - lastCPU).total_milliseconds >= 500
                cpuUsage = Utils.getCpu
                lastCPU = Time.monotonic
              end
            {% end %}

            # Update progress bar information
            pb.label = "Rend: #{numDone}/#{totalFiles}"
            {% if flag?(:linux) %}
              pb.postLabel = formatStr("~:d Smp/s, CPU: ~a", avgSamples, sprintf("%5.1f%%", cpuUsage))
            {% else %}
              pb.postLabel = formatStr("~:d Smp/s, CPU: ???    ", avgSamples)
            {% end %}
            pb.step = samplesDone
            pb.refresh
          end
          Fiber.yield
        end
      end

      chan
    end

    private def generateCue(jobs : Array(Job), filename : String|Path) : Nil
      puts "Generating CUE file..." unless Benben.args["quiet"].called?
      cuePath : Path = Path[Path[filename].dirname] # Used to determine the path to the rendered file
      cue : RemiAudio::Cue = RemiAudio::Cue.new
      cue.catalog = "0000000000000"

      jobs.each_with_index do |job, idx|
        next if job.errored?
        job.generateCue(cue, cuePath, idx)
      end

      # Write the CUE to the file
      File.open(filename, "w") do |file|
        file << "REM Created with Benben v#{VERSION}\n"
        cue.write(file)
      end
    end
  end
end
