package lang

import "strings"

func normalizeClassReferences(program *Program) {
	classes := map[string]bool{}
	for _, class := range program.Classes {
		classes[class.Name] = true
	}
	normalize := func(value string) string { return normalizeReferenceType(value, classes) }
	for i := range program.Interfaces {
		for j := range program.Interfaces[i].Methods {
			normalizeSignature(&program.Interfaces[i].Methods[j], normalize)
		}
	}
	for i := range program.Enums {
		for j := range program.Enums[i].Variants {
			for k := range program.Enums[i].Variants[j].PayloadTypes {
				program.Enums[i].Variants[j].PayloadTypes[k] = normalize(program.Enums[i].Variants[j].PayloadTypes[k])
			}
		}
	}
	for i := range program.Classes {
		for j := range program.Classes[i].Fields {
			program.Classes[i].Fields[j].Type = normalize(program.Classes[i].Fields[j].Type)
		}
		for j := range program.Classes[i].Methods {
			normalizeFunction(&program.Classes[i].Methods[j], normalize)
		}
	}
	for i := range program.Workers {
		for j := range program.Workers[i].Fields {
			program.Workers[i].Fields[j].Type = normalize(program.Workers[i].Fields[j].Type)
			normalizeExprTypes(program.Workers[i].Fields[j].Value, normalize)
		}
		for j := range program.Workers[i].Methods {
			normalizeFunction(&program.Workers[i].Methods[j], normalize)
		}
	}
	for i := range program.Functions {
		normalizeFunction(&program.Functions[i], normalize)
	}
}

func normalizeSignature(signature *FunctionSignature, normalize func(string) string) {
	for i := range signature.Params {
		signature.Params[i].Type = normalize(signature.Params[i].Type)
	}
	signature.ReturnType = normalize(signature.ReturnType)
}

func normalizeFunction(function *FunctionDecl, normalize func(string) string) {
	for i := range function.Params {
		function.Params[i].Type = normalize(function.Params[i].Type)
	}
	function.ReturnType = normalize(function.ReturnType)
	normalizeStmtTypes(function.Body, normalize)
}

func normalizeStmtTypes(statements []Stmt, normalize func(string) string) {
	for index, statement := range statements {
		switch value := statement.(type) {
		case VarDecl:
			value.Type = normalize(value.Type)
			normalizeExprTypes(value.Value, normalize)
			statements[index] = value
		case MultiVarDecl:
			normalizeExprTypes(value.Value, normalize)
		case AssignStmt:
			normalizeExprTypes(value.Value, normalize)
		case AddAssignStmt:
			normalizeExprTypes(value.Value, normalize)
		case MultiAssignStmt:
			normalizeExprTypes(value.Value, normalize)
		case ReturnStmt:
			if value.Value != nil {
				normalizeExprTypes(value.Value, normalize)
			}
		case ThrowStmt:
			normalizeExprTypes(value.Value, normalize)
		case GoStmt:
			normalizeExprTypes(value.Value, normalize)
		case DeferStmt:
			normalizeExprTypes(value.Value, normalize)
		case ExprStmt:
			normalizeExprTypes(value.Value, normalize)
		case IfStmt:
			normalizeExprTypes(value.Cond, normalize)
			normalizeStmtTypes(value.Then, normalize)
			normalizeStmtTypes(value.Else, normalize)
		case WhileStmt:
			normalizeExprTypes(value.Cond, normalize)
			normalizeStmtTypes(value.Body, normalize)
		case ForEachStmt:
			normalizeExprTypes(value.Source, normalize)
			normalizeStmtTypes(value.Body, normalize)
		case SelectStmt:
			for _, c := range value.Cases {
				normalizeExprTypes(c.Source, normalize)
				normalizeStmtTypes(c.Body, normalize)
			}
		case MatchStmt:
			normalizeExprTypes(value.Value, normalize)
			for _, c := range value.Cases {
				normalizeStmtTypes(c.Body, normalize)
			}
		case TryCatchStmt:
			value.CatchType = normalize(value.CatchType)
			normalizeStmtTypes(value.TryBody, normalize)
			normalizeStmtTypes(value.CatchBody, normalize)
			statements[index] = value
		}
	}
}

func normalizeExprTypes(expression Expr, normalize func(string) string) {
	switch value := expression.(type) {
	case UnaryExpr:
		normalizeExprTypes(value.Value, normalize)
	case NonNullExpr:
		normalizeExprTypes(value.Value, normalize)
	case BinaryExpr:
		normalizeExprTypes(value.Left, normalize)
		normalizeExprTypes(value.Right, normalize)
	case SelectorExpr:
		normalizeExprTypes(value.Receiver, normalize)
	case SafeSelectorExpr:
		normalizeExprTypes(value.Receiver, normalize)
	case IndexExpr:
		normalizeExprTypes(value.Receiver, normalize)
		normalizeExprTypes(value.Index, normalize)
	case EnumVariantExpr:
		for _, item := range value.Values {
			normalizeExprTypes(item, normalize)
		}
	case MatchExpr:
		normalizeExprTypes(value.Value, normalize)
		for _, item := range value.Cases {
			normalizeExprTypes(item.Value, normalize)
		}
	case LambdaExpr:
		for i := range value.Params {
			value.Params[i].Type = normalize(value.Params[i].Type)
		}
		normalizeStmtTypes(value.Body, normalize)
	case CallExpr:
		skipTypeArgs := false
		if selector, ok := value.Callee.(SelectorExpr); ok {
			if root, ok := selector.Receiver.(IdentExpr); ok && root.Name == "sql" {
				skipTypeArgs = true
			}
		}
		if !skipTypeArgs {
			for i := range value.TypeArgs {
				value.TypeArgs[i] = normalize(value.TypeArgs[i])
			}
		}
		normalizeExprTypes(value.Callee, normalize)
		for _, item := range value.Args {
			normalizeExprTypes(item, normalize)
		}
		for _, item := range value.NamedArgs {
			normalizeExprTypes(item.Value, normalize)
		}
	}
}

func normalizeReferenceType(value string, classes map[string]bool) string {
	if value == "" {
		return value
	}
	nullable := strings.HasSuffix(value, "?")
	if nullable {
		value = strings.TrimSuffix(value, "?")
	}
	explicitPointer := strings.HasPrefix(value, "*")
	if explicitPointer {
		value = strings.TrimPrefix(value, "*")
	}
	if params, result, ok := parseFunctionType(value); ok {
		for i := range params {
			params[i] = normalizeReferenceType(params[i], classes)
		}
		value = "(" + strings.Join(params, ", ") + ") -> " + normalizeReferenceType(result, classes)
	} else if base, args, ok := parseGenericType(value); ok {
		for i := range args {
			args[i] = normalizeReferenceType(args[i], classes)
		}
		value = base + "<" + strings.Join(args, ", ") + ">"
	} else if classes[value] || explicitPointer {
		value = "*" + value
	}
	if nullable {
		value += "?"
	}
	return value
}