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
	drop 1
	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
	drop 1
	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

	drop 1

	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_int 1
	push_int 5
	lte

	jf @end

	push_string "hello "
	push_string "world"
	add

	drop 1

	@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
	drop 1
	jmp @end

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

	@else:
	push_int 3
	drop 1

	@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
	drop 1
	jmp @end

	@elif:
	push_true
	jf @end
	push_int 2
	drop 1

	@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
	drop 1
	jmp @end

	@else:
	push_int 2
	drop 1
	
	@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
	drop 1
	jmp @inner_start

	@inner_end:
	jmp @start

	@end
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestForCondBreakContinue(t *testing.T) {
	src := `
	var sum = 0
	var x = 0
	for {
		if x % 2 == 0 {
			continue
		}

		sum = sum + x

		if x == 100 {
			break
		}
	}
	`

	expected := `
	push_int 0
	push_int 0

	@continue:
	get_local 1
	push_int 2
	mod
	push_int 0
	eq
	jf @sum
	jmp @continue

	@sum:
	get_local 0
	get_local 1
	add
	set_local 0

	get_local 1
	push_int 100
	eq
	jf @repeat
	jmp @break

	@repeat:
	jmp @continue

	@break:
	halt
	`

	mustCompileTo(t, src, expected)
}

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

		"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

	get_local 0
	get_member "push"
	push_string "oops"
	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

	get_local 2
	push_string "oops"
	eq
	jf @say

	jmp @end

	@say:
	push_string "say"
	get_local 2
	call 1
	drop 1

	jmp @check
	@end:
	drop 3
	halt
	`

	mustCompileTo(t, src, expected)
}

func TestSimpleFunction(t *testing.T) {
	src := `
	fn the_meaning_of_life() {
		return 42
	}

	var result = the_meaning_of_life()
	`

	expected := `
	push_function @the_meaning_of_life

	get_local 0
	call 0
	halt

	@the_meaning_of_life:
	push_int 42
	ret
	`

	mustCompileTo(t, src, expected)
}

func TestFunctionArgs(t *testing.T) {
	// TODO: Are arguments in the correct order?
	src := `
	fn add(a, b) {
		return a + b
	}
	
	add(4, 5)
	`

	expected := `
	push_function @add
	set_arg_count 2

	get_local 0
	push_int 4
	push_int 5
	call 2
	drop 1
	halt

	@add:
	get_local 1
	get_local 0
	add
	ret
	`

	mustCompileTo(t, src, expected)
}

func TestClosureEnv(t *testing.T) {
	src := `
	fn create() {
		var x = 0
		fn closure() {
			x = x + 1
			return x
		}

		return closure
	}
	`

	expected := `
	push_function @create
	halt

	@closure:
	get_env 0
	push_int 1
	add
	set_env 0
	get_env 0
	ret

	@create:
	push_int 0
	push_function @closure
	add_to_env 0

	get_local 1
	ret
	`

	mustCompileTo(t, src, expected)
}

func TestType(t *testing.T) {
	src := `
	type Cat {
		(name, age) {
			this.name = name
			this.age = age
		}

		fn meow(this) {
			return this.name + " says Meow!"
		}
	}

	var kitty = Cat("Kitty", 3)
	kitty.meow()
	`

	expected := `
	push_type "Cat"

	get_local 0
	get_member "$add_method"

	push_string "$init"
	push_function @Cat:$init
	set_arg_count 2

	call 2
	drop 1

	get_local 0
	get_member "$add_method"

	push_string "meow"
	push_function @Cat:meow

	call 2
	drop 1

	get_local 0
	push_string "Kitty"
	push_int 3
	call 2

	get_local 1
	get_member "meow"
	call 0
	drop 1

	halt

	@Cat:$init:
		push_object
		get_env 0
		anchor_type

		get_local 2
		get_local 1
		set_member "name"

		get_local 2
		get_local 0
		set_member "age"

		get_local 2
		ret
	@Cat:meow:
		get_env 0
		get_member "name"

		push_string " says Meow!"

		add
		ret
	`

	mustCompileTo(t, src, expected)
}

func TestGlobals(t *testing.T) {
	src := `
	use mul from hello by mel
	use pi from math

	var result = mul(pi, 2)
	`

	expected := `
	get_global "mel:hello:mul"
	get_global ":math:pi"
	push_int 2
	call 2
	halt
	`

	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("test_program", "test", 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().Code()) {
		d1 := text.NewDecompiler(expectedResult)
		d2 := text.NewDecompiler(*testResult.Code())
		require.Equal(t, d1.Decompile(), d2.Decompile())
	}
}