import Foundation

/// Handles sharded SafeTensors models (with model.safetensors.index.json).
public final class SafeTensorsIndex {
    public let weightMap: [String: String]
    public let baseDir: String

    /// Load the index file from a model directory.
    public init(modelDir: String) throws {
        let indexURL = URL(fileURLWithPath: modelDir).appendingPathComponent("model.safetensors.index.json")
        let data = try Data(contentsOf: indexURL)
        guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
              let weightMap = json["weight_map"] as? [String: String]
        else {
            throw WeightError.invalidHeader("Index file missing weight_map")
        }
        self.weightMap = weightMap
        self.baseDir = modelDir
    }

    /// All unique shard filenames referenced by the index.
    public var shardFiles: Set<String> {
        Set(weightMap.values)
    }

    /// Resolve a tensor name to its shard file path.
    public func shardPath(for tensor: String) -> String? {
        guard let shard = weightMap[tensor] else { return nil }
        return (baseDir as NSString).appendingPathComponent(shard)
    }

    /// List all tensor names.
    public var allTensors: [String] { Array(weightMap.keys) }
}