package compiler_test

import (
	"jinx/pkg/lang/compiler"
	"jinx/pkg/lang/parser"
	"jinx/pkg/lang/scanner"
	"jinx/pkg/lang/vm/text"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestSimpleAddExpr(t *testing.T) {
	src := `
	1 + 2
	`

	expected := `
	push_int 1
	push_int 2
	add
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestOperationOrder(t *testing.T) {
	grouped := `
	(((1 + 2) - 3) + 4)
	`

	free := `
	1 + 2 - 3 + 4
	`

	expected := `
	push_int 1
	push_int 2
	add
	push_int 3
	sub
	push_int 4
	add
	halt
	`

	mustCompileTo(t, grouped, expected)
	mustCompileTo(t, free, expected)
}

func TestNestedExpr(t *testing.T) {
	// Yeah, we don't compile identifiers yet, so here we call and index directly on literals.
	src := `
	1 + 2 - 3 - (123("a", "b", "c") + 456[321])
	`

	expected := `
	push_int 1
	push_int 2
	add
	
	push_int 3
	sub

	push_int 123
	push_string "a"
	push_string "b"
	push_string "c"
	call 3

	push_int 456
	push_int 321
	index

	add
	sub

	halt
	`

	mustCompileTo(t, src, expected)
}

func TestVars(t *testing.T) {
	src := `
	var x = 10
	var y = 25
	x = x + 7

	var res = x + y
	`

	expected := `
	push_int 10

	push_int 25

	get_local 0
	push_int 7
	add
	set_local 0

	get_local 0
	get_local 1
	add

	halt
	`

	mustCompileTo(t, src, expected)
}

func TestArrayLit(t *testing.T) {
	src := `var x = [1, 2, 3]`

	expected := `
	push_array

	get_local 0
	get_member "push"
	push_int 1
	call 1

	get_local 0
	get_member "push"
	push_int 2
	call 1

	get_local 0
	get_member "push"
	push_int 3
	call 1
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestIf(t *testing.T) {
	src := `
	":("
	if 1 <= 5 {
		"hello " + "world"
	}
	`

	expected := `
	push_string ":("

	push_int 1
	push_int 5
	lte

	jf @end

	push_string "hello "
	push_string "world"
	add

	@end:
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestIfElifElse(t *testing.T) {
	src := `
		if false {
			1
		} elif true {
			2
		} else {
			3
		}
	`

	expected := `
	push_false
	jf @elif

	push_int 1
	jmp @end

	@elif:
	push_true
	jf @else
	push_int 2
	jmp @end

	@else:
	push_int 3

	@end:
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestIfNoElse(t *testing.T) {
	src := `
		if false {
			1
		} elif true {
			2
		}
	`

	expected := `
	push_false
	jf @elif

	push_int 1
	jmp @end

	@elif:
	push_true
	jf @end
	push_int 2

	@end:
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestNestedIfs(t *testing.T) {
	src := `
	if true {
		if false {
			1
		} else {
			2
		}
	}
	`

	expected := `
	push_true
	jf @end

	push_false
	jf @else

	push_int 1
	jmp @end

	@else:
	push_int 2
	
	@end:
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestForCond(t *testing.T) {
	src := `
	for true {
		for false {
			1
		}
	}
	`

	expected := `
	@start:
	push_true
	jf @end

	@inner_start:
	push_false
	jf @inner_end

	push_int 1
	jmp @inner_start

	@inner_end:
	jmp @start

	@end
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestForIn(t *testing.T) {
	src := `
	for x in [1, 2, 3] {
		"say"(x)
	}
	`

	expected := `
	push_array

	get_local 0
	get_member "push"
	push_int 1
	call 1

	get_local 0
	get_member "push"
	push_int 2
	call 1

	get_local 0
	get_member "push"
	push_int 3
	call 1

	push_int 0
	push_null

	@check:
	get_local 1
	get_local 0
	get_member "length"
	call 0
	lt

	jf @end

	get_local 0
	get_local 1
	index
	set_local 2

	get_local 1
	push_int 1
	add
	set_local 1

	push_string "say"
	get_local 2
	call 1

	jmp @check
	@end:
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestSimpleFunction(t *testing.T) {
	src := `
	var result = the_meaning_of_life()

	fn the_meaning_of_life() {
		return 42
	}
	`

	expected := `
	push_function @the_meaning_of_life
	call 0
	halt

	@the_meaning_of_life:
	push_int 42
	ret
	`

	mustCompileTo(t, src, expected)
}

func TestFunctionArgs(t *testing.T) {
	src := `
	fn add(a, b) {
		return a + b
	}
	
	add(4, 5)
	`

	expected := `
	push_function @add
	set_arg_count 2
	push_int 4
	push_int 5
	call 2
	halt

	@add:
	get_local 0
	get_local 1
	add
	ret
	`

	mustCompileTo(t, src, expected)
}

func mustCompileTo(t *testing.T, src, expected string) {
	scanner := scanner.New(strings.NewReader(src))
	tokens, err := scanner.Scan()
	require.NoError(t, err)

	parser := parser.New(tokens)
	program, err := parser.Parse()
	require.NoError(t, err)

	langCompiler := compiler.New(program)
	testResult, err := langCompiler.Compile()
	require.NoError(t, err)

	textCompiler := text.NewCompiler(strings.NewReader(expected))
	expectedResult, err := textCompiler.Compile()
	require.NoError(t, err)

	if !assert.Equal(t, expectedResult.Code(), testResult.Code()) {
		d1 := text.NewDecompiler(expectedResult)
		d2 := text.NewDecompiler(testResult)
		require.Equal(t, d1.Decompile(), d2.Decompile())
	}
}