package main

import (
	"errors"
	"strconv"
	"unicode"
)

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 safeASCIIIdentifier(identifier string) string {
	ascii := ""
	for _, rune := range identifier {
		if rune < unicode.MaxASCII && (unicode.IsLetter(rune) || unicode.IsDigit(rune)) {
			ascii += string(rune)
			continue
		}

		ascii += "$" + strconv.Itoa(int(rune))
	}

	return ascii
}

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:
		ass := expr.Value.(AssignmentExpression)

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

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

		local := strconv.Itoa(block.Locals[ass.Variable].Index)
		getLocal := "local.get $" + local + "\n"
		setLocal := "local.set $" + local + "\n"

		return exprWAT + cast + setLocal + getLocal, nil
	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_Binary:
		binary := expr.Value.(BinaryExpression)

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

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

		// TODO: cast produces unnecessary/wrong cast, make sure to upcast to target type
		castLeft, err := castPrimitiveWAT(binary.Left.ValueType.Value.(PrimitiveType), resultType)
		if err != nil {
			return "", err
		}

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

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

		op := ""

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

		switch binary.Operation {
		case Operation_Add:
			op = getPrimitiveWATType(resultType) + ".add\n"
		case Operation_Sub:
			op = getPrimitiveWATType(resultType) + ".sub\n"
		case Operation_Mul:
			op = getPrimitiveWATType(resultType) + ".mul\n"
		case Operation_Div:
			op = getPrimitiveWATType(resultType) + ".div_" + suffix + "\n"
		case Operation_Mod:
			op = getPrimitiveWATType(resultType) + ".rem_" + suffix + "\n"
		case Operation_Greater:
			op = getPrimitiveWATType(resultType) + ".gt_" + suffix + "\n"
		case Operation_Less:
			op = getPrimitiveWATType(resultType) + ".lt_" + suffix + "\n"
		case Operation_GreaterEquals:
			op = getPrimitiveWATType(resultType) + ".ge_" + suffix + "\n"
		case Operation_LessEquals:
			op = getPrimitiveWATType(resultType) + ".le_" + suffix + "\n"
		case Operation_NotEquals:
			op = getPrimitiveWATType(resultType) + ".ne\n"
		case Operation_Equals:
			op = getPrimitiveWATType(resultType) + ".eq\n"
		default:
			panic("operation not implemented")
		}

		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
		}

		// TODO: upcast to return type for non-primitive types

		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
		}

		// TODO: make sure to upcast to target type
		return wat + "local.set $" + strconv.Itoa(block.Locals[dlv.Variable].Index) + "\n", nil
	case Statement_If:
		ifS := stmt.Value.(IfStatement)

		conditionWAT, err := compileExpressionWAT(ifS.Condition, block)
		if err != nil {
			return "", err
		}

		condBlockWAT, err := compileBlockWAT(ifS.ConditionalBlock)
		if err != nil {
			return "", err
		}

		wat := ""

		if ifS.ElseBlock != nil {
			wat += "block\n"
		}

		// condition
		wat += "block\n"

		wat += conditionWAT
		wat += "i32.eqz\n" // logical not
		wat += "br_if 0\n"

		// condition is true
		wat += condBlockWAT

		if ifS.ElseBlock != nil {
			wat += "br 1\n" // jump over else block
		}

		wat += "end\n"

		if ifS.ElseBlock != nil {
			// condition is false
			elseWAT, err := compileBlockWAT(ifS.ElseBlock)
			if err != nil {
				return "", err
			}

			wat += elseWAT
			wat += "end\n"
		}

		return wat, nil
	}

	panic("stmt not implemented")
}

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 $" + safeASCIIIdentifier(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 $" + safeASCIIIdentifier(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
}