import commandline.options.*;
import console.color.*;

record TestStatus (
    passed: UInt,
    failed: UInt,
    pending: UInt,
);

overload TestStatus() = TestStatus(UInt(0), UInt(0), UInt(0));

record TestCase (
    name: String,
    _case: Lambda[[TestStatus], []],
    pending?: Bool,
);

[T | String?(T) and T != String]
overload TestCase(name : T, _case, pending?) : TestCase
    = TestCase(String(name), _case, pending?);

overload TestCase(name, fn) = TestCase(name, Lambda[[TestStatus], []](fn), false);

PendingTestCase(name, fn) = TestCase(name, Lambda[[TestStatus], []](fn), true);

record TestSuite (
    name: String,
    testCases: Vector[TestCase]
);

[S, T | S != String or T != Vector[TestCase]] 
overload TestSuite(name : S, testCases : T) = TestSuite(String(name), Vector[TestCase](testCases));

alias TEST_CASE_INDENT = "  ";
alias TEST_FAILURE_INDENT = TEST_CASE_INDENT + TEST_CASE_INDENT;

accumulateTestStatus(test: TestStatus, newTest: TestStatus) {
    test.passed  += newTest.passed;
    test.failed  += newTest.failed;
    test.pending += newTest.pending;
}

reportTest(test: TestStatus, ..prefix) {
    println(
        ..prefix,
        test.passed, " expectations passed; ",
        test.failed, " failed (",
        test.pending, " pending",
        ")"
    );
}

runTestSuite(suite: TestSuite) : TestStatus {
    var test = TestStatus();
    println("Testing ", suite.name, ":");
    for (_case in suite.testCases)
      if(runCase?(_case))
        accumulateTestStatus(test, runTestCase(_case));

    reportTest(test, "-- ", suite.name, ": ");
    return test;
}

pendingMessage(pending?) {
    if (pending?) {
        return "(pending) ";
    } else {
        return "";
    }
}

printFailure(..xs) { 
  printlnColored(Red, TEST_FAILURE_INDENT, ..xs); 
}

runTestCase(_case: TestCase) : TestStatus {
    print(TEST_CASE_INDENT, pendingMessage(_case.pending?), _case.name, ":");
    flush(stdout);
    
    var test = TestStatus();
    try {
        _case._case(test);
    } catch(e:RequiredExpectationFailed) {
        println();
        printFailure("required expectation did not pass");
    } catch(e) {
        failed(test);
        println();
        printFailure("uncaught exception:");
        printFailure(TEST_CASE_INDENT, *e);
    }

    if (test.failed > 0) {
        if (_case.pending?) {
            test.pending = test.failed;
            test.failed = 0;
        } else {
            printFailure("-- FAILED");
        }
    } else {
        if (_case.pending?) {
            println();
            printFailure("-- FIXED (should clear pending flag)");
            test.failed = 1;
        } else {
            printlnColored(Green, " passed");
        }
    }

    return test;
}

passed(test) { test.passed += 1; }
failed(test) { test.failed += 1; }

// primitive assertion function
expect(test: TestStatus, passed?: Bool, failReportFn) {
    if (passed?) {
        passed(test);
    } else {
        failed(test);
        println();
        failReportFn();
    }
}

expectTrue(test: TestStatus, name, value) {
    expect(test, value, () => {
        printFailure(name, " is expected to be true but isn't");
    });
}

expectFalse(test: TestStatus, name, value) {
    expect(test, not value, () => {
        printFailure(name, " is expected to be false but isn't");
    });
}

alias expectType(test: TestStatus, name, ExpectedType, result) {
    expectType2(test, name, ExpectedType, Type(result));
}

private expectType2(test:TestStatus, name, ExpectedType, ResultType) {
    expect(test, ResultType == ExpectedType, () => {
        printFailure(name, " is expected to be of type ", ExpectedType, " but is of type ", ResultType); 
    });
}

expectEqual(test: TestStatus, name, expected, result) {
    expect(
        test,
        expected == result,
        () => {
            printFailure(name, " is expected to equal ", expected, " but is ", result);
        }
    );
}

expectNotEqual(test: TestStatus, name, expected, result) {
    expect(
        test,
        expected != result,
        () => {
            printFailure(name, " is expected not to equal ", expected, " but is ", result);
        }
    );
}

define expectedResults?;
overload expectedResults?(expected, result) = false;
[T | Tuple?(T)]
overload expectedResults?(expected: T, result: T) = expected == result;

expectEqualValues(test: TestStatus, name, expectedTuple, ..results) {
    var resultsTuple = Tuple(..results);
    expect(
        test,
        expectedResults?(expectedTuple, resultsTuple),
        () => {
            printFailure(name, " are expected to equal:");
            printFailure(TEST_CASE_INDENT, ..weaveValues(", ", ..unpack(expectedTuple)));
            printFailure("but are:");
            printFailure(TEST_CASE_INDENT, ..weaveValues(", ", ..unpack(resultsTuple)));
        }
    );
}

expectCallDefined(test: TestStatus, Proc, ..Types) {
    var typesTuple = Tuple(..Types);
    expect(
        test,
        CallDefined?(Proc, ..Types),
        () => {
            printFailure(Proc, " isn't callable with the expected argument types:");
            printFailure(TEST_CASE_INDENT, ..weaveValues(", ", ..unpack(typesTuple)));
        }
    );
}

expectCallUndefined(test: TestStatus, Proc, ..Types) {
    var typesTuple = Tuple(..Types);
    expect(
        test,
        not CallDefined?(Proc, ..Types),
        () => {
            printFailure(Proc, " shouldn't be callable with the expected argument types:");
            printFailure(TEST_CASE_INDENT, ..weaveValues(", ", ..unpack(typesTuple)));
        }
    );
}

expectIterator(test: TestStatus, name, iter, ..expected) {
    var expectedV = Vector[Type(next(iter))](..expected);
    var resultV = Vector[Type(next(iter))]();
    while (hasNext?(iter))
        push(resultV, next(iter));
    expect(
        test,
        expectedV == resultV,
        () => {
            printFailure(name, " is expected to contain:");
            for (expectedX in expectedV)
                printFailure(TEST_CASE_INDENT, expectedX);
            printFailure("but contains:");
            for (resultX in resultV)
                printFailure(TEST_CASE_INDENT, resultX);
        }
    );
}

expectSequence(test: TestStatus, name, seq, ..expected) {
    expectIterator(test, name, iterator(seq), ..expected);
}

expectExceptionType(test: TestStatus, name, T, fn) {
    try {
        fn();
    } catch(x: T) {
        expect(test, true, () => {});
        return;
    } catch(x) {
        expect(test, false, () => {
            printFailure(name, " is expected to throw exception of type ", T, " but throws ", x);
        });
        return;
    }

    expect(test, false, () => {
        printFailure(name, " is expected to throw exception of type ", T, " but doesn't throw");
    });
}

expectException(test: TestStatus, name, expectedException, fn) {
    try {
        fn();
    } catch(x: Type(expectedException)) {
        expect(test, x == expectedException, () => {
            printFailure(name, " is expected to throw ", expectedException, " but throws ", x);
        });
        return;
    } catch(x) {
        expect(test, false, () => {
            printFailure(name, " is expected to throw ", expectedException, " but throws ", x);
        });
        return;
    }

    expect(test, false, () => {
        printFailure(name, " is expected to throw ", expectedException, " but doesn't throw");
    });
}

expectPrinted(test: TestStatus, name, expectedOutput, forward ..values) {
    expectEqual(test, name, expectedOutput, str(..values));
}

record RequiredExpectationFailed();
instance Exception (RequiredExpectationFailed);

required(expectFn, test: TestStatus, forward ..args) {
    var preFailed = test.failed;
    expectFn(test, ..args);
    if (test.failed != preFailed)
        throw RequiredExpectationFailed();
}

var testOptionSpecs = OptionSpecs(
  OptionSpec("s", "suite", HAS_VALUE, ALLOWS_REPEATS),
  OptionSpec("c", "case", HAS_VALUE, ALLOWS_REPEATS),
  OptionSpec("n", "nc", "no-color")
);

var testOptionsInitialized? = false;
var _testOptions = OptionTable();

private testOptions() {
  if(not testOptionsInitialized?){
    _testOptions = CommandlineOptionTable(testOptionSpecs);
    testOptionsInitialized? = true;
  }

  return ref _testOptions;
}

[..T] ifPasses(expectFn, test: TestStatus, ..args:T) {
    forward fn = dropValues(static countValues(..T) - 1, ..args);
    var preFailed = test.failed;
    expectFn(test, ..takeValues(static countValues(..T) - 1, ..args));
    if (test.failed == preFailed)
        fn();
}

testMain(..suites) {
    var test = TestStatus();
    eachValue(
        suite -> { 
          if(runSuite?(suite))
            accumulateTestStatus(test, runTestSuite(suite)); 
        },
        ..suites
    );

    reportTest(test, "Overall: ");

    if (test.failed > 0) {
        println("TESTS FAILED");
        return 1;
    } else {
        return 0;
    }
}

private printlnColored(c : Color, ..values){
  if(colored?()) textColor(c);
  println(..values);
  if(colored?()) resetConsole();
}

private runSuite?(suite) = optionsInclude?("suite", suite.name);
private runCase?(_case) = optionsInclude?("case", _case.name);
private colored?() = not(optionGiven?(testOptions(), "no-color"));

private optionsInclude?(name, elt){
  forward seq = getOptions(testOptions(), name);
  return empty?(seq) or include?(seq, elt);
}

private include?(seq, elt){
  for(x in seq) if(x == elt) return true;
  return false;
}
