// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// A prepared dataset folder: train.bin/val.bin (llm.c-style header
// {magic 20240520, version, num_tokens} then uint16 tokens) + a forgebpe
// tokenizer model whose vocab size must match model.vocab_size.
import Foundation

struct Dataset: Identifiable, Equatable {
    var id: String { path }
    var path: String
    var name: String
    var trainTokens: Int?
    var valTokens: Int?
    var vocabSize: Int?
    var sizeOnDisk: Int64

    static let binMagic: Int32 = 20_240_520

    // num_tokens lives at int32 offset 2 of the 256-int32 header.
    static func tokenCount(binURL: URL) -> Int? {
        guard let fh = try? FileHandle(forReadingFrom: binURL),
              let data = try? fh.read(upToCount: 12), data.count == 12 else {
            return nil
        }
        try? fh.close()
        let ints = data.withUnsafeBytes { $0.load(fromByteOffset: 0, as: (Int32, Int32, Int32).self) }
        guard ints.0 == binMagic else { return nil }
        return Int(ints.2)
    }

    // forgebpe v1 header: line 1 "forgebpe v1", line 2 vocab size.
    static func vocabSize(tokenizerURL: URL) -> Int? {
        guard let fh = try? FileHandle(forReadingFrom: tokenizerURL),
              let data = try? fh.read(upToCount: 64),
              let head = String(data: data, encoding: .utf8) else { return nil }
        try? fh.close()
        let lines = head.split(separator: "\n", maxSplits: 2)
        guard lines.count >= 2, lines[0].hasPrefix("forgebpe") else { return nil }
        return Int(lines[1].trimmingCharacters(in: .whitespaces))
    }

    static func scan(directory: URL) -> Dataset? {
        let fm = FileManager.default
        let train = directory.appendingPathComponent("train.bin")
        guard fm.fileExists(atPath: train.path) else { return nil }
        let val = directory.appendingPathComponent("val.bin")
        let tok = (try? fm.contentsOfDirectory(at: directory,
                                               includingPropertiesForKeys: nil))?
            .first { $0.pathExtension == "model" }
        var size: Int64 = 0
        if let items = try? fm.contentsOfDirectory(at: directory,
                                                   includingPropertiesForKeys: [.fileSizeKey]) {
            for f in items {
                size += Int64((try? f.resourceValues(forKeys: [.fileSizeKey]).fileSize) ?? 0)
            }
        }
        return Dataset(
            path: directory.path,
            name: directory.lastPathComponent,
            trainTokens: tokenCount(binURL: train),
            valTokens: fm.fileExists(atPath: val.path) ? tokenCount(binURL: val) : nil,
            vocabSize: tok.flatMap { vocabSize(tokenizerURL: $0) },
            sizeOnDisk: size)
    }
}

struct Checkpoint: Identifiable, Equatable {
    var id: String { path }
    var path: String
    var step: Int
    var sizeBytes: Int64
    var modifiedAt: Date
    var valLossAtStep: Double?
}