import Foundation

/// Loads a precomputed RoPE table from an .npy file.
/// Parses the numpy float16 binary format and returns Float32 values.
func loadRopeTable(_ name: String) throws -> [Float] {
    let url = PlapreConfig.repoRoot.appendingPathComponent(name)

    guard FileManager.default.fileExists(atPath: url.path) else {
        throw DataLoadingError.fileNotFound(url.path)
    }

    let data = try Data(contentsOf: url)

    // Find header end (newline after header)
    var headerEnd = 0
    for i in 0..<data.count {
        if data[i] == 0x0A {
            if i > 5 {
                headerEnd = i + 1
                let remaining = data.count - headerEnd
                if remaining % 2 == 0 {  // float16 = 2 bytes
                    break
                }
            }
        }
    }

    let rawData = data.subdata(in: headerEnd..<data.count)
    let count = rawData.count / 2
    var result = [Float](repeating: 0, count: count)

    rawData.withUnsafeBytes { ptr in
        let f16 = ptr.bindMemory(to: Float16.self)
        for i in 0..<count {
            result[i] = Float(f16[i])
        }
    }

    return result
}

/// Loads a speaker embedding from speakers.json.
func loadSpeaker(_ name: String) throws -> [Float] {
    let url = PlapreConfig.repoRoot.appendingPathComponent("speakers.json")

    guard FileManager.default.fileExists(atPath: url.path) else {
        throw DataLoadingError.fileNotFound(url.path)
    }

    let data = try Data(contentsOf: url)
    let json = try JSONSerialization.jsonObject(with: data) as? [String: [Double]]

    guard let speakers = json else {
        throw DataLoadingError.invalidFormat("speakers.json")
    }

    guard let embedding = speakers[name] else {
        throw DataLoadingError.speakerNotFound(name, speakers.keys.sorted())
    }

    return embedding.map { Float($0) }
}

enum DataLoadingError: LocalizedError {
    case fileNotFound(String)
    case invalidFormat(String)
    case speakerNotFound(String, [String])

    var errorDescription: String? {
        switch self {
        case .fileNotFound(let path):
            return "File not found: \(path)"
        case .invalidFormat(let file):
            return "Invalid format in: \(file)"
        case .speakerNotFound(let name, let available):
            return "Speaker '\(name)' not found. Available: \(available.joined(separator: ", "))"
        }
    }
}