package coffee.liz.lambda.parser;

import coffee.liz.lambda.LambdaDriver;
import coffee.liz.lambda.ast.Expression;
import coffee.liz.lambda.ast.Expression.AbstractionExpression;
import coffee.liz.lambda.ast.Expression.ApplicationExpression;
import coffee.liz.lambda.ast.Expression.IdentifierExpression;
import coffee.liz.lambda.ast.LambdaProgram;
import coffee.liz.lambda.ast.Macro;
import coffee.liz.lambda.ast.SourceCode;
import coffee.liz.lambda.ast.SourceComment;
import org.junit.jupiter.api.Test;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;

class ParserTest {

	@Test
	public void testTrivial() {
		final LambdaProgram lambda = LambdaDriver.parse(SourceCode.ofLambda("λx.x"));
		final LambdaProgram arrow = LambdaDriver.parse(SourceCode.ofArrow("x -> x"));

		assertStructurallyEqual(lambda, arrow);
		assertEquals(0, lambda.macros().size());
		assertInstanceOf(AbstractionExpression.class, lambda.expression());
		assertEquals("x", ((AbstractionExpression) lambda.expression()).parameter());
	}

	@Test
	public void testApplication() {
		final LambdaProgram lambda = LambdaDriver.parse(SourceCode.ofLambda("(λx.x) y"));
		final LambdaProgram arrow = LambdaDriver.parse(SourceCode.ofArrow("(x -> x)(y)"));

		assertStructurallyEqual(lambda, arrow);

		final ApplicationExpression app = (ApplicationExpression) lambda.expression();
		assertInstanceOf(AbstractionExpression.class, app.applicable());
		assertInstanceOf(IdentifierExpression.class, app.argument());
	}

	@Test
	public void testChainedLambdas() {
		final LambdaProgram lambda = LambdaDriver.parse(SourceCode.ofLambda("λx.λy.λz.x"));
		final LambdaProgram arrow = LambdaDriver.parse(SourceCode.ofArrow("x -> y -> z -> x"));

		assertStructurallyEqual(lambda, arrow);

		final AbstractionExpression outer = (AbstractionExpression) lambda.expression();
		assertEquals("x", outer.parameter());
		final AbstractionExpression middle = (AbstractionExpression) outer.body();
		assertEquals("y", middle.parameter());
		final AbstractionExpression inner = (AbstractionExpression) middle.body();
		assertEquals("z", inner.parameter());
	}

	@Test
	public void testChainedApplication() {
		final LambdaProgram lambda = LambdaDriver.parse(SourceCode.ofLambda("((f x) y) z"));
		final LambdaProgram arrow = LambdaDriver.parse(SourceCode.ofArrow("f(x)(y)(z)"));

		assertStructurallyEqual(lambda, arrow);

		final ApplicationExpression app1 = (ApplicationExpression) lambda.expression();
		assertEquals("z", ((IdentifierExpression) app1.argument()).name());
		final ApplicationExpression app2 = (ApplicationExpression) app1.applicable();
		assertEquals("y", ((IdentifierExpression) app2.argument()).name());
		final ApplicationExpression app3 = (ApplicationExpression) app2.applicable();
		assertEquals("f", ((IdentifierExpression) app3.applicable()).name());
		assertEquals("x", ((IdentifierExpression) app3.argument()).name());
	}

	@Test
	public void testMacros() {
		final LambdaProgram lambda = LambdaDriver.parse(SourceCode.ofLambda("""
				let id = λx.x;
				id
				"""));
		final LambdaProgram arrow = LambdaDriver.parse(SourceCode.ofArrow("""
				let id = x -> x;
				id
				"""));

		assertStructurallyEqual(lambda, arrow);
		assertEquals(1, lambda.macros().size());
		assertEquals("id", lambda.macros().getFirst().name());
		assertInstanceOf(IdentifierExpression.class, lambda.expression());
	}

	@Test
	public void testLineComments() {
		final LambdaProgram lambda = LambdaDriver.parse(SourceCode.ofLambda("""
				-- The identity function
				let id = λx.x; -- returns its argument
				id -- use it
				"""));
		final LambdaProgram arrow = LambdaDriver.parse(SourceCode.ofArrow("""
				-- The identity function
				let id = x -> x; -- returns its argument
				id -- use it
				"""));

		assertStructurallyEqual(lambda, arrow);
		assertEquals(1, lambda.macros().size());
		assertEquals("id", lambda.macros().getFirst().name());
	}

	@Test
	public void testComplexProgram() {
		final LambdaProgram lambda = LambdaDriver.parse(SourceCode.ofLambda("""
				let zero = λf.λx.x;
				let one  = λf.λx.f x;
				let succ = λn.λf.λx.f (n f x);
				let add  = λm.λn.λf.λx.m f (n f x);

				succ (add one zero)
				"""));

		final LambdaProgram arrow = LambdaDriver.parse(SourceCode.ofArrow("""
				let zero = f -> x -> x;
				let one  = f -> x -> f(x);
				let succ = n -> f -> x -> f(n(f)(x));
				let add  = m -> n -> f -> x -> m(f)(n(f)(x));

				succ(add(one)(zero))
				"""));

		assertStructurallyEqual(lambda, arrow);

		assertEquals(4, lambda.macros().size());
		assertEquals("zero", lambda.macros().get(0).name());
		assertEquals("one", lambda.macros().get(1).name());
		assertEquals("succ", lambda.macros().get(2).name());
		assertEquals("add", lambda.macros().get(3).name());
	}

	@Test
	public void testOmegaCombinator() {
		final LambdaProgram lambda = LambdaDriver.parse(SourceCode.ofLambda("(λx.x x)(λx.x x)"));
		final LambdaProgram arrow = LambdaDriver.parse(SourceCode.ofArrow("(x -> x(x))(x -> x(x))"));

		assertStructurallyEqual(lambda, arrow);
	}

	private static void assertStructurallyEqual(final LambdaProgram expected, final LambdaProgram actual) {
		assertEquals(expected.macros().size(), actual.macros().size(), "Macro count mismatch");
		for (int i = 0; i < expected.macros().size(); i++) {
			assertStructurallyEqual(expected.macros().get(i), actual.macros().get(i));
		}
		assertStructurallyEqual(expected.expression(), actual.expression());
	}

	private static void assertStructurallyEqual(final Macro expected, final Macro actual) {
		assertEquals(expected.name(), actual.name(), "Macro name mismatch");
		assertEquals(expected.comment().map(SourceComment::text), actual.comment().map(SourceComment::text),
				"Macro comment mismatch");
		assertStructurallyEqual(expected.expression(), actual.expression());
	}

	private static void assertStructurallyEqual(final Expression expected, final Expression actual) {
		assertEquals(expected.getClass(), actual.getClass(), "Expression type mismatch");
		assertEquals(expected.comment().map(SourceComment::text), actual.comment().map(SourceComment::text),
				"Expression comment mismatch");

		switch (expected) {
			case IdentifierExpression e ->
				assertEquals(e.name(), ((IdentifierExpression) actual).name(), "Identifier name mismatch");
			case AbstractionExpression e -> {
				assertEquals(e.parameter(), ((AbstractionExpression) actual).parameter(), "Parameter mismatch");
				assertStructurallyEqual(e.body(), ((AbstractionExpression) actual).body());
			}
			case ApplicationExpression e -> {
				assertStructurallyEqual(e.applicable(), ((ApplicationExpression) actual).applicable());
				assertStructurallyEqual(e.argument(), ((ApplicationExpression) actual).argument());
			}
		}
	}
}