// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Contract tests: config codec against the REAL forge configs, the
// param-count formula against `forge info`, LTTB invariants, EMA bias
// correction, CSV parser against real and mutated lines, state machine.
import XCTest
@testable import ForgeStudio

final class ForgeStudioTests: XCTestCase {
    let forgeRepo = URL(fileURLWithPath: NSHomeDirectory() + "/Desktop/forge")

    func testConfigRoundTripAgainstRealConfigs() throws {
        let configsDir = forgeRepo.appendingPathComponent("configs")
        let files = (try? FileManager.default.contentsOfDirectory(
            at: configsDir, includingPropertiesForKeys: nil))?
            .filter { $0.pathExtension == "json" } ?? []
        try XCTSkipIf(files.isEmpty, "forge repo introuvable")
        for url in files {
            let cfg = try ForgeConfig.load(from: url)
            // Re-export and re-parse: field names must be byte-compatible.
            let data = try cfg.exportJSON()
            let reparsed = try JSONDecoder().decode(ForgeConfig.self, from: data)
            XCTAssertEqual(cfg, reparsed, url.lastPathComponent)
            XCTAssertTrue(cfg.validationErrors.isEmpty,
                          "\(url.lastPathComponent): \(cfg.validationErrors)")
        }
    }

    func testParamCountMatchesForgeInfo() throws {
        // Known-good values printed by `forge info` for the shipped configs.
        let expected: [(String, Double)] = [
            ("gpt-50m.json", 52.20), ("gpt-50m-moe.json", 151.75),
            ("gpt-50m-deep.json", 50.60), ("gpt-50m-mistral.json", 48.10),
        ]
        for (file, millions) in expected {
            let url = forgeRepo.appendingPathComponent("configs/\(file)")
            guard let cfg = try? ForgeConfig.load(from: url) else { continue }
            XCTAssertEqual(Double(cfg.model.paramCount) / 1e6, millions,
                           accuracy: 0.01, file)
        }
    }

    func testLTTBInvariants() {
        let pts = (0..<10_000).map {
            Downsampler.XY(x: Double($0), y: sin(Double($0) / 50) + Double($0) * 0.001)
        }
        let ds = Downsampler.lttb(pts, threshold: 500)
        XCTAssertEqual(ds.count, 500)
        XCTAssertEqual(ds.first, pts.first) // endpoints preserved exactly
        XCTAssertEqual(ds.last, pts.last)
        XCTAssertTrue(zip(ds, ds.dropFirst()).allSatisfy { $0.x < $1.x }) // monotonic
        XCTAssertEqual(Downsampler.lttb(pts, threshold: 20_000).count, pts.count)
    }

    func testEMABiasCorrection() {
        // A constant series must stay exactly constant under bias-corrected EMA.
        let ema = Smoothing.ema([Double](repeating: 3.5, count: 100), smoothing: 0.9)
        XCTAssertTrue(ema.allSatisfy { abs($0 - 3.5) < 1e-9 })
        XCTAssertEqual(Smoothing.ema([1, 2, 3], smoothing: 0), [1, 2, 3])
    }

    func testCSVParser() {
        var p = LogParser()
        XCTAssertNil(p.parseCSVLine("step,loss,lr,grad_norm,tokens_per_sec,val_loss,elapsed_s"))
        let pt = p.parseCSVLine("584,2.296893,3.680000e-04,0.558962,13692.0,-1.000000,3861.402")
        XCTAssertEqual(pt?.step, 584)
        XCTAssertEqual(pt!.trainLoss, 2.296893, accuracy: 1e-9)
        XCTAssertNil(pt?.valLoss) // -1 sentinel
        XCTAssertEqual(pt!.elapsedS!, 3861.402, accuracy: 1e-6)
        let withVal = p.parseCSVLine("599,2.19,3.05e-04,0.55,13748.0,2.3299,3960.1")
        XCTAssertEqual(withVal!.valLoss!, 2.3299, accuracy: 1e-9)
        // Legacy 6-column header, garbage, truncation: never crash.
        var legacy = LogParser()
        _ = legacy.parseCSVLine("step,loss,lr,grad_norm,tokens_per_sec,val_loss")
        let old = legacy.parseCSVLine("10,5.2,1e-4,0.9,42000.0,-1.0")
        XCTAssertNil(old?.elapsedS)
        XCTAssertNil(legacy.parseCSVLine("garbage,🤖,,"))
        XCTAssertNil(legacy.parseCSVLine("1,2"))
        XCTAssertNil(legacy.parseCSVLine(""))
    }

    func testStdoutEvents() {
        XCTAssertEqual(
            LogParser.parseStdout("checkpoint saved: runs/x/ckpt_000200.bin"),
            .checkpointSaved(path: "runs/x/ckpt_000200.bin"))
        if case .banner(let params, let steps, let tps, let backend)? =
            LogParser.parseStdout(
                "training gpt-50m: 52196480 params, 1170 steps, 65536 tokens/step, backend=metal, opt=adamw, sched=cosine")
        {
            XCTAssertEqual(params, 52_196_480)
            XCTAssertEqual(steps, 1170)
            XCTAssertEqual(tps, 65536)
            XCTAssertEqual(backend, "metal")
        } else {
            XCTFail("banner non parsé")
        }
    }

    func testRunStateMachine() {
        // Terminal states allow nothing; queued can't jump to finished.
        XCTAssertTrue(RunState.transitions[.finished]!.isEmpty)
        XCTAssertTrue(RunState.transitions[.failed]!.isEmpty)
        XCTAssertFalse(RunState.transitions[.queued]!.contains(.finished))
        XCTAssertTrue(RunState.transitions[.running]!.contains(.finishing))
    }

    func testLRSchedulePreviewMatchesForgeMath() {
        var t = TrainConfig()
        t.lr = 5e-4; t.warmupSteps = 100; t.maxSteps = 1000
        XCTAssertEqual(t.lrAt(step: 0), 5e-4 * 1.0 / 101.0, accuracy: 1e-12)
        XCTAssertEqual(t.lrAt(step: 999) > t.lr * t.minLrRatio, true)
        XCTAssertEqual(t.lrAt(step: 1000), t.lr * t.minLrRatio, accuracy: 1e-12)
        t.schedule = "wsd"
        XCTAssertEqual(t.lrAt(step: 500), t.lr, accuracy: 1e-15) // plateau
    }
}