import Foundation

// ─────────────────────────────────────────────────────────────
// Float16 Weight Conversion Tool
// ─────────────────────────────────────────────────────────────

public struct Float16Converter {
    /// Convert Float32 safetensors to Float16
    public static func convertModel(
        from sourceDir: String,
        to targetDir: String
    ) throws {
        print("Converting model from Float32 to Float16...")
        print("Source: \(sourceDir)")
        print("Target: \(targetDir)")
        
        // Create target directory
        try FileManager.default.createDirectory(
            at: URL(fileURLWithPath: targetDir),
            withIntermediateDirectories: true
        )
        
        // Copy config files
        let configFiles = ["config.json", "tokenizer.json", "tokenizer.model"]
        for file in configFiles {
            let source = (sourceDir as NSString).appendingPathComponent(file)
            let target = (targetDir as NSString).appendingPathComponent(file)
            
            if FileManager.default.fileExists(atPath: source) {
                try FileManager.default.copyItem(atPath: source, toPath: target)
                print("✓ Copied \(file)")
            }
        }
        
        // Convert safetensors
        let sourceFile = (sourceDir as NSString).appendingPathComponent("model.safetensors")
        if FileManager.default.fileExists(atPath: sourceFile) {
            try convertSafetensors(from: sourceFile, to: (targetDir as NSString).appendingPathComponent("model.safetensors"))
        }
        
        print("✓ Conversion complete!")
    }
    
    /// Convert single safetensors file
    private static func convertSafetensors(from source: String, to target: String) throws {
        print("Converting \(source)...")
        
        // Load safetensors
        let data = try Data(contentsOf: URL(fileURLWithPath: source))
        let headerSize = data.prefix(8).withUnsafeBytes { $0.load(as: UInt64.self) }
        let header = String(data: data.subdata(in: 8..<Int(headerSize) + 8), encoding: .utf8)!
        
        // Parse header
        guard let headerJson = try JSONSerialization.jsonObject(with: Data(header.utf8)) as? [String: Any] else {
            throw ConversionError.invalidHeader
        }
        
        // Convert tensors
        var newHeader: [String: Any] = [:]
        var newData = Data()
        
        let headerSizeInt = Int(headerSize)
        var offset = headerSizeInt + 8
        
        for (name, info) in headerJson {
            guard let tensorInfo = info as? [String: Any],
                  let dtype = tensorInfo["dtype"] as? String,
                  let shape = tensorInfo["shape"] as? [Int],
                  let dataOffsets = tensorInfo["data_offsets"] as? [Int] else {
                continue
            }
            
            // Only convert Float32 tensors
            if dtype == "F32" {
                let start = dataOffsets[0]
                let end = dataOffsets[1]
                let tensorData = data.subdata(in: offset + start..<offset + end)
                
                // Convert to Float16
                let floatValues = tensorData.withUnsafeBytes { Array($0.bindMemory(to: Float.self)) }
                let halfValues = floatValues.map { Float16($0) }
                let halfData = halfValues.withUnsafeBytes { Data($0) }
                
                // Update header
                var newTensorInfo = tensorInfo
                newTensorInfo["dtype"] = "F16"
                newTensorInfo["data_offsets"] = [newData.count, newData.count + halfData.count]
                newHeader[name] = newTensorInfo
                
                // Append data
                newData.append(halfData)
                
                print("  ✓ Converted \(name) (\(floatValues.count) F32 → \(halfValues.count) F16)")
            } else {
                // Keep other dtypes as-is
                let start = dataOffsets[0]
                let end = dataOffsets[1]
                let tensorData = data.subdata(in: offset + start..<offset + end)
                
                newHeader[name] = tensorInfo
                
                // Update offsets
                if let offsets = newHeader[name] as? [String: Any] {
                    var newOffsets = offsets
                    newOffsets["data_offsets"] = [newData.count, newData.count + tensorData.count]
                    newHeader[name] = newOffsets
                }
                
                newData.append(tensorData)
            }
        }
        
        // Write new safetensors
        let newHeaderJson = try JSONSerialization.data(withJSONObject: newHeader)
        var outputData = Data()
        
        // Write header size
        var newHeaderSize = UInt64(newHeaderJson.count)
        outputData.append(Data(bytes: &newHeaderSize, count: 8))
        
        // Write header
        outputData.append(newHeaderJson)
        
        // Write tensor data
        outputData.append(newData)
        
        try outputData.write(to: URL(fileURLWithPath: target))
        print("  ✓ Saved to \(target)")
    }
}

/// Conversion errors
public enum ConversionError: Error, LocalizedError {
    case invalidHeader
    case invalidTensor
    case conversionFailed(String)
    
    public var errorDescription: String? {
        switch self {
        case .invalidHeader:
            return "Invalid safetensors header"
        case .invalidTensor:
            return "Invalid tensor data"
        case .conversionFailed(let detail):
            return "Conversion failed: \(detail)"
        }
    }
}