import Foundation
import Accelerate
import OnnxRuntimeBindings

// MARK: - Available Languages

let AVAILABLE_LANGS = ["en", "ko", "es", "pt", "fr"]

func isValidLang(_ lang: String) -> Bool {
    return AVAILABLE_LANGS.contains(lang)
}

// 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], _ langList: [String]) -> (textIds: [[Int64]], textMask: [[[Float]]]) {
        var processedTexts = [String]()
        for (i, text) in textList.enumerated() {
            processedTexts.append(preprocessText(text, lang: langList[i]))
        }
        
        // Use unicodeScalars.count for correct length after NFKD decomposition
        var textIdsLengths = [Int]()
        for text in processedTexts {
            textIdsLengths.append(text.unicodeScalars.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, lang: String) -> String {
    // Use NFKD (decomposed) for proper Hangul Jamo decomposition
    var text = text.decomposedStringWithCompatibilityMapping

    // Remove emojis (wide Unicode range)
    // Swift NSRegularExpression doesn't support Unicode escapes above \uFFFF
    // Use character filtering instead
    text = text.unicodeScalars.filter { scalar in
        let value = scalar.value
        return !((value >= 0x1F600 && value <= 0x1F64F) ||
                 (value >= 0x1F300 && value <= 0x1F5FF) ||
                 (value >= 0x1F680 && value <= 0x1F6FF) ||
                 (value >= 0x1F700 && value <= 0x1F77F) ||
                 (value >= 0x1F780 && value <= 0x1F7FF) ||
                 (value >= 0x1F800 && value <= 0x1F8FF) ||
                 (value >= 0x1F900 && value <= 0x1F9FF) ||
                 (value >= 0x1FA00 && value <= 0x1FA6F) ||
                 (value >= 0x1FA70 && value <= 0x1FAFF) ||
                 (value >= 0x2600 && value <= 0x26FF) ||
                 (value >= 0x2700 && value <= 0x27BF) ||
                 (value >= 0x1F1E6 && value <= 0x1F1FF))
    }.map { String($0) }.joined()

    // Replace various dashes and symbols
    let replacements: [String: String] = [
        "–": "-",      // en dash
        "‑": "-",      // non-breaking hyphen
        "—": "-",      // em dash
        "_": " ",      // underscore
        "\u{201C}": "\"",     // left double quote
        "\u{201D}": "\"",     // right double quote
        "\u{2018}": "'",      // left single quote
        "\u{2019}": "'",      // right single quote
        "´": "'",      // acute accent
        "`": "'",      // grave accent
        "[": " ",      // left bracket
        "]": " ",      // right bracket
        "|": " ",      // vertical bar
        "/": " ",      // slash
        "#": " ",      // hash
        "→": " ",      // right arrow
        "←": " ",      // left arrow
    ]

    for (old, new) in replacements {
        text = text.replacingOccurrences(of: old, with: new)
    }

    // Remove special symbols
    let specialSymbols = ["♥", "☆", "♡", "©", "\\"]
    for symbol in specialSymbols {
        text = text.replacingOccurrences(of: symbol, with: "")
    }

    // Replace known expressions
    let exprReplacements: [String: String] = [
        "@": " at ",
        "e.g.,": "for example, ",
        "i.e.,": "that is, ",
    ]

    for (old, new) in exprReplacements {
        text = text.replacingOccurrences(of: old, with: new)
    }

    // Fix spacing around punctuation
    text = text.replacingOccurrences(of: " ,", with: ",")
    text = text.replacingOccurrences(of: " .", with: ".")
    text = text.replacingOccurrences(of: " !", with: "!")
    text = text.replacingOccurrences(of: " ?", with: "?")
    text = text.replacingOccurrences(of: " ;", with: ";")
    text = text.replacingOccurrences(of: " :", with: ":")
    text = text.replacingOccurrences(of: " '", with: "'")

    // Remove duplicate quotes
    while text.contains("\"\"") {
        text = text.replacingOccurrences(of: "\"\"", with: "\"")
    }
    while text.contains("''") {
        text = text.replacingOccurrences(of: "''", with: "'")
    }
    while text.contains("``") {
        text = text.replacingOccurrences(of: "``", with: "`")
    }

    // Remove extra spaces
    let whitespacePattern = try! NSRegularExpression(pattern: "\\s+")
    let whitespaceRange = NSRange(text.startIndex..., in: text)
    text = whitespacePattern.stringByReplacingMatches(in: text, range: whitespaceRange, withTemplate: " ")
    text = text.trimmingCharacters(in: .whitespacesAndNewlines)

    // If text doesn't end with punctuation, quotes, or closing brackets, add a period
    if !text.isEmpty {
        let punctPattern = try! NSRegularExpression(pattern: "[.!?;:,'\"\\u201C\\u201D\\u2018\\u2019)\\]}…。」』】〉》›»]$")
        let punctRange = NSRange(text.startIndex..., in: text)
        if punctPattern.firstMatch(in: text, range: punctRange) == nil {
            text += "."
        }
    }

    // Validate language
    guard isValidLang(lang) else {
        fatalError("Invalid language: \(lang). Available: \(AVAILABLE_LANGS.joined(separator: ", "))")
    }

    // Wrap text with language tags
    text = "<\(lang)>\(text)</\(lang)>"

    return text
}

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: - Text Chunking

let MAX_CHUNK_LENGTH = 300
let ABBREVIATIONS = [
    "Dr.", "Mr.", "Mrs.", "Ms.", "Prof.", "Sr.", "Jr.",
    "St.", "Ave.", "Rd.", "Blvd.", "Dept.", "Inc.", "Ltd.",
    "Co.", "Corp.", "etc.", "vs.", "i.e.", "e.g.", "Ph.D."
]

func chunkText(_ text: String, maxLen: Int = 0) -> [String] {
    let actualMaxLen = maxLen > 0 ? maxLen : MAX_CHUNK_LENGTH
    let trimmedText = text.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines)
    
    if trimmedText.isEmpty {
        return [""]
    }
    
    // Split by paragraphs using regex
    let paraPattern = try! NSRegularExpression(pattern: "\\n\\s*\\n")
    let paraRange = NSRange(trimmedText.startIndex..., in: trimmedText)
    var paragraphs = [String]()
    var lastEnd = trimmedText.startIndex
    
    paraPattern.enumerateMatches(in: trimmedText, range: paraRange) { match, _, _ in
        if let match = match, let range = Range(match.range, in: trimmedText) {
            paragraphs.append(String(trimmedText[lastEnd..<range.lowerBound]))
            lastEnd = range.upperBound
        }
    }
    if lastEnd < trimmedText.endIndex {
        paragraphs.append(String(trimmedText[lastEnd...]))
    }
    if paragraphs.isEmpty {
        paragraphs = [trimmedText]
    }
    
    var chunks = [String]()
    
    for para in paragraphs {
        let trimmedPara = para.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines)
        if trimmedPara.isEmpty {
            continue
        }
        
        if trimmedPara.count <= actualMaxLen {
            chunks.append(trimmedPara)
            continue
        }
        
        // Split by sentences
        let sentences = splitSentences(trimmedPara)
        var current = ""
        var currentLen = 0
        
        for sentence in sentences {
            let trimmedSentence = sentence.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines)
            if trimmedSentence.isEmpty {
                continue
            }
            
            let sentenceLen = trimmedSentence.count
            if sentenceLen > actualMaxLen {
                // If sentence is longer than maxLen, split by comma or space
                if !current.isEmpty {
                    chunks.append(current.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
                    current = ""
                    currentLen = 0
                }
                
                // Try splitting by comma
                let parts = trimmedSentence.components(separatedBy: ",")
                for part in parts {
                    let trimmedPart = part.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines)
                    if trimmedPart.isEmpty {
                        continue
                    }
                    
                    let partLen = trimmedPart.count
                    if partLen > actualMaxLen {
                        // Split by space as last resort
                        let words = trimmedPart.components(separatedBy: CharacterSet.whitespaces).filter { !$0.isEmpty }
                        var wordChunk = ""
                        var wordChunkLen = 0
                        
                        for word in words {
                            let wordLen = word.count
                            if wordChunkLen + wordLen + 1 > actualMaxLen && !wordChunk.isEmpty {
                                chunks.append(wordChunk.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
                                wordChunk = ""
                                wordChunkLen = 0
                            }
                            
                            if !wordChunk.isEmpty {
                                wordChunk += " "
                                wordChunkLen += 1
                            }
                            wordChunk += word
                            wordChunkLen += wordLen
                        }
                        
                        if !wordChunk.isEmpty {
                            chunks.append(wordChunk.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
                        }
                    } else {
                        if currentLen + partLen + 1 > actualMaxLen && !current.isEmpty {
                            chunks.append(current.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
                            current = ""
                            currentLen = 0
                        }
                        
                        if !current.isEmpty {
                            current += ", "
                            currentLen += 2
                        }
                        current += trimmedPart
                        currentLen += partLen
                    }
                }
                continue
            }
            
            if currentLen + sentenceLen + 1 > actualMaxLen && !current.isEmpty {
                chunks.append(current.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
                current = ""
                currentLen = 0
            }
            
            if !current.isEmpty {
                current += " "
                currentLen += 1
            }
            current += trimmedSentence
            currentLen += sentenceLen
        }
        
        if !current.isEmpty {
            chunks.append(current.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines))
        }
    }
    
    return chunks.isEmpty ? [""] : chunks
}

func splitSentences(_ text: String) -> [String] {
    // Swift's regex doesn't support lookbehind reliably, so we use a simpler approach
    // Split on sentence boundaries and then check if they're abbreviations
    let regex = try! NSRegularExpression(pattern: "([.!?])\\s+")
    let range = NSRange(text.startIndex..., in: text)
    
    // Find all matches
    let matches = regex.matches(in: text, range: range)
    if matches.isEmpty {
        return [text]
    }
    
    var sentences = [String]()
    var lastEnd = text.startIndex
    
    for match in matches {
        guard let matchRange = Range(match.range, in: text) else { continue }
        
        // Get the text before the punctuation
        let beforePunc = String(text[lastEnd..<matchRange.lowerBound])
        
        // Get the punctuation character
        let puncRange = Range(NSRange(location: match.range.location, length: 1), in: text)!
        let punc = String(text[puncRange])
        
        // Check if this ends with an abbreviation
        var isAbbrev = false
        let combined = beforePunc.trimmingCharacters(in: CharacterSet.whitespaces) + punc
        for abbrev in ABBREVIATIONS {
            if combined.hasSuffix(abbrev) {
                isAbbrev = true
                break
            }
        }
        
        if !isAbbrev {
            // This is a real sentence boundary
            sentences.append(String(text[lastEnd..<matchRange.upperBound]))
            lastEnd = matchRange.upperBound
        }
    }
    
    // Add the remaining text
    if lastEnd < text.endIndex {
        sentences.append(String(text[lastEnd...]))
    }
    
    return sentences.isEmpty ? [text] : sentences
}

// 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
    }
    
    private func _infer(_ textList: [String], _ langList: [String], _ style: Style, _ totalStep: Int, speed: Float = 1.05) throws -> (wav: [Float], duration: [Float]) {
        let bsz = textList.count
        
        // Process text
        let (textIds, textMask) = textProcessor.call(textList, langList)
        
        // 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
        var duration = durationData.withUnsafeBytes { ptr in
            Array(ptr.bindMemory(to: Float.self))
        }
        
        // Apply speed factor to duration
        for i in 0..<duration.count {
            duration[i] /= speed
        }
        
        // 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)
    }
    
    func call(_ text: String, _ lang: String, _ style: Style, _ totalStep: Int, speed: Float = 1.05, silenceDuration: Float = 0.3) throws -> (wav: [Float], duration: Float) {
        let maxLen = lang == "ko" ? 120 : 300
        let chunks = chunkText(text, maxLen: maxLen)
        let langList = Array(repeating: lang, count: chunks.count)
        
        var wavCat = [Float]()
        var durCat: Float = 0.0
        
        for (i, chunk) in chunks.enumerated() {
            let result = try _infer([chunk], [langList[i]], style, totalStep, speed: speed)
            
            let dur = result.duration[0]
            let wavLen = Int(Float(sampleRate) * dur)
            let wavChunk = Array(result.wav.prefix(wavLen))
            
            if i == 0 {
                wavCat = wavChunk
                durCat = dur
            } else {
                let silenceLen = Int(silenceDuration * Float(sampleRate))
                let silence = [Float](repeating: 0.0, count: silenceLen)
                
                wavCat.append(contentsOf: silence)
                wavCat.append(contentsOf: wavChunk)
                durCat += silenceDuration + dur
            }
        }
        
        return (wavCat, durCat)
    }
    
    func batch(_ textList: [String], _ langList: [String], _ style: Style, _ totalStep: Int, speed: Float = 1.05) throws -> (wav: [Float], duration: [Float]) {
        return try _infer(textList, langList, style, totalStep, speed: speed)
    }
}

// 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)
}