package vm_test

import (
	"jinx/pkg/lang/modules"
	"jinx/pkg/lang/vm"
	"jinx/pkg/lang/vm/code"
	"jinx/pkg/lang/vm/text"
	"os"
	"strings"
	"testing"

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

func TestSimpleSub(t *testing.T) {
	src := `
	push_int 2
	push_int 1
	sub
	`

	test(t, src, "1")
}

func TestGetLocal(t *testing.T) {
	src := `
	push_int 404
	push_int 2
	push_int 1
	add
	get_local 1
	`

	test(t, src, "3")
}

func TestFibonacci(t *testing.T) {
	src := `
	# Array stored in local 0
	push_array 

	get_local 0
	push_int 1
	temp_arr_push

	get_local 0
	push_int 1
	temp_arr_push

	@fib_loop:
	get_local 0
	temp_arr_len
	
	push_int 1

	# Index of the last element stored in local 1
	sub

	get_local 0
	get_local 0
	# This is the last element
	get_local 1
	# Store the last element in local 2
	index

	get_local 0

	get_local 1
	push_int 1
	# Index of the second last element stored in local 2
	sub

	# Store the second last element in local 3
	index

	add
	temp_arr_push

	# Drop local 1, which was the length of the array, which we no longer need
	drop 1

	get_local 0
	temp_arr_len

	push_int 10

	# Jump if the array is larger than 10.
	lte
	jt @fib_loop
	`

	test(t, src, "[1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89]")
}

func TestFunction(t *testing.T) {
	src := `
	push_function @subtract_two
	set_arg_count 1
	push_int 44
	call 1
	halt

	@subtract_two:
		push_int 2
		sub
		ret
	`

	test(t, src, "42")
}

func TestSimpleEnv(t *testing.T) {
	src := `
    push_int 404
    push_int 1
    push_int 405

    push_function @test
    # Add the local 1 to the environment.
    add_to_env 1
    call 0
    halt

    @test:
		get_env 0
        push_int 2
        add
        ret
    `

	test(t, src, "3")
}

func TestEscapedEnv(t *testing.T) {
	/*
	   fn create() {
	       var x = 0
	       return fn () {
	           x = x + 1
	           return x
	       }
	   }

	   var f = create()
	   var res = [f(), f(), f()]
	*/

	src := `    
    push_function @create
    call 0

    push_array

    get_local 1
    get_local 0
    call 0
    temp_arr_push

    get_local 1
    get_local 0
    call 0
    temp_arr_push
    
    get_local 1
    get_local 0
    call 0
    temp_arr_push
    halt

    @create:
        push_int 0
        push_int 404
        push_int 405
        push_function @create:anon_0
        add_to_env 0
        ret
    @create:anon_0:
		get_env 0
        push_int 1
        add
        set_env 0
        get_env 0
        ret
    `

	test(t, src, "[1, 2, 3]")
}

func TestEnvNotEscaped(t *testing.T) {
	/*
		var x = 1

		fn f() {
		    x = x + 1
		    return
		}

		f()
		var result = x
	*/

	src := `
	push_int 1

	push_function @f
	add_to_env 0
	get_local 1

	call 0
	drop 1

	get_local 0
	halt

	@f:
	get_env 0
	push_int 1
	add
	set_env 0

	push_null
	ret
	`

	test(t, src, "2")
}

func TestMember(t *testing.T) {
	src := `
	push_array

	get_local 0
	push_int 1
	temp_arr_push

	get_local 0
	push_int 2
	temp_arr_push

	get_local 0
	push_int 3
	temp_arr_push

	get_member "length"
	call 0
	`

	test(t, src, "3")
}

func TestObject(t *testing.T) {
	src := `
	push_object

	get_local 0
	push_string "Petronij"
	set_member "name"

	get_local 0
	push_int 10
	set_member "age"

	get_local 0
	get_member "name"
	`

	test(t, src, "\"Petronij\"")
}

func TestTypeConstruct(t *testing.T) {
	/*
		type Cat {
			(name, age) {
				this.name = name
				this.age = age
			}

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

		var cat = Cat("Kitty", 3)
		cat.meow()
	*/

	src := `
	# Create a new type Cat
	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
	
	# Create a new instance of Cat
	get_local 0
	get_member "$init"
	push_string "Kitty"
	push_int 3
	call 2

	# Call the meow method
	get_local 1
	get_member "meow"
	call 0
	halt

	@Cat:$init:
		push_object
		get_env 0
		anchor_type

		get_local 2
		get_local 0
		set_member "name"

		get_local 2
		get_local 1
		set_member "age"

		ret
	@Cat:meow:
		get_env 0
		get_member "name"
	
		push_string " says Meow!"

		add
		ret
	`

	test(t, src, "\"Kitty says Meow!\"")
}

func TestPrimes(t *testing.T) {
	/*
		var i = 2
		var primes = []

		for i <= 10 {
			var factors = []

			var j = 1
			for j <= i {
				if i % j <= 0 {
					factors.push(j)
				}
				j = j + 1
			}

			if factors.length() <= 2 {
				primes.push(i)
			}

			i = i + 1
		}

		print(primes)
	*/

	// TODO: This test uses OpLte instead of OpEq, which works,
	// but once OpEq is implemented, this test should be updated.

	src := `
	push_int 2
	push_array
	@main_loop:
		get_local 0
		push_int 10
		lte
		jf @end

		push_array
		push_int 1

	@factor_loop:
		get_local 3
		get_local 0
		lte
		jf @prime_check

		
		get_local 0
		get_local 3
		mod
		
		push_int 0

		lte
		jf @factor_loop_inc

		get_local 2
		get_local 3
		temp_arr_push

	@factor_loop_inc:
		get_local 3
		push_int 1
		add
		set_local 3
		jmp @factor_loop

	@prime_check:
		get_local 2
		get_member "length"
		call 0
		push_int 2
		lte

		jf @main_loop_inc

		get_local 1
		get_local 0
		temp_arr_push

	@main_loop_inc:
		get_local 0
		push_int 1
		add
		set_local 0

		drop 2

		jmp @main_loop
	@end:
		halt
	`

	test(t, src, "[2, 3, 5, 7]")
}

func TestModules(t *testing.T) {
	t.Skip("We can't pass in random modules anymore, they have to go through the resolver")

	// mainSrc := `
	// get_global "add_one"
	// get_global "x"
	// call 1
	// halt
	// `

	// librarySrc1 := `
	// push_int 41
	// add_global "x"
	// halt
	// `

	// librarySrc2 := `
	// push_function @add_one
	// set_arg_count 1
	// add_global "add_one"
	// halt

	// @add_one:
	// get_local 0
	// push_int 1
	// add
	// ret
	// `

	// main := compile(t, mainSrc)
	// library1 := compile(t, librarySrc1)
	// library2 := compile(t, librarySrc2)

	// vm := vm.New(&main, []*code.Code{&library1, &library2})
	// err := vm.Run()
	// require.NoError(t, err)

	// res, err := vm.GetResult()
	// require.NoError(t, err)

	// require.Equal(t, "42", res)
}

func TestCoreSay(t *testing.T) {
	src := `
	get_global ":core:say"
	push_string "Meow!!"
	call 1
	push_int 0
	halt
	`

	previous := os.Stdout
	r, w, _ := os.Pipe()
	defer w.Close()
	os.Stdout = w
	test(t, src, "0")
	os.Stdout = previous

	result := make([]byte, 8)
	_, err := r.Read(result)
	require.NoError(t, err)

	require.Equal(t, "\"Meow!!\"", string(result))
}

func test(t *testing.T, src string, expected string) {
	bc := compile(t, src)
	vm := vm.New(modules.NewUnknownModule(&bc))
	err := vm.Run()
	require.NoError(t, err)

	res, err := vm.GetResult()
	require.NoError(t, err)

	require.Equal(t, expected, res)
}

func compile(t *testing.T, src string) code.Code {
	comp := text.NewCompiler(strings.NewReader(src))
	bc, err := comp.Compile()
	require.NoErrorf(t, err, "compilation from text format failed: %s", err)

	return bc
}