//
//  ToneWriter.swift
//  ToneShaperPreview
//
//  Created by Joseph Pagliaro on 12/7/23.
//

import SwiftUI
import AVFoundation

let kUseSimpsons = false

class ToneWriter {
    
    var useSimpsons: Bool
    
    let kAudioWriterExpectsMediaDataInRealTime = false
    let kToneGeneratorQueue = "com.limit-point.tone-generator-queue"
    
    var scale: ((Double)->Double)? // scale factor range in [0,1]
    
    var integral:[Double] = []
    
    init(useSimpsons: Bool) {
        self.useSimpsons = useSimpsons
    }
    
    deinit {
        print("ToneWriter deinit")
    }
    
    func audioSamplesForRange(sampleRate:Int, sampleRange:ClosedRange<Int>) -> [Int16] {
        
        var samples:[Int16] = []
        
        let delta_t:Double = 1.0 / Double(sampleRate)
        
        for i in sampleRange.lowerBound...sampleRange.upperBound {
            let t = Double(i) * delta_t
            
            let x = integral[i]
            var value = sin(2 * .pi * x) * Double(Int16.max)
            if let scale = scale {
                value = scale(t) * value
            }
            let valueInt16 = Int16(max(min(value, Double(Int16.max)), Double(Int16.min)))
            samples.append(valueInt16)
        }
        
        return samples
    }
    
    func rangeForIndex(bufferIndex:Int, bufferSize:Int, samplesRemaining:Int?) -> ClosedRange<Int> {
        let start = bufferIndex * bufferSize
        
        if let samplesRemaining = samplesRemaining {
            return start...(start + samplesRemaining - 1)
        }
        
        return start...(start + bufferSize - 1)
    }
    
    func sampleBufferForSamples(audioSamples:[Int16], bufferIndex:Int, sampleRate:Int, bufferSize:Int) -> CMSampleBuffer? {
        
        var sampleBuffer:CMSampleBuffer?
        
        let bytesInt16 = MemoryLayout<Int16>.stride
        let dataSize = audioSamples.count * bytesInt16
        
        var samplesBlock:CMBlockBuffer? 
        
        let memoryBlock:UnsafeMutableRawPointer = UnsafeMutableRawPointer.allocate(
            byteCount: dataSize,
            alignment: MemoryLayout<Int16>.alignment)
        
        let _ = audioSamples.withUnsafeBufferPointer { buffer in
            memoryBlock.initializeMemory(as: Int16.self, from: buffer.baseAddress!, count: buffer.count)
        }
        
        if CMBlockBufferCreateWithMemoryBlock(
            allocator: kCFAllocatorDefault, 
            memoryBlock: memoryBlock, 
            blockLength: dataSize, 
            blockAllocator: nil, 
            customBlockSource: nil, 
            offsetToData: 0, 
            dataLength: dataSize, 
            flags: 0, 
            blockBufferOut:&samplesBlock
        ) == kCMBlockBufferNoErr, let samplesBlock = samplesBlock {
            
            var asbd = AudioStreamBasicDescription()
            asbd.mSampleRate = Float64(sampleRate)
            asbd.mFormatID = kAudioFormatLinearPCM
            asbd.mFormatFlags = kAudioFormatFlagIsSignedInteger | kAudioFormatFlagIsPacked
            asbd.mBitsPerChannel = 16
            asbd.mChannelsPerFrame = 1
            asbd.mFramesPerPacket = 1
            asbd.mBytesPerFrame = 2
            asbd.mBytesPerPacket = 2
            
            var formatDesc: CMAudioFormatDescription?
            
            let sampleDuration = CMTimeMakeWithSeconds((1.0 / Float64(sampleRate)), preferredTimescale: Int32.max)
            
            if CMAudioFormatDescriptionCreate(allocator: nil, asbd: &asbd, layoutSize: 0, layout: nil, magicCookieSize: 0, magicCookie: nil, extensions: nil, formatDescriptionOut: &formatDesc) == noErr, let formatDesc = formatDesc {
                
                let sampleTime = CMTimeMultiply(sampleDuration, multiplier: Int32(bufferIndex * bufferSize))
                
                let timingInfo = CMSampleTimingInfo(duration: sampleDuration, presentationTimeStamp: sampleTime, decodeTimeStamp: .invalid)
                
                if CMSampleBufferCreate(allocator: kCFAllocatorDefault, dataBuffer: samplesBlock, dataReady: true, makeDataReadyCallback: nil, refcon: nil, formatDescription: formatDesc, sampleCount: audioSamples.count, sampleTimingEntryCount: 1, sampleTimingArray: [timingInfo], sampleSizeEntryCount: 0, sampleSizeArray: nil, sampleBufferOut: &sampleBuffer) == noErr, let sampleBuffer = sampleBuffer {
                    
                    guard sampleBuffer.isValid, sampleBuffer.numSamples == audioSamples.count else {
                        return nil
                    }
                }
            }
        }
        
        return sampleBuffer
    }
    
    func sampleBufferForComponent(sampleRate:Int, bufferSize: Int, bufferIndex:Int, samplesRemaining:Int?) -> CMSampleBuffer? {
        
        let audioSamples = audioSamplesForRange(sampleRate: sampleRate, sampleRange: rangeForIndex(bufferIndex:bufferIndex, bufferSize: bufferSize, samplesRemaining: samplesRemaining))
        
        return sampleBufferForSamples(audioSamples: audioSamples, bufferIndex: bufferIndex, sampleRate: sampleRate, bufferSize: bufferSize)
    }
    
    func saveComponentSamplesToFile(userIFCurve:[CGPoint], duration:Double = 3, sampleRate:Int = 44100, bufferSize:Int = 8192, destinationURL:URL, completion: @escaping (URL?, String?) -> ())  {
        
        
        let sampleCount = Int(duration * Double(sampleRate))
        let stepSize = 1.0 / Double(sampleRate)
        
        let piecewise_integrator = PiecewiseIntegrator(userIFCurve: userIFCurve, sampleCount: sampleCount, rangeCount: bufferSize, delta: stepSize, useSimpsons: useSimpsons)
        
        while let piecewiseIntegral = piecewise_integrator.nextIntegral() {
            integral = integral + piecewiseIntegral
        }
        
        var nbrSampleBuffers = Int(duration * Double(sampleRate)) / bufferSize
        
        let samplesRemaining = Int(duration * Double(sampleRate)) % bufferSize
        
        if samplesRemaining > 0 {
            nbrSampleBuffers += 1
        }
        
        print(nbrSampleBuffers)
        print(piecewise_integrator.partitions.count)
        
        print("samplesRemaining = \(samplesRemaining)")
        
        
        guard let sampleBuffer = sampleBufferForComponent(sampleRate: sampleRate, bufferSize:  bufferSize, bufferIndex: 0, samplesRemaining: nil) else {
            completion(nil, "Invalid first sample buffer.")
            return
        }
        
        var actualDestinationURL = destinationURL
        
        if actualDestinationURL.pathExtension != "wav" {
            actualDestinationURL.deletePathExtension()
            actualDestinationURL.appendPathExtension("wav")
        }
        
        try? FileManager.default.removeItem(at: actualDestinationURL)
        
        guard let assetWriter = try? AVAssetWriter(outputURL: actualDestinationURL, fileType: AVFileType.wav) else {
            completion(nil, "Can't create asset writer.")
            return
        }
        
        let sourceFormat = CMSampleBufferGetFormatDescription(sampleBuffer)
        
        let audioCompressionSettings = [AVFormatIDKey: kAudioFormatLinearPCM] as [String : Any]
        
        if assetWriter.canApply(outputSettings: audioCompressionSettings, forMediaType: AVMediaType.audio) == false {
            completion(nil, "Can't apply compression settings to asset writer.")
            return
        }
        
        let audioWriterInput = AVAssetWriterInput(mediaType: AVMediaType.audio, outputSettings:audioCompressionSettings, sourceFormatHint: sourceFormat)
        
        audioWriterInput.expectsMediaDataInRealTime = kAudioWriterExpectsMediaDataInRealTime
        
        if assetWriter.canAdd(audioWriterInput) {
            assetWriter.add(audioWriterInput)
            
        } else {
            completion(nil, "Can't add audio input to asset writer.")
            return
        }
        
        let serialQueue: DispatchQueue = DispatchQueue(label: kToneGeneratorQueue)
        
        assetWriter.startWriting()
        assetWriter.startSession(atSourceTime: CMTime.zero)
        
        func finishWriting() {
            assetWriter.finishWriting {
                switch assetWriter.status {
                    case .failed:
                        
                        var errorMessage = ""
                        if let error = assetWriter.error {
                            
                            let nserr = error as NSError
                            
                            let description = nserr.localizedDescription
                            errorMessage = description
                            
                            if let failureReason = nserr.localizedFailureReason {
                                print("error = \(failureReason)")
                                errorMessage += ("Reason " + failureReason)
                            }
                        }
                        completion(nil, errorMessage)
                        print("saveComponentsSamplesToFile errorMessage = \(errorMessage)")
                        return
                    case .completed:
                        print("saveComponentsSamplesToFile completed : \(actualDestinationURL)")
                        completion(actualDestinationURL, nil)
                        return
                    default:
                        print("saveComponentsSamplesToFile other failure?")
                        completion(nil, nil)
                        return
                }
            }
        }
        
        var bufferIndex = 0
        
        audioWriterInput.requestMediaDataWhenReady(on: serialQueue) { [weak self] in
            
            while audioWriterInput.isReadyForMoreMediaData, bufferIndex < nbrSampleBuffers {
                
                var currentSampleBuffer:CMSampleBuffer?
                
                if samplesRemaining > 0 {
                    if bufferIndex < nbrSampleBuffers-1 {
                        currentSampleBuffer = self?.sampleBufferForComponent(sampleRate: sampleRate, bufferSize: bufferSize, bufferIndex: bufferIndex, samplesRemaining: nil)
                    }
                    else {
                        currentSampleBuffer = self?.sampleBufferForComponent(sampleRate: sampleRate, bufferSize: bufferSize, bufferIndex: bufferIndex, samplesRemaining: samplesRemaining)
                    }
                }
                else {
                    currentSampleBuffer = self?.sampleBufferForComponent(sampleRate: sampleRate, bufferSize: bufferSize, bufferIndex: bufferIndex, samplesRemaining: nil)
                }
                
                if let currentSampleBuffer = currentSampleBuffer {
                    audioWriterInput.append(currentSampleBuffer)
                }
                
                bufferIndex += 1
                
                if bufferIndex == nbrSampleBuffers {
                    audioWriterInput.markAsFinished()
                    finishWriting()
                }
            }
        }
    }
}

var toneWriter = ToneWriter(useSimpsons: kUseSimpsons)

func GenerateToneShapeAudio() {
    
    let duration:Double = 1.0
    let scale:((Double)->Double) = {t in 1 - (t / duration)}
    
        // equally spaced
    let N:Double = 6
    let DF:Double = N/6
    let x:[Double] = [0 * DF, 1 * DF, 2 * DF, 3 * DF, 4 * DF, 5 * DF, 6 * DF]
    let y:[Double] = [30,440,50,440,50,440,50]
    
    let userIFCurve:[CGPoint] = [CGPoint(x: x[0], y: y[0]), CGPoint(x: x[1], y: y[1]), CGPoint(x: x[2], y: y[2]), CGPoint(x: x[3], y: y[3]), CGPoint(x: x[4], y: y[4]), CGPoint(x: x[5], y: y[5]), CGPoint(x: x[6], y: y[6])]
    
    var audioFileURL:URL
    
    do {
        let documentsDirectory = try FileManager.default.url(for: .documentDirectory, in: .userDomainMask, appropriateFor: nil, create: false)
        audioFileURL = documentsDirectory.appendingPathComponent("ToneShape.wav")
    } catch {
        return 
    }
    
    toneWriter.scale = scale
    
    toneWriter.saveComponentSamplesToFile(userIFCurve: userIFCurve, duration: duration,  destinationURL: audioFileURL) { resultURL, message in
        if let resultURL = resultURL {
            let asset = AVAsset(url: resultURL)
            Task {
                do {
                    let duration = try await asset.load(.duration)
                    print("ToneWriter : audio duration = \(duration.seconds)")
                }
                catch {
                }
            }
        }
        else {
            print("An error occurred : \(message ?? "No error message available.")")
        }
    }
}
