import Foundation
import Accelerate
import OnnxRuntimeBindings

// MARK: - Configuration Structures

struct Config: Codable {
    struct AEConfig: Codable {
        let sample_rate: Int
        let base_chunk_size: Int
    }
    
    struct TTLConfig: Codable {
        let chunk_compress_factor: Int
        let latent_dim: Int
    }
    
    let ae: AEConfig
    let ttl: TTLConfig
}

// MARK: - Voice Style Data Structure

struct VoiceStyleData: Codable {
    struct StyleComponent: Codable {
        let data: [[[Float]]]
        let dims: [Int]
        let type: String
    }
    
    let style_ttl: StyleComponent
    let style_dp: StyleComponent
}

// MARK: - Unicode Text Processor

class UnicodeProcessor {
    let indexer: [Int64]
    
    init(unicodeIndexerPath: String) throws {
        let data = try Data(contentsOf: URL(fileURLWithPath: unicodeIndexerPath))
        self.indexer = try JSONDecoder().decode([Int64].self, from: data)
    }
    
    func call(_ textList: [String]) -> (textIds: [[Int64]], textMask: [[[Float]]]) {
        let processedTexts = textList.map { preprocessText($0) }
        
        var textIdsLengths = [Int]()
        for text in processedTexts {
            textIdsLengths.append(text.count)
        }
        
        let maxLen = textIdsLengths.max() ?? 0
        
        var textIds = [[Int64]]()
        for text in processedTexts {
            var row = Array(repeating: Int64(0), count: maxLen)
            let unicodeValues = Array(text.unicodeScalars.map { Int($0.value) })
            for (j, val) in unicodeValues.enumerated() {
                if val < indexer.count {
                    row[j] = indexer[val]
                } else {
                    row[j] = -1
                }
            }
            textIds.append(row)
        }
        
        let textMask = getTextMask(textIdsLengths)
        return (textIds, textMask)
    }
}

func preprocessText(_ text: String) -> String {
    return text.precomposedStringWithCompatibilityMapping
}

func lengthToMask(_ lengths: [Int], maxLen: Int? = nil) -> [[[Float]]] {
    let actualMaxLen = maxLen ?? (lengths.max() ?? 0)
    
    var mask = [[[Float]]]()
    for len in lengths {
        var row = Array(repeating: Float(0.0), count: actualMaxLen)
        for j in 0..<min(len, actualMaxLen) {
            row[j] = 1.0
        }
        mask.append([row])
    }
    return mask
}

func getTextMask(_ textIdsLengths: [Int]) -> [[[Float]]] {
    let maxLen = textIdsLengths.max() ?? 0
    return lengthToMask(textIdsLengths, maxLen: maxLen)
}

func sampleNoisyLatent(duration: [Float], sampleRate: Int, baseChunkSize: Int, chunkCompress: Int, latentDim: Int) -> (noisyLatent: [[[Float]]], latentMask: [[[Float]]]) {
    let bsz = duration.count
    let maxDur = duration.max() ?? 0.0
    
    let wavLenMax = Int(maxDur * Float(sampleRate))
    var wavLengths = [Int]()
    for d in duration {
        wavLengths.append(Int(d * Float(sampleRate)))
    }
    
    let chunkSize = baseChunkSize * chunkCompress
    let latentLen = (wavLenMax + chunkSize - 1) / chunkSize
    let latentDimVal = latentDim * chunkCompress
    
    var noisyLatent = [[[Float]]]()
    for _ in 0..<bsz {
        var batch = [[Float]]()
        for _ in 0..<latentDimVal {
            var row = [Float]()
            for _ in 0..<latentLen {
                // Box-Muller transform
                let u1 = Float.random(in: 0.0001...1.0)
                let u2 = Float.random(in: 0.0...1.0)
                let val = sqrt(-2.0 * log(u1)) * cos(2.0 * Float.pi * u2)
                row.append(val)
            }
            batch.append(row)
        }
        noisyLatent.append(batch)
    }
    
    var latentLengths = [Int]()
    for len in wavLengths {
        latentLengths.append((len + chunkSize - 1) / chunkSize)
    }
    
    let latentMask = lengthToMask(latentLengths, maxLen: latentLen)
    
    // Apply mask
    for b in 0..<bsz {
        for d in 0..<latentDimVal {
            for t in 0..<latentLen {
                noisyLatent[b][d][t] *= latentMask[b][0][t]
            }
        }
    }
    
    return (noisyLatent, latentMask)
}

func getLatentMask(_ wavLengths: [Int64], _ cfgs: Config) -> [[[Float]]] {
    let baseChunkSize = cfgs.ae.base_chunk_size
    let chunkCompressFactor = cfgs.ttl.chunk_compress_factor
    let latentSize = baseChunkSize * chunkCompressFactor
    
    var latentLengths = [Int]()
    for len in wavLengths {
        latentLengths.append((Int(len) + latentSize - 1) / latentSize)
    }
    
    let maxLen = latentLengths.max() ?? 0
    return lengthToMask(latentLengths, maxLen: maxLen)
}

// MARK: - WAV File I/O

func writeWavFile(_ filename: String, _ audioData: [Float], _ sampleRate: Int) throws {
    let url = URL(fileURLWithPath: filename)
    
    // Convert float to int16
    let int16Data = audioData.map { sample -> Int16 in
        let clamped = max(-1.0, min(1.0, sample))
        return Int16(clamped * 32767.0)
    }
    
    // Create WAV header
    let numChannels: UInt16 = 1
    let bitsPerSample: UInt16 = 16
    let byteRate = UInt32(sampleRate) * UInt32(numChannels) * UInt32(bitsPerSample) / 8
    let blockAlign = numChannels * bitsPerSample / 8
    let dataSize = UInt32(int16Data.count * 2)
    
    var data = Data()
    
    // RIFF chunk
    data.append("RIFF".data(using: .ascii)!)
    withUnsafeBytes(of: UInt32(36 + dataSize).littleEndian) { data.append(contentsOf: $0) }
    data.append("WAVE".data(using: .ascii)!)
    
    // fmt chunk
    data.append("fmt ".data(using: .ascii)!)
    withUnsafeBytes(of: UInt32(16).littleEndian) { data.append(contentsOf: $0) }
    withUnsafeBytes(of: UInt16(1).littleEndian) { data.append(contentsOf: $0) } // PCM
    withUnsafeBytes(of: numChannels.littleEndian) { data.append(contentsOf: $0) }
    withUnsafeBytes(of: UInt32(sampleRate).littleEndian) { data.append(contentsOf: $0) }
    withUnsafeBytes(of: byteRate.littleEndian) { data.append(contentsOf: $0) }
    withUnsafeBytes(of: blockAlign.littleEndian) { data.append(contentsOf: $0) }
    withUnsafeBytes(of: bitsPerSample.littleEndian) { data.append(contentsOf: $0) }
    
    // data chunk
    data.append("data".data(using: .ascii)!)
    withUnsafeBytes(of: dataSize.littleEndian) { data.append(contentsOf: $0) }
    
    // audio data
    int16Data.withUnsafeBytes { data.append(contentsOf: $0) }
    
    try data.write(to: url)
}

// MARK: - Utility Functions

func timer<T>(_ name: String, _ f: () throws -> T) rethrows -> T {
    let start = Date()
    print("\(name)...")
    let result = try f()
    let elapsed = Date().timeIntervalSince(start)
    print(String(format: "  -> %@ completed in %.2f sec", name, elapsed))
    return result
}

func sanitizeFilename(_ text: String, maxLen: Int) -> String {
    let truncated = text.count > maxLen ? String(text.prefix(maxLen)) : text
    return truncated.map { char in
        if char.isLetter || char.isNumber {
            return char
        } else {
            return Character("_")
        }
    }.map(String.init).joined()
}

func loadCfgs(_ onnxDir: String) throws -> Config {
    let cfgPath = "\(onnxDir)/tts.json"
    let data = try Data(contentsOf: URL(fileURLWithPath: cfgPath))
    let config = try JSONDecoder().decode(Config.self, from: data)
    return config
}

// MARK: - ONNX Runtime Integration

struct Style {
    let ttl: ORTValue
    let dp: ORTValue
}

class TextToSpeech {
    let cfgs: Config
    let textProcessor: UnicodeProcessor
    let dpOrt: ORTSession
    let textEncOrt: ORTSession
    let vectorEstOrt: ORTSession
    let vocoderOrt: ORTSession
    let sampleRate: Int
    
    init(cfgs: Config, textProcessor: UnicodeProcessor,
         dpOrt: ORTSession, textEncOrt: ORTSession,
         vectorEstOrt: ORTSession, vocoderOrt: ORTSession) {
        self.cfgs = cfgs
        self.textProcessor = textProcessor
        self.dpOrt = dpOrt
        self.textEncOrt = textEncOrt
        self.vectorEstOrt = vectorEstOrt
        self.vocoderOrt = vocoderOrt
        self.sampleRate = cfgs.ae.sample_rate
    }
    
    func call(_ textList: [String], _ style: Style, _ totalStep: Int) throws -> (wav: [Float], duration: [Float]) {
        let bsz = textList.count
        
        // Process text
        let (textIds, textMask) = textProcessor.call(textList)
        
        // Flatten text IDs
        let textIdsFlat = textIds.flatMap { $0 }
        let textIdsShape: [NSNumber] = [NSNumber(value: bsz), NSNumber(value: textIds[0].count)]
        let textIdsValue = try ORTValue(tensorData: NSMutableData(bytes: textIdsFlat, length: textIdsFlat.count * MemoryLayout<Int64>.size),
                                        elementType: .int64,
                                        shape: textIdsShape)
        
        // Flatten text mask
        let textMaskFlat = textMask.flatMap { $0.flatMap { $0 } }
        let textMaskShape: [NSNumber] = [NSNumber(value: bsz), 1, NSNumber(value: textMask[0][0].count)]
        let textMaskValue = try ORTValue(tensorData: NSMutableData(bytes: textMaskFlat, length: textMaskFlat.count * MemoryLayout<Float>.size),
                                         elementType: .float,
                                         shape: textMaskShape)
        
        // Predict duration
        let dpOutputs = try dpOrt.run(withInputs: ["text_ids": textIdsValue, "style_dp": style.dp, "text_mask": textMaskValue],
                                      outputNames: ["duration"],
                                      runOptions: nil)
        
        let durationData = try dpOutputs["duration"]!.tensorData() as Data
        let duration = durationData.withUnsafeBytes { ptr in
            Array(ptr.bindMemory(to: Float.self))
        }
        
        // Encode text
        let textEncOutputs = try textEncOrt.run(withInputs: ["text_ids": textIdsValue, "style_ttl": style.ttl, "text_mask": textMaskValue],
                                                outputNames: ["text_emb"],
                                                runOptions: nil)
        
        let textEmbValue = textEncOutputs["text_emb"]!
        
        // Sample noisy latent
        var (xt, latentMask) = sampleNoisyLatent(duration: duration, sampleRate: sampleRate,
                                                  baseChunkSize: cfgs.ae.base_chunk_size,
                                                  chunkCompress: cfgs.ttl.chunk_compress_factor,
                                                  latentDim: cfgs.ttl.latent_dim)
        
        // Prepare constant arrays
        let totalStepArray = Array(repeating: Float(totalStep), count: bsz)
        let totalStepValue = try ORTValue(tensorData: NSMutableData(bytes: totalStepArray, length: totalStepArray.count * MemoryLayout<Float>.size),
                                          elementType: .float,
                                          shape: [NSNumber(value: bsz)])
        
        // Denoising loop
        for step in 0..<totalStep {
            let currentStepArray = Array(repeating: Float(step), count: bsz)
            let currentStepValue = try ORTValue(tensorData: NSMutableData(bytes: currentStepArray, length: currentStepArray.count * MemoryLayout<Float>.size),
                                                elementType: .float,
                                                shape: [NSNumber(value: bsz)])
            
            // Flatten xt
            let xtFlat = xt.flatMap { $0.flatMap { $0 } }
            let xtShape: [NSNumber] = [NSNumber(value: bsz), NSNumber(value: xt[0].count), NSNumber(value: xt[0][0].count)]
            let xtValue = try ORTValue(tensorData: NSMutableData(bytes: xtFlat, length: xtFlat.count * MemoryLayout<Float>.size),
                                       elementType: .float,
                                       shape: xtShape)
            
            // Flatten latent mask
            let latentMaskFlat = latentMask.flatMap { $0.flatMap { $0 } }
            let latentMaskShape: [NSNumber] = [NSNumber(value: bsz), 1, NSNumber(value: latentMask[0][0].count)]
            let latentMaskValue = try ORTValue(tensorData: NSMutableData(bytes: latentMaskFlat, length: latentMaskFlat.count * MemoryLayout<Float>.size),
                                               elementType: .float,
                                               shape: latentMaskShape)
            
            let vectorEstOutputs = try vectorEstOrt.run(withInputs: [
                "noisy_latent": xtValue,
                "text_emb": textEmbValue,
                "style_ttl": style.ttl,
                "latent_mask": latentMaskValue,
                "text_mask": textMaskValue,
                "current_step": currentStepValue,
                "total_step": totalStepValue
            ], outputNames: ["denoised_latent"], runOptions: nil)
            
            let denoisedData = try vectorEstOutputs["denoised_latent"]!.tensorData() as Data
            let denoisedFlat = denoisedData.withUnsafeBytes { ptr in
                Array(ptr.bindMemory(to: Float.self))
            }
            
            // Reshape to 3D
            let latentDimVal = xt[0].count
            let latentLen = xt[0][0].count
            xt = []
            var idx = 0
            for _ in 0..<bsz {
                var batch = [[Float]]()
                for _ in 0..<latentDimVal {
                    var row = [Float]()
                    for _ in 0..<latentLen {
                        row.append(denoisedFlat[idx])
                        idx += 1
                    }
                    batch.append(row)
                }
                xt.append(batch)
            }
        }
        
        // Generate waveform
        let finalXtFlat = xt.flatMap { $0.flatMap { $0 } }
        let finalXtShape: [NSNumber] = [NSNumber(value: bsz), NSNumber(value: xt[0].count), NSNumber(value: xt[0][0].count)]
        let finalXtValue = try ORTValue(tensorData: NSMutableData(bytes: finalXtFlat, length: finalXtFlat.count * MemoryLayout<Float>.size),
                                        elementType: .float,
                                        shape: finalXtShape)
        
        let vocoderOutputs = try vocoderOrt.run(withInputs: ["latent": finalXtValue],
                                                outputNames: ["wav_tts"],
                                                runOptions: nil)
        
        let wavData = try vocoderOutputs["wav_tts"]!.tensorData() as Data
        let wav = wavData.withUnsafeBytes { ptr in
            Array(ptr.bindMemory(to: Float.self))
        }
        
        return (wav, duration)
    }
}

// MARK: - Component Loading Functions

func loadVoiceStyle(_ voiceStylePaths: [String], verbose: Bool) throws -> Style {
    let bsz = voiceStylePaths.count
    
    // Read first file to get dimensions
    let firstData = try Data(contentsOf: URL(fileURLWithPath: voiceStylePaths[0]))
    let firstStyle = try JSONDecoder().decode(VoiceStyleData.self, from: firstData)
    
    let ttlDims = firstStyle.style_ttl.dims
    let dpDims = firstStyle.style_dp.dims
    
    let ttlDim1 = ttlDims[1]
    let ttlDim2 = ttlDims[2]
    let dpDim1 = dpDims[1]
    let dpDim2 = dpDims[2]
    
    // Pre-allocate arrays with full batch size
    let ttlSize = bsz * ttlDim1 * ttlDim2
    let dpSize = bsz * dpDim1 * dpDim2
    var ttlFlat = [Float](repeating: 0.0, count: ttlSize)
    var dpFlat = [Float](repeating: 0.0, count: dpSize)
    
    // Fill in the data
    for (i, path) in voiceStylePaths.enumerated() {
        let data = try Data(contentsOf: URL(fileURLWithPath: path))
        let voiceStyle = try JSONDecoder().decode(VoiceStyleData.self, from: data)
        
        // Flatten TTL data
        let ttlOffset = i * ttlDim1 * ttlDim2
        var idx = 0
        for batch in voiceStyle.style_ttl.data {
            for row in batch {
                for val in row {
                    ttlFlat[ttlOffset + idx] = val
                    idx += 1
                }
            }
        }
        
        // Flatten DP data
        let dpOffset = i * dpDim1 * dpDim2
        idx = 0
        for batch in voiceStyle.style_dp.data {
            for row in batch {
                for val in row {
                    dpFlat[dpOffset + idx] = val
                    idx += 1
                }
            }
        }
    }
    
    let ttlShape: [NSNumber] = [NSNumber(value: bsz), NSNumber(value: ttlDim1), NSNumber(value: ttlDim2)]
    let dpShape: [NSNumber] = [NSNumber(value: bsz), NSNumber(value: dpDim1), NSNumber(value: dpDim2)]
    
    let ttlValue = try ORTValue(tensorData: NSMutableData(bytes: &ttlFlat, length: ttlFlat.count * MemoryLayout<Float>.size),
                                elementType: .float,
                                shape: ttlShape)
    let dpValue = try ORTValue(tensorData: NSMutableData(bytes: &dpFlat, length: dpFlat.count * MemoryLayout<Float>.size),
                               elementType: .float,
                               shape: dpShape)
    
    if verbose {
        print("Loaded \(bsz) voice styles\n")
    }
    
    return Style(ttl: ttlValue, dp: dpValue)
}

func loadTextToSpeech(_ onnxDir: String, _ useGpu: Bool, _ env: ORTEnv) throws -> TextToSpeech {
    if useGpu {
        throw NSError(domain: "TTS", code: 1, userInfo: [NSLocalizedDescriptionKey: "GPU mode is not supported yet"])
    }
    print("Using CPU for inference\n")
    
    let cfgs = try loadCfgs(onnxDir)
    
    let sessionOptions = try ORTSessionOptions()
    
    let dpPath = "\(onnxDir)/duration_predictor.onnx"
    let textEncPath = "\(onnxDir)/text_encoder.onnx"
    let vectorEstPath = "\(onnxDir)/vector_estimator.onnx"
    let vocoderPath = "\(onnxDir)/vocoder.onnx"
    
    let dpOrt = try ORTSession(env: env, modelPath: dpPath, sessionOptions: sessionOptions)
    let textEncOrt = try ORTSession(env: env, modelPath: textEncPath, sessionOptions: sessionOptions)
    let vectorEstOrt = try ORTSession(env: env, modelPath: vectorEstPath, sessionOptions: sessionOptions)
    let vocoderOrt = try ORTSession(env: env, modelPath: vocoderPath, sessionOptions: sessionOptions)
    
    let unicodeIndexerPath = "\(onnxDir)/unicode_indexer.json"
    let textProcessor = try UnicodeProcessor(unicodeIndexerPath: unicodeIndexerPath)
    
    return TextToSpeech(cfgs: cfgs, textProcessor: textProcessor,
                       dpOrt: dpOrt, textEncOrt: textEncOrt,
                       vectorEstOrt: vectorEstOrt, vocoderOrt: vocoderOrt)
}