//
//  EngineTests.swift
//  Metrika
//
//  Author:  Simon-Pierre Boucher
//  Contact: contact@spboucher.ai
//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
//

import Testing
import ZQEngine

/// End-to-end tests: full command lines through the session, checking the
/// returned scalar results against the R fixtures.
@Suite("ZQEngine end-to-end", .serialized)
struct EngineTests {
    let fixtures: Fixtures
    let session: ZQSession

    init() async throws {
        self.fixtures = try Fixtures()
        self.session = try ZQSession(discoverUserCommands: false)
        _ = try await session.execute("use \(fixtures.datasetURL.path)")
        _ = try await session.execute("gen log_rev = ln(revenue)")
    }

    @Test("use reports dataset dimensions")
    func loadedDimensions() async throws {
        let result = try await session.execute("count")
        #expect(result.scalars["N"] == 60)
    }

    @Test("sysuse loads bundled samples and lists them")
    func sysuse() async throws {
        let fresh = try ZQSession(discoverUserCommands: false)
        let listing = try await fresh.execute("sysuse")
        #expect(listing.text.contains("sales"))
        let sales = try await fresh.execute("sysuse sales")
        #expect(sales.scalars["N"] == 1000)
        let reg = try await fresh.execute("reg revenue price")
        #expect(reg.scalars["N"] == 1000)
        let cars = try await fresh.execute("sysuse mtcars")
        #expect(cars.scalars["N"] == 32)
        await #expect(throws: ZQEngineError.self) {
            _ = try await fresh.execute("sysuse nosuchsample")
        }
    }

    @Test("regress with listwise deletion matches R")
    func regress() async throws {
        let result = try await session.execute("reg log_rev price")
        #expect(result.scalars["N"] == fixtures["ols_n"])
        expectClose(try #require(result.scalars["b_price"]), fixtures["ols_b_price"], "b[price]")
        expectClose(try #require(result.scalars["b__cons"]), fixtures["ols_b_cons"], "b[_cons]")
        expectClose(
            try #require(result.scalars["se_price"]),
            fixtures["ols_se_classical_price"], "se[price]"
        )
        expectClose(try #require(result.scalars["r2"]), fixtures["ols_r2"], "R²")
        #expect(result.text.contains("observations dropped due to missing values"))
    }

    @Test("regress, robust matches R HC1")
    func regressRobust() async throws {
        let result = try await session.execute("reg log_rev price, robust")
        expectClose(
            try #require(result.scalars["se_price"]),
            fixtures["ols_se_hc1_price"], "robust se[price]"
        )
    }

    @Test("regress, cluster matches R with Stata small-sample factor")
    func regressCluster() async throws {
        let result = try await session.execute("reg log_rev price, cluster(firm_id)")
        #expect(result.scalars["N_clust"] == fixtures["ols_G"])
        expectClose(
            try #require(result.scalars["se_price"]),
            fixtures["ols_se_cluster_price"], "cluster se[price]"
        )
    }

    @Test("factor variable expansion i.region matches R factor()")
    func regressFactor() async throws {
        let result = try await session.execute("reg log_rev price i.region")
        expectClose(try #require(result.scalars["b_price"]), fixtures["ols2_b_price"], "b[price]")
        expectClose(
            try #require(result.scalars["b_2.region"]),
            fixtures["ols2_b_region2"], "b[2.region]"
        )
        expectClose(
            try #require(result.scalars["b_3.region"]),
            fixtures["ols2_b_region3"], "b[3.region]"
        )
        expectClose(try #require(result.scalars["b__cons"]), fixtures["ols2_b_cons"], "b[_cons]")
        expectClose(try #require(result.scalars["r2"]), fixtures["ols2_r2"], "R²")
    }

    @Test("summarize scalars match R")
    func summarize() async throws {
        let result = try await session.execute("summarize revenue")
        expectClose(
            try #require(result.scalars["mean"]),
            fixtures["sum_revenue_mean"], "mean"
        )
        expectClose(try #require(result.scalars["sd"]), fixtures["sum_revenue_sd"], "sd")
    }

    @Test("if qualifier and count")
    func conditionalCount() async throws {
        let all = try await session.execute("count")
        let some = try await session.execute("count if missing(price)")
        #expect(all.scalars["N"] == 60)
        #expect(some.scalars["N"] == fixtures["ols_dropped"])
    }

    @Test("bootstrap is reproducible for a fixed seed")
    func bootstrapReproducible() async throws {
        let first = try await session.execute(
            "bootstrap, reps(200) seed(42): reg log_rev price"
        )
        let second = try await session.execute(
            "bootstrap, reps(200) seed(42): reg log_rev price"
        )
        #expect(first.scalars["se_price"] == second.scalars["se_price"])
        // Bootstrap SE should be in the neighborhood of the analytic one.
        let analytic = fixtures["ols_se_classical_price"]
        let bootstrap = try #require(first.scalars["se_price"])
        #expect(bootstrap > analytic / 3 && bootstrap < analytic * 3)
    }

    @Test("generate honors if, replace counts changes")
    func generateReplace() async throws {
        _ = try await session.execute("gen flag = 1 if price > 10")
        let count = try await session.execute("count if flag == 1")
        let expected = try await session.execute("count if price > 10")
        #expect(count.scalars["N"] == expected.scalars["N"])
        let replaced = try await session.execute("replace flag = 0 if missing(flag)")
        #expect(replaced.text.contains("real changes"))
        _ = try await session.execute("drop flag")
    }

    @Test("graph scatter produces a plot spec")
    func graphScatter() async throws {
        _ = try await session.execute("graph scatter log_rev price, by(region)")
        let plot = try #require(await session.lastPlot)
        #expect(plot.kind == .scatter)
        #expect(plot.series.count == 3)  // three regions
        #expect(plot.xLabel == "price")
    }

    @Test("logit through the console matches R")
    func logitCommand() async throws {
        let result = try await session.execute("logit purchase price")
        expectClose(
            try #require(result.scalars["b_price"]), fixtures["logit_b_price"], "b[price]"
        )
        expectClose(try #require(result.scalars["ll"]), fixtures["logit_ll"], "ll")
        #expect(result.text.contains("Logistic regression"))
        #expect(result.text.contains("Pseudo R2"))
    }

    @Test("poisson with robust SE through the console")
    func poissonCommand() async throws {
        let result = try await session.execute("poisson orders price, robust")
        expectClose(
            try #require(result.scalars["b_price"]), fixtures["pois_b_price"], "b[price]"
        )
        expectClose(
            try #require(result.scalars["se_price"]),
            fixtures["pois_se_hc0_price"], "robust se[price]"
        )
    }

    @Test("tabulate one-way counts region levels")
    func tabulateOneWay() async throws {
        let result = try await session.execute("tab region")
        #expect(result.scalars["N"] == 60)
        #expect(result.scalars["r"] == 3)
        #expect(result.text.contains("Total"))
        // region cycles 1,2,3 over 60 rows → 20 each.
        #expect(result.text.contains("20"))
    }

    @Test("tabulate two-way region by purchase")
    func tabulateTwoWay() async throws {
        let result = try await session.execute("tabulate region purchase")
        #expect(result.scalars["r"] == 3)
        #expect(result.scalars["c"] == 2)
        #expect(result.scalars["N"] == 60)
    }

    @Test("correlate matches R and reports listwise obs")
    func correlateCommand() async throws {
        let result = try await session.execute("correlate revenue price")
        expectClose(
            try #require(result.scalars["rho"]),
            fixtures["corr_rev_price"], "cor(revenue, price)"
        )
        #expect(result.scalars["N"] == fixtures["ols_n"])
    }

    @Test("help lists every category; help <verb> shows the entry")
    func help() async throws {
        let index = try await session.execute("help")
        for category in ZQCommandReference.categories {
            #expect(index.text.contains(category), Comment(rawValue: category))
        }
        let entry = try await session.execute("help regress")
        #expect(entry.text.contains("cluster(var)"))
        #expect(entry.text.contains("reg log_rev price, robust"))
        let byAbbreviation = try await session.execute("help reg")
        #expect(byAbbreviation.text.contains("Ordinary least squares"))
        let unknown = try await session.execute("help nosuchthing")
        #expect(unknown.text.contains("no entry"))
    }

    @Test("every dispatched estimation and data verb has a manual entry")
    func manualCoverage() {
        let documented = Set(ZQCommandReference.all.map(\.verb))
        for verb in [
            "use", "sysuse", "save", "clear", "describe", "list", "count", "generate",
            "replace", "drop", "keep", "summarize", "tabulate", "correlate",
            "regress", "logit", "probit", "poisson", "ivregress", "xtreg",
            "xtset", "lasso", "elasticnet", "boost", "bootstrap", "permute", "predict",
            "margins", "bayes", "scatter", "histogram", "kdensity", "graph", "display",
            "set", "log", "help", "zscore",
        ] {
            #expect(documented.contains(verb), Comment(rawValue: "missing manual entry: \(verb)"))
        }
    }

    @Test("display evaluates scalar expressions")
    func display() async throws {
        let result = try await session.execute("display 2 + 2 * 3")
        #expect(result.text == "8")
    }

    @Test("unknown variable yields a helpful error")
    func unknownVariable() async throws {
        await #expect(throws: (any Error).self) {
            _ = try await session.execute("summarize nonexistent_var")
        }
    }
}