// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Datasets tab: the registered datasets (token counts, vocab, size) and the
// "new dataset" flow driving tools/prepare_data.py (TinyStories) or
// tools/prepare_hf_data.py (HF sources & research presets) with a live
// progress console. Python and the tools live in the forge repo, located
// from the configured binary.
import SwiftUI

struct DatasetsView: View {
    @Environment(AppModel.self) private var app
    @State private var datasets: [Dataset] = []
    @State private var showNew = false

    var body: some View {
        VStack(alignment: .leading, spacing: 12) {
            HStack {
                Text("Datasets").font(.title2.weight(.semibold))
                Spacer()
                Button {
                    showNew = true
                } label: {
                    Label("Nouveau dataset", systemImage: "plus")
                }
            }
            if datasets.isEmpty {
                ContentUnavailableView(
                    "Aucun dataset",
                    systemImage: "cylinder.split.1x2",
                    description: Text("Préparez TinyStories ou un mélange Hugging Face — tout se passe ici, sans terminal."))
            } else {
                Table(datasets) {
                    TableColumn("Nom") { Text($0.name) }
                    TableColumn("Tokens train") {
                        Text($0.trainTokens?.formatted(.number.notation(.compactName)) ?? "?")
                    }
                    TableColumn("Tokens val") {
                        Text($0.valTokens?.formatted(.number.notation(.compactName)) ?? "?")
                    }
                    TableColumn("Vocab") { Text($0.vocabSize.map(String.init) ?? "?") }
                    TableColumn("Taille") {
                        Text(ByteCountFormatter.string(fromByteCount: $0.sizeOnDisk,
                                                       countStyle: .file))
                    }
                    TableColumn("Chemin") {
                        Text($0.path).font(.caption.monospaced())
                            .foregroundStyle(.secondary)
                    }
                }
            }
        }
        .padding()
        .onAppear { datasets = app.datasets() }
        .sheet(isPresented: $showNew, onDismiss: { datasets = app.datasets() }) {
            NewDatasetSheet()
        }
        .navigationTitle("Datasets")
    }
}

struct NewDatasetSheet: View {
    @Environment(AppModel.self) private var app
    @Environment(\.dismiss) private var dismiss

    enum Kind: String, CaseIterable {
        case tinystories = "TinyStories"
        case hf = "Hugging Face"
    }

    @State private var kind: Kind = .tinystories
    @State private var name = "tinystories"
    @State private var vocabSize = 4096
    @State private var maxTrainMB = 100.0
    @State private var hfChoice = "preset:smollm-web"
    @State private var console: [String] = []
    @State private var running = false
    @State private var failed = false

    // Mirrors SOURCES/PRESETS in tools/prepare_hf_data.py.
    private let hfChoices: [(label: String, value: String)] = [
        ("Preset · smollm-web (60% FineWeb-Edu + 40% DCLM)", "preset:smollm-web"),
        ("Preset · textbooks (FineWeb-Edu + Cosmopedia)", "preset:textbooks"),
        ("Preset · smol-full (edu+cosmo+math)", "preset:smol-full"),
        ("Preset · decay-anneal (fin de run WSD)", "preset:decay-anneal"),
        ("Source · fineweb-edu", "source:fineweb-edu"),
        ("Source · cosmopedia", "source:cosmopedia"),
        ("Source · finemath", "source:finemath"),
        ("Source · wikipedia-fr", "source:wikipedia-fr"),
        ("Source · tinystories (via HF)", "source:tinystories"),
    ]

    var body: some View {
        VStack(alignment: .leading, spacing: 12) {
            Text("Nouveau dataset").font(.title2.weight(.semibold))
            Picker("Type", selection: $kind) {
                ForEach(Kind.allCases, id: \.self) { Text($0.rawValue) }
            }
            .pickerStyle(.segmented)
            Form {
                TextField("Nom (dossier dans data/)", text: $name)
                if kind == .hf {
                    Picker("Corpus", selection: $hfChoice) {
                        ForEach(hfChoices, id: \.value) { Text($0.label).tag($0.value) }
                    }
                }
                TextField("vocab_size", value: $vocabSize, format: .number)
                TextField("Taille max train (MB)", value: $maxTrainMB, format: .number)
            }
            .formStyle(.grouped)
            .frame(height: kind == .hf ? 180 : 150)
            .disabled(running)

            if !console.isEmpty {
                ConsoleView(lines: console)
            }

            HStack {
                if failed {
                    Label("Échec — voir la console.", systemImage: "xmark.circle")
                        .foregroundStyle(.red)
                }
                Spacer()
                Button(running ? "Fermer (continue en arrière-plan)" : "Fermer") {
                    dismiss()
                }
                Button(running ? "Préparation…" : "Préparer") { prepare() }
                    .keyboardShortcut(.defaultAction)
                    .disabled(running || name.isEmpty)
            }
        }
        .padding()
        .frame(minWidth: 640, minHeight: 480)
    }

    private var forgeRepo: URL? {
        // <repo>/build/forge → <repo>
        ForgeBinaryLocator.savedBinaryURL?
            .deletingLastPathComponent().deletingLastPathComponent()
    }

    private func prepare() {
        guard let repo = forgeRepo else {
            console = ["⚠︎ Binaire forge non configuré — Réglages."]
            failed = true
            return
        }
        let outDir = app.store.workspaceURL.appendingPathComponent("data/\(name)")
        var args: [String]
        switch kind {
        case .tinystories:
            args = ["tools/prepare_data.py", "--out", outDir.path,
                    "--vocab-size", String(vocabSize),
                    "--max-train-mb", String(maxTrainMB)]
        case .hf:
            let parts = hfChoice.split(separator: ":", maxSplits: 1).map(String.init)
            args = ["tools/prepare_hf_data.py", "--\(parts[0])", parts[1],
                    "--out", outDir.path, "--vocab-size", String(vocabSize),
                    "--max-train-mb", String(maxTrainMB)]
        }
        running = true
        failed = false
        console = ["$ python3 " + args.joined(separator: " ")]
        Task {
            do {
                let runner = ProcessRunner()
                let (lines, exit) = try await runner.launch(
                    executable: URL(fileURLWithPath: "/usr/bin/env"),
                    arguments: ["python3"] + args,
                    currentDirectory: repo)
                for await (line, isErr) in lines {
                    console.append(isErr ? "⚠︎ " + line : line)
                    if console.count > 2000 { console.removeFirst(500) }
                }
                let status = await exit.value
                running = false
                failed = status.code != 0
                if !failed { console.append("✓ Dataset prêt : \(outDir.path)") }
            } catch {
                running = false
                failed = true
                console.append("⚠︎ " + error.localizedDescription)
            }
        }
    }
}