package main

import (
	"errors"
	"log"
	"strconv"
)

func getPrimitiveWATType(primitive PrimitiveType) string {
	switch primitive {
	case Primitive_I8, Primitive_I16, Primitive_I32, Primitive_U8, Primitive_U16, Primitive_U32:
		return "i32"
	case Primitive_I64, Primitive_U64:
		return "i64"
	case Primitive_F32:
		return "f32"
	case Primitive_F64:
		return "f64"
	case Primitive_Bool:
		return "i32"
	}

	panic("unhandled type")
}

func getWATType(t Type) string {
	// TODO: tuples?

	if t.Type != Type_Primitive {
		panic("not implemented") // TODO: non-primitive types
	}

	primitive := t.Value.(PrimitiveType)
	return getPrimitiveWATType(primitive)
}

func getTypeCast(primitive PrimitiveType) string {
	switch primitive {
	case Primitive_I8:
		return "i32.extend8_s\n"
	case Primitive_U8:
		return "i32.const 255\ni32.and\n"
	case Primitive_I16:
		return "i32.extend16_s\n"
	case Primitive_U16:
		return "i32.const 65535\ni32.and\n"
	case Primitive_Bool:
		return "i32.const 1\ni32.and\n"
	}

	return ""
}

func pushConstantNumberWAT(primitive PrimitiveType, value any) string {
	switch primitive {
	case Primitive_I8, Primitive_I16, Primitive_I32:
		return "i32.const " + strconv.FormatInt(value.(int64), 10) + "\n"
	case Primitive_U8, Primitive_U16, Primitive_U32:
		return "i32.const " + strconv.FormatUint(value.(uint64), 10) + "\n"
	case Primitive_I64:
		return "i64.const " + strconv.FormatInt(value.(int64), 10) + "\n"
	case Primitive_U64:
		return "i64.const " + strconv.FormatUint(value.(uint64), 10) + "\n"
	case Primitive_F32:
		return "f32.const " + strconv.FormatFloat(value.(float64), 'f', -1, 32) + "\n"
	case Primitive_F64:
		return "f64.const " + strconv.FormatFloat(value.(float64), 'f', -1, 64) + "\n"
	}

	panic("invalid type")
}

func castPrimitiveWAT(from PrimitiveType, to PrimitiveType) (string, error) {
	if from == to {
		return "", nil
	}

	if from == Primitive_Bool || to == Primitive_Bool {
		return "", errors.New("cannot upcast from or to bool")
	}

	fromFloat := isFloatingPoint(from)
	toFloat := isFloatingPoint(to)
	if fromFloat && toFloat {
		if to == Primitive_F32 {
			return "f32.demote_f64\n", nil
		} else {
			return "f64.promote_f32\n", nil
		}
	}

	if toFloat {
		suffix := ""
		if isUnsignedInt(to) {
			suffix = "u"
		} else {
			suffix = "s"
		}

		return getPrimitiveWATType(to) + ".convert_" + getPrimitiveWATType(from) + "_" + suffix + "\n", nil
	}

	if fromFloat {
		suffix := ""

		if isUnsignedInt(to) {
			suffix = "u"
		} else {
			suffix = "s"
		}

		return getPrimitiveWATType(to) + ".trunc_" + getPrimitiveWATType(from) + "_" + suffix + "\n", nil
	}

	if getBits(from) == getBits(to) {
		return "", nil
	}

	if getPrimitiveWATType(from) == getPrimitiveWATType(to) {
		if getBits(to) < getBits(from) {
			return getTypeCast(to), nil
		}
	}

	if getBits(from) < 64 && getBits(to) == 64 {
		suffix := ""
		if isUnsignedInt(from) {
			suffix = "u"
		} else {
			suffix = "s"
		}

		return "i64.extend_i32_" + suffix + "\n", nil
	}

	return "i32.wrap_i64\n" + getTypeCast(to), nil
}

func compileExpressionWAT(expr Expression, block Block) (string, error) {
	var err error

	switch expr.Type {
	case Expression_Assignment:
		// TODO
	case Expression_Literal:
		lit := expr.Value.(LiteralExpression)
		switch lit.Literal.Type {
		case Literal_Number:
			return pushConstantNumberWAT(lit.Literal.Primitive, lit.Literal.Value), nil
		case Literal_Boolean:
			if lit.Literal.Value.(bool) {
				return "i32.const 1\n", nil
			} else {
				return "i32.const 0\n", nil
			}
		case Literal_String:
			panic("not implemented")
		}
	case Expression_VariableReference:
		ref := expr.Value.(VariableReferenceExpression)

		cast := ""
		if expr.ValueType.Type == Type_Primitive {
			cast = getTypeCast(expr.ValueType.Value.(PrimitiveType))
		}

		return "local.get $" + strconv.Itoa(block.Locals[ref.Variable].Index) + "\n" + cast, nil
	case Expression_Arithmetic:
		arith := expr.Value.(ArithmeticExpression)

		log.Printf("%+#v", arith)

		// TODO: currently fine, only allowed for primitive types, but should be expanded to allow e.g. strings
		exprType := expr.ValueType.Value.(PrimitiveType)

		watLeft, err := compileExpressionWAT(arith.Left, block)
		if err != nil {
			return "", err
		}

		castLeft, err := castPrimitiveWAT(arith.Left.ValueType.Value.(PrimitiveType), exprType)
		if err != nil {
			return "", err
		}

		watRight, err := compileExpressionWAT(arith.Right, block)
		if err != nil {
			return "", err
		}

		castRight, err := castPrimitiveWAT(arith.Right.ValueType.Value.(PrimitiveType), exprType)
		if err != nil {
			return "", err
		}

		op := ""

		suffix := ""
		if isUnsignedInt(exprType) {
			suffix = "u"
		} else {
			suffix = "s"
		}

		switch arith.Operation {
		case Arithmetic_Add:
			op = getPrimitiveWATType(exprType) + ".add\n"
		case Arithmetic_Sub:
			op = getPrimitiveWATType(exprType) + ".sub\n"
		case Arithmetic_Mul:
			op = getPrimitiveWATType(exprType) + ".mul\n"
		case Arithmetic_Div:
			op = getPrimitiveWATType(exprType) + ".div_" + suffix + "\n"
		case Arithmetic_Mod:
			op = getPrimitiveWATType(exprType) + ".rem_" + suffix + "\n"
		}

		return watLeft + castLeft + watRight + castRight + op + getTypeCast(exprType), nil
	case Expression_Tuple:
		tuple := expr.Value.(TupleExpression)

		wat := ""
		for _, member := range tuple.Members {
			memberWAT, err := compileExpressionWAT(member, block)
			if err != nil {
				return "", err
			}

			wat += memberWAT
		}

		return wat, nil
	case Expression_FunctionCall:
		fc := expr.Value.(FunctionCallExpression)

		wat := ""
		if fc.Parameters != nil {
			wat, err = compileExpressionWAT(*fc.Parameters, block)
			if err != nil {
				return "", err
			}
		}

		return wat + "call $" + fc.Function + "\n", nil
	case Expression_Negate:
		neg := expr.Value.(NegateExpression)
		exprType := expr.ValueType.Value.(PrimitiveType)

		wat, err := compileExpressionWAT(neg.Value, block)
		if err != nil {
			return "", err
		}

		watType := getPrimitiveWATType(exprType)
		if isSignedInt(exprType) || isUnsignedInt(exprType) {
			return watType + ".const 0\n" + wat + watType + ".sub\n", nil
		}

		if isFloatingPoint(exprType) {
			return watType + ".neg\n", nil
		}
	}

	panic("expr not implemented")
}

func compileStatementWAT(stmt Statement, block Block) (string, error) {
	switch stmt.Type {
	case Statement_Expression:
		expr := stmt.Value.(ExpressionStatement)
		wat, err := compileExpressionWAT(expr.Expression, block)
		if err != nil {
			return "", err
		}

		return wat + "drop\n", nil
	case Statement_Block:
		block := stmt.Value.(BlockStatement)
		wat, err := compileBlockWAT(block.Block)
		if err != nil {
			return "", err
		}

		return wat, nil
	case Statement_Return:
		ret := stmt.Value.(ReturnStatement)
		wat, err := compileExpressionWAT(*ret.Value, block)
		if err != nil {
			return "", err
		}

		return wat + "return\n", nil
	case Statement_DeclareLocalVariable:
		dlv := stmt.Value.(DeclareLocalVariableStatement)
		if dlv.Initializer == nil {
			return "", nil
		}

		wat, err := compileExpressionWAT(*dlv.Initializer, block)
		if err != nil {
			return "", err
		}

		return wat + "local.set $" + strconv.Itoa(block.Locals[dlv.Variable].Index) + "\n", nil
	}

	return "", nil
}

func compileBlockWAT(block Block) (string, error) {
	blockWAT := ""

	for _, stmt := range block.Statements {
		wat, err := compileStatementWAT(stmt, block)
		if err != nil {
			return "", err
		}

		blockWAT += wat
	}

	return blockWAT, nil
}

func compileFunctionWAT(function ParsedFunction) (string, error) {
	funcWAT := "(func $" + function.Name + "\n"

	for _, local := range function.Locals {
		if !local.IsParameter {
			continue
		}
		funcWAT += "\t(param $" + strconv.Itoa(local.Index) + " " + getWATType(local.Type) + ")\n"
	}

	// TODO: tuples
	returnTypes := []Type{}
	if function.ReturnType != nil {
		returnTypes = []Type{*function.ReturnType}
		if function.ReturnType.Type == Type_Tuple {
			returnTypes = function.ReturnType.Value.(TupleType).Types
		}
	}

	for _, t := range returnTypes {
		funcWAT += "\t(result " + getWATType(t) + ")\n"
	}

	for _, local := range function.Locals {
		if local.IsParameter {
			continue
		}
		funcWAT += "\t(local $" + strconv.Itoa(local.Index) + " " + getWATType(local.Type) + ")\n"
	}

	wat, err := compileBlockWAT(function.Body)
	if err != nil {
		return "", err
	}

	funcWAT += wat

	return funcWAT + ") (export \"" + function.Name + "\" (func $" + function.Name + "))\n", nil
}

func backendWAT(file ParsedFile) (string, error) {
	module := "(module (memory 1)\n"

	for _, function := range file.Functions {
		wat, err := compileFunctionWAT(function)
		if err != nil {
			return "", err
		}

		module += wat
	}

	module += ")"
	return module, nil
}