package org.zstack.identity;

import org.springframework.beans.factory.annotation.Autowired;
import org.zstack.core.Platform;
import org.zstack.core.db.EntityMetadata;
import org.zstack.header.identity.AccountConstant;
import org.zstack.header.identity.SessionInventory;
import org.zstack.header.zql.ASTNode;
import org.zstack.header.zql.MarshalZQLASTTreeExtensionPoint;
import org.zstack.header.zql.RestrictByExprExtensionPoint;
import org.zstack.header.zql.ZQLExtensionContext;
import org.zstack.zql.ZQLContext;
import org.zstack.zql.ast.ZQLMetadata;

public class IdentityZQLExtension implements MarshalZQLASTTreeExtensionPoint, RestrictByExprExtensionPoint {
    private static final String ENTITY_NAME = "__ACCOUNT_FILTER__";
    private static final String ENTITY_FIELD = "__ACCOUNT_FILTER_FIELD__";

    public static final String SKIP_IDENTITY_FILTER = "__SKIP_IDENTITY_FILTER __";

    @Autowired
    private AccountManager acntMgr;

    @Override
    public void marshalZQLASTTree(ASTNode.Query node) {
        SessionInventory session = ZQLContext.getAPISession();
        if (session == null) {
            return;
        }

        ASTNode.RestrictExpr expr = new ASTNode.RestrictExpr();
        expr.setEntity(ENTITY_NAME);
        expr.setField(ENTITY_FIELD);

        node.addRestrictExpr(expr);
    }

    @Override
    public String restrictByExpr(ZQLExtensionContext context, ASTNode.RestrictExpr expr) {
        if (!ENTITY_NAME.equals(expr.getEntity()) || !ENTITY_FIELD.equals(expr.getField())) {
            // not for us
            return null;
        }

        Boolean skip = (Boolean) ZQLContext.getCustomizedContext(SKIP_IDENTITY_FILTER);
        if (skip != null && skip) {
            throw new SkipThisRestrictExprException();
        }

        ZQLMetadata.InventoryMetadata src = ZQLMetadata.getInventoryMetadataByName(context.getQueryTargetInventoryName());
        String accountUuid = context.getAPISession().getAccountUuid();
        if (AccountConstant.INITIAL_SYSTEM_ADMIN_UUID.equals(accountUuid)) {
            throw new SkipThisRestrictExprException();
        }

        if (!acntMgr.isResourceHavingAccountReference(src.inventoryAnnotation.mappingVOClass())) {
            throw new SkipThisRestrictExprException();
        }

        String primaryKey = EntityMetadata.getPrimaryKeyField(src.inventoryAnnotation.mappingVOClass()).getName();

        return String.format("(%s.%s IN (SELECT accountresourcerefvo.resourceUuid FROM AccountResourceRefVO accountresourcerefvo WHERE" +
                        "  accountresourcerefvo.ownerAccountUuid = '%s' OR (accountresourcerefvo.resourceUuid" +
                        " IN (SELECT sharedresourcevo.resourceUuid FROM SharedResourceVO sharedresourcevo WHERE" +
                        " sharedresourcevo.receiverAccountUuid = '%s' OR sharedresourcevo.toPublic = 1))))",
                src.simpleInventoryName(), primaryKey, accountUuid, accountUuid);
    }
}

@@ -547,16 +547,104 @@
 adata;%0A%0A
+import java.util.ArrayList;%0Aimport java.util.List;%0Aimport java.util.stream.Collectors;%0A%0A
 public c
@@ -1397,24 +1397,131 @@
 pr);%0A    %7D%0A%0A
+    protected List%3CString%3E getRestrictAccountUuids(SessionInventory session) %7B%0A        return null;%0A    %7D%0A%0A
     @Overrid
@@ -1971,118 +1971,89 @@
     
-ZQLMetadata.InventoryMetadata src = ZQLMetadata.getInventoryMetadataByName(context.getQueryTargetInventoryName
+List%3CString%3E restrictAccountUuids = getRestrictAccountUuids(context.getAPISession
 ());
@@ -2197,16 +2197,48 @@
 untUuid)
+ && restrictAccountUuids == null
 ) %7B%0A    
@@ -2291,32 +2291,341 @@
 n();%0A        %7D%0A%0A
+        if (restrictAccountUuids == null) %7B%0A            restrictAccountUuids = new ArrayList%3C%3E();%0A        %7D%0A%0A        restrictAccountUuids.add(context.getAPISession().getAccountUuid());%0A%0A        ZQLMetadata.InventoryMetadata src = ZQLMetadata.getInventoryMetadataByName(context.getQueryTargetInventoryName());%0A
         if (!acn
@@ -2885,24 +2885,192 @@
 getName();%0A%0A
+        String accountStr = restrictAccountUuids.stream()%0A                .map(uuid -%3E String.format(%22'%25s'%22, uuid))%0A                .collect(Collectors.joining(%22,%22));%0A%0A
         retu
@@ -3256,30 +3256,31 @@
 AccountUuid 
-= '%25s'
+in (%25s)
  OR (account
@@ -3491,14 +3491,15 @@
 uid 
-= '%25s'
+in (%25s)
  OR 
@@ -3596,20 +3596,19 @@
  account
-Uuid
+Str
 , accoun
@@ -3608,20 +3608,19 @@
  account
-Uuid
+Str
 );%0A    %7D

Add function for addtional account restrict