package com.bonus.canteen.core.starter.threadpool;

import com.bonus.canteen.core.common.utils.TenantContextHolder;
import com.bonus.canteen.core.config.WebContext;
import com.bonus.canteen.core.dataset.rule.RoutingRule;
import com.github.pagehelper.PageHelper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.task.TaskDecorator;

public class CleanTenantIdTaskDecorator implements TaskDecorator {
   private static final Logger log = LoggerFactory.getLogger(CleanTenantIdTaskDecorator.class);

   @Override
   public Runnable decorate(Runnable runnable) {
      Long tenantId = TenantContextHolder.getTenantId();
      RoutingRule.Context routeContext = RoutingRule.getContext().clone();
      WebContext.Context context = WebContext.get();
      String invokeThreadName = Thread.currentThread().getName();
      return () -> {
         if (invokeThreadName.equals(Thread.currentThread().getName())) {
            runnable.run();
         } else {
            try {
               TenantContextHolder.setTenantId(tenantId);
               RoutingRule.setContext(routeContext);
               WebContext.set(context);
               runnable.run();
            } finally {
               TenantContextHolder.clear();
               RoutingRule.reset();
               WebContext.reset();
               PageHelper.clearPage();
            }

         }
      };
   }
}
