Refactor customer form and service for company ID validation.

internal/customer/service.go

@@ -100,23 +100,13 @@ func (s *customerService) Register(ctx context.Context, form *CreateCustomerForm
 		return nil, errors.New("invalid username format")
 	}
 
-	// Verify that all company IDs exist
 	var Companies []string
 	if len(form.Companies) > 0 {
-		for _, companyID := range form.Companies {
-			if _, err := bson.ObjectIDFromHex(companyID); err != nil {
-				return nil, errors.New("invalid company ID format: " + companyID)
-			}
-			_, err := s.companyService.GetByID(ctx, companyID)
-			if err != nil {
-				s.logger.Error("Invalid company ID provided", map[string]interface{}{
-					"error":      err.Error(),
-					"company_id": companyID,
-				})
-				return nil, errors.New("invalid company ID: " + companyID)
-			}
+		validated, err := s.validateCompanyIDs(ctx, form.Companies)
+		if err != nil {
+			return nil, err
 		}
-		Companies = form.Companies
+		Companies = validated
 	}
 
 	// Hash password
@@ -362,36 +352,21 @@ func (s *customerService) Update(ctx context.Context, id string, form *UpdateCus
 		changeUsername = true
 	}
 
-	// Handle company assignments
-	companies := make([]*CompanySummary, 0)
 	companiesChanged := false
-	if len(form.Companies) > 0 {
-		// Check if companies changed
-		if len(form.Companies) != len(customer.Companies) {
-			companiesChanged = true
-		} else {
-			companyMap := make(map[string]bool)
-			for _, id := range customer.Companies {
-				companyMap[id] = true
-			}
-			for _, id := range form.Companies {
-				if !companyMap[id] {
-					companiesChanged = true
-					break
-				}
-			}
-		}
-		customer.Companies = form.Companies
-		companies, err = s.loadCompaniesForCustomer(ctx, customer)
-		if err != nil {
-			s.logger.Error("Failed to load companies for customer", map[string]interface{}{
-				"error":       err.Error(),
-				"customer_id": customer.ID,
-			})
+	if form.Companies != nil {
+		companyIDs, validateErr := s.validateCompanyIDs(ctx, *form.Companies)
+		if validateErr != nil {
+			return nil, validateErr
 		}
+		companiesChanged = !companyIDsEqual(customer.Companies, companyIDs)
+		customer.Companies = companyIDs
 	}
 
 	infoChanged := false
+	if form.Type != "" && customer.Type != CustomerType(form.Type) {
+		customer.Type = CustomerType(form.Type)
+		infoChanged = true
+	}
 	if form.FullName != nil && (customer.FullName == nil || *form.FullName != *customer.FullName) {
 		customer.FullName = form.FullName
 		infoChanged = true
@@ -458,7 +433,16 @@ func (s *customerService) Update(ctx context.Context, id string, form *UpdateCus
 		"email":       customer.Email,
 	})
 
-	return customer.ToResponse(companies), nil
+	responseCompanies, err := s.loadCompaniesForCustomer(ctx, customer)
+	if err != nil {
+		s.logger.Error("Failed to load companies for customer response", map[string]interface{}{
+			"error":       err.Error(),
+			"customer_id": customer.ID,
+		})
+		responseCompanies = []*CompanySummary{}
+	}
+
+	return customer.ToResponse(responseCompanies), nil
 }
 
 // Delete deletes a customer (soft delete)
@@ -1269,6 +1253,61 @@ func (s *customerService) isValidPassword(password string) bool {
 	return hasUpper && hasLower && hasSpecial
 }
 
+func (s *customerService) validateCompanyIDs(ctx context.Context, companyIDs []string) ([]string, error) {
+	if len(companyIDs) == 0 {
+		return []string{}, nil
+	}
+
+	validated := make([]string, 0, len(companyIDs))
+	seen := make(map[string]struct{}, len(companyIDs))
+	for _, companyID := range companyIDs {
+		if companyID == "" {
+			continue
+		}
+		if _, exists := seen[companyID]; exists {
+			continue
+		}
+		if _, err := bson.ObjectIDFromHex(companyID); err != nil {
+			return nil, errors.New("invalid company ID format: " + companyID)
+		}
+		if _, err := s.companyService.GetByID(ctx, companyID); err != nil {
+			s.logger.Error("Invalid company ID provided", map[string]interface{}{
+				"error":      err.Error(),
+				"company_id": companyID,
+			})
+			return nil, errors.New("invalid company ID: " + companyID)
+		}
+		seen[companyID] = struct{}{}
+		validated = append(validated, companyID)
+	}
+
+	return validated, nil
+}
+
+func companyIDsEqual(existing, updated []string) bool {
+	if len(existing) != len(updated) {
+		return false
+	}
+
+	counts := make(map[string]int, len(existing))
+	for _, id := range existing {
+		counts[id]++
+	}
+	for _, id := range updated {
+		counts[id]--
+		if counts[id] < 0 {
+			return false
+		}
+	}
+	for _, count := range counts {
+		if count != 0 {
+			return false
+		}
+	}
+
+	return true
+}
+
 // loadCompaniesForCustomer loads company summaries for a customer
 func (s *customerService) loadCompaniesForCustomer(ctx context.Context, customer *Customer) ([]*CompanySummary, error) {
 	if len(customer.Companies) == 0 {