/*
 * JBoss, Home of Professional Open Source.
 * See the COPYRIGHT.txt file distributed with this work for information
 * regarding copyright ownership.  Some portions may be licensed
 * to Red Hat, Inc. under one or more contributor license agreements.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301 USA.
 * Modified by
 * @author  Data Virtuality GmbH
 * @date    $Date$
 *
 */
// https://issues.jboss.org/browse/TEIID-247, 2166, 2151, 2234, 2581, 2571, 2369, 1386, 2226, 2233, 2224, 2225, 2257
// https://issues.jboss.org/browse/TEIID-2175
// https://github.com/teiid/teiid/commit/ac4d4be06719f8f7757870eae37ec7bc9de44a08
// https://issues.jboss.org/browse/TEIID-2471
// https://issues.jboss.org/browse/TEIID-1092
// https://issues.jboss.org/browse/TEIID-2507
// https://issues.jboss.org/browse/TEIID-2534
// https://issues.jboss.org/browse/TEIID-2750
// https://issues.jboss.org/browse/TEIID-2402
// https://issues.redhat.com/browse/TEIID-2384
// https://issues.redhat.com/browse/TEIID-5480
// https://issues.redhat.com/browse/TEIID-2577
// https://issues.redhat.com/browse/TEIID-2732
// https://issues.redhat.com/browse/TEIID-3664
package org.teiid.dqp.internal.process;

import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.Cache;
import java.util.concurrent.locks.ReentrantLock;

import com.datavirtuality.dv.core.interfaces.concurrent.RefreshSemaphore;
import com.datavirtuality.dv.core.services.DVServices;
import org.teiid.adminapi.DataPolicy;
import org.teiid.adminapi.impl.SessionMetadata;
import org.teiid.adminapi.impl.VDBMetaData;
import org.teiid.client.BatchSerializer;
import org.teiid.client.security.SessionToken;
import org.teiid.core.util.PropertiesUtils;
import org.teiid.dqp.message.RequestID;
import org.teiid.jdbc.EmbeddedProfile;
import org.teiid.logging.LogManager;
import org.teiid.query.function.FunctionMethods;
import org.teiid.query.metadata.TransformationMetadata;
import org.teiid.query.util.DVUsersHandlerIntf;
import org.teiid.security.SecurityHelper;

import java.io.Serializable;
import java.util.*;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.FutureTask;
import org.wildfly.security.auth.server.SecurityIdentity;
import org.wildfly.security.authz.Roles;

public class DQPWorkContext implements Serializable {

	private static final String TEIID_VDB = "teiid-vdb"; //$NON-NLS-1$

	private static final String TEIID_SESSION = "teiid-session"; //$NON-NLS-1$

	private static final long serialVersionUID = -6389893410233192977L;

	private static final boolean longDatesTimes = PropertiesUtils.getBooleanProperty(System.getProperties(), "org.teiid.longDatesTimes", false); //$NON-NLS-1$

	public enum Version {
		SEVEN_1("7.1", (byte)0), //$NON-NLS-1$
		SEVEN_3("7.3", (byte)0), //$NON-NLS-1$
		SEVEN_4("7.4", (byte)0), //$NON-NLS-1$
		EIGHT_0("8.0", (byte)(longDatesTimes?0:1)), //$NON-NLS-1$
		EIGHT_2("8.2", (byte)2), //$NON-NLS-1$
		EIGHT_4("8.4.0.CR3", (byte)2), //$NON-NLS-1$
		EIGHT_6("08.06.00.Beta3", (byte)3),
		EIGHT_7("08.07.00.Beta2", (byte)3), //$NON-NLS-1$
		EIGHT_10("08.10.00.Alpha3", BatchSerializer.VERSION_GEOMETRY), //$NON-NLS-1$;
		ELEVEN_2("11.02", BatchSerializer.VERSION_GEOGRAPHY); //$NON-NLS-1$

		private String string;
		private byte clientSerializationVersion;

		private Version(String string, byte clientSerializationVersion) {
			this.string = string;
			this.clientSerializationVersion = clientSerializationVersion;
		}

		public byte getClientSerializationVersion() {
			return clientSerializationVersion;
		}

		@Override
		public String toString() {
			return string;
		}

		private static TreeMap<String, Version> versionMap = new TreeMap<String, Version>();
		static {
			for (Version v : Version.values()) {
				versionMap.put(v.toString(), v);
			}
		}

		public static Version getVersion(String version) {
			Map.Entry<String, Version> v = versionMap.floorEntry(version);
			if (v == null) {
				return SEVEN_1;
			}
			return v.getValue();
		}

		public static Version latest() {
			return versionMap.lastEntry().getValue();
		}
	}

	private static ThreadLocal<DQPWorkContext> CONTEXTS = new ThreadLocal<DQPWorkContext>() {
		protected DQPWorkContext initialValue() {
			return new DQPWorkContext();
		}
	};

	public static DQPWorkContext getWorkContext() {
		return CONTEXTS.get();
	}

	public static void setWorkContext(DQPWorkContext context) {
		LogManager.removeMdc(TEIID_SESSION);
		LogManager.removeMdc(TEIID_VDB);
		if (context == null) {
			CONTEXTS.remove();
		} else {
			if (context.session != null) {
				if (context.session.getSessionId() != null)
					LogManager.putMdc(TEIID_SESSION, getSessionTag(context.session.getSessionId()));
				if (context.session.getVdb() != null) {
					LogManager.putMdc(TEIID_VDB, getVdbTag(context.session.getVdb().getFullName()));
				}
			}
			CONTEXTS.set(context);
		}
	}

	private static String getSessionTag(String sessionId) {
		return "(session:" + sessionId + ")";
	}

	private static String getVdbTag(String fullName) {
		return "(vdb:" + fullName + ")";
	}

	private volatile SessionMetadata session = new SessionMetadata();
    private String clientAddress;
    private String clientHostname;
    private SecurityHelper securityHelper;
    private HashMap<String, DataPolicy> policies;
    private boolean useCallingThread;
    private Version clientVersion = Version.latest();
    private boolean admin;
    private boolean noViewProcFuncDefs = false;
    private RefreshSemaphore dwhSemaphore = null;
	private Thread dwhSemaphoreFlag = null;
	private String translatedSql = null;
	private List<Exception> warnings = null;

	private transient EmbeddedProfile connectionProfile = new EmbeddedProfile();

    public DQPWorkContext() {
		POLICY_UPDATE.mark(this);
		METADATA_UPDATE.mark(this);
	}

    public boolean useCallingThread() {
		return useCallingThread;
	}

    public void setUseCallingThread(boolean useCallingThread) {
		this.useCallingThread = useCallingThread;
	}

    public SessionMetadata getSession() {
		return session;
	}

    public void setSession(SessionMetadata session) {
		this.session = session;
		this.policies = null;
	}

    public void setSecurityHelper(SecurityHelper securityHelper) {
		this.securityHelper = securityHelper;
	}

    public SecurityHelper getSecurityHelper() {
		return securityHelper;
	}

    /**
     * @return
     */
    public String getUserName() {
		return session.getUserName();
    }

    public SecurityIdentity getSubject() {
		return session.getSubject();
    }

    /**
     * @return
     */
    public String getVdbName() {
        return session.getVDBName();
    }

    /**
     * @return
     */
    public int getVdbVersion() {
        return session.getVDBVersion();
    }

	public String getSessionId() {
		return this.session.getSessionId();
	}

	public String getAppName() {
		return session.getApplicationName();
	}

	public RequestID getRequestID(long exeuctionId) {
		return new RequestID(this.getSessionId(), exeuctionId);
	}

	public SessionToken getSessionToken() {
		return session.getSessionToken();
	}

	public void setClientAddress(String clientAddress) {
		this.clientAddress = clientAddress;
	}

	/**
	 * Get the client address from the socket transport - not as reported from the client
	 * @return
	 */
	public String getClientAddress() {
		return clientAddress;
	}

	public void setClientHostname(String clientHostname) {
		this.clientHostname = clientHostname;
	}

	/**
	 * Get the client hostname from the socket transport - not as reported from the client
	 * @return
	 */
	public String getClientHostname() {
		return clientHostname;
	}

	public String getSecurityDomain() {
		return this.session.getSecurityDomain();
	}

	public VDBMetaData getVDB() {
		return session.getVdb();
	}

	public <V> V runInContext(Callable<V> callable) throws Throwable {
		FutureTask<V> task = new FutureTask<V>(callable);
		runInContext(task);
		try {
			return task.get();
		} catch (ExecutionException e) {
			throw e.getCause();
		}
	}

	public void runInContext(final Runnable runnable) {
		DQPWorkContext previous = DQPWorkContext.getWorkContext();
		DQPWorkContext.setWorkContext(this);
		SecurityIdentity subject = null;
		if (this.session != null) {
			if (this.session.getSessionId() != null)
				LogManager.putMdc(TEIID_SESSION, getSessionTag(this.session.getSessionId()));
			if (this.session.getVdb() != null) {
				LogManager.putMdc(TEIID_VDB, getVdbTag(this.session.getVdb().getFullName()));
			}
			subject = getSubject();
		}
		try {
			if (subject != null)
				subject.runAs(runnable);
			else
			runnable.run();
		} finally {
			DQPWorkContext.setWorkContext(previous);
			LogManager.removeMdc(TEIID_SESSION);
			LogManager.removeMdc(TEIID_VDB);
		}
	}

	public HashMap<String, DataPolicy> getAllowedDataPolicies() {
		VDBMetaData vdb = getVDB();
		if (POLICY_UPDATE.consumeIfPending(this)) {
			this.policies = new HashMap<>();
			if (this.securityHelper != null) {
				String baseUser = FunctionMethods.trimSecurityDomain(this.session.getUserName());
				List<String> roles = DVServices.callne( DVUsersHandlerIntf.class, s -> s.getDomainOrUserRoleNames(this.session.getUserName()), Collections.emptyList() );
				SecurityIdentity si = this.securityHelper.createAdHocIdentity( getSecurityDomain(), baseUser, roles );
				this.session.setSubject(si);
			} else {
				this.session.setSubject(null);
			}
		}
		if (vdb != null && (this.policies == null || this.policies.isEmpty() || METADATA_UPDATE.consumeIfPending(this)))
			refreshRolesAndPolicies(session);
		return this.policies;
	}

	public List<String> refreshRolesAndPolicies(SessionMetadata session) {
		METADATA_UPDATE.mark(this);
		List<String> roles = DVServices.callne(DVUsersHandlerIntf.class, s -> s.getDomainOrUserRoleNames(session.getUserName()), Collections.emptyList());
		this.policies = getAllowedDataPolicies(session, roles);

		return roles;
	}

	public HashMap<String, DataPolicy> getAllowedDataPolicies(SessionMetadata session) {
		return getAllowedDataPolicies(session, null);
    }

	public HashMap<String, DataPolicy> getAllowedDataPolicies(SessionMetadata session, List<String> roles) {
		VDBMetaData vdb = getVDB();
		HashMap<String, DataPolicy> policies = new HashMap<>();
		Set<String> userRoles = (roles != null && !roles.isEmpty()) ? new HashSet<>(roles) : getUserRoles(session);
		if (this.policies != null && this.policies.isEmpty()) {
			DVServices.callne(DVUsersHandlerIntf.class, s -> s.invalidatePermissions(), null);
		}
		DVServices.callne(DVUsersHandlerIntf.class, s -> s.resurrectInvalidatedPermissions(false, false), null); // resurrect permissions that could have been changed by other dvserver instances
		// get data roles from the VDB
		TransformationMetadata metadata = vdb.getAttachment(TransformationMetadata.class);
		Collection<? extends DataPolicy> allPolicies;
		if (metadata == null) {
			allPolicies = vdb.getDataPolicies();
		} else {
			allPolicies = metadata.getPolicies().values();
		}

		for (DataPolicy policy : allPolicies) {
			if (matchesPrincipal(userRoles, policy)) {
				policies.put(policy.getName(), policy);
			}
		}
		return policies;
	}

	public void setPolicies(HashMap<String, DataPolicy> policies) {
		this.policies = policies;
	}

	private static boolean matchesPrincipal(Set<String> userRoles, DataPolicy policy) {
		if (policy.isAnyAuthenticated()) {
			return true;
		}
		return !Collections.disjoint(policy.getMappedRoleNames(), userRoles);
	}

	public Set<String> getRoleSetsInSession() {
		return getRoleSetsInSession(session);
	}

	public Set<String> getRoleSetsInSession(SessionMetadata session) {
		return getUserRoles(session);
	}

	private Set<String> getUserRoles(SessionMetadata session) {
		if (session == null || session.getSubject() == null) {
			return Collections.emptySet();
		}
		Set<String> roles = new HashSet<>();
		Roles siRoles = session.getSubject().getRoles();
		siRoles.forEach(r -> roles.add(r));

		if (getVDB() != null && METADATA_UPDATE.consumeIfPending(this)) {
			roles.addAll(refreshRolesAndPolicies(session));
		}
		
		if (this.policies != null && this.session.getUserName().equals(session.getUserName())) { // apply user's roles from cache unless the session is being executed on behalf of another user
			roles.addAll(this.policies.keySet());
		}
		
		return roles;
	}

	public Version getClientVersion() {
		return clientVersion;
	}

	public void setClientVersion(Version clientVersion) {
		this.clientVersion = clientVersion;
	}

	public void setAdmin(boolean admin) {
		this.admin = admin;
	}

	public boolean isAdmin() {
		return admin;
	}

	public void setConnectionProfile(EmbeddedProfile connectionProfile) {
		this.connectionProfile = connectionProfile;
	}

	public EmbeddedProfile getConnectionProfile() {
		return connectionProfile;
	}

	public boolean isNoViewProcFuncDefs() {
		return noViewProcFuncDefs;
	}

	public void setNoViewProcFuncDefs(boolean noViewProcFuncDefs) {
		this.noViewProcFuncDefs = noViewProcFuncDefs;
	}

	public RefreshSemaphore getDwhSemaphore() {
		return dwhSemaphore;
	}

	public void setDwhSemaphore(RefreshSemaphore dwhSemaphore) {
		this.dwhSemaphore = dwhSemaphore;
	}

	public Thread getDwhSemaphoreFlag() {
		return dwhSemaphoreFlag;
	}

	public void setDwhSemaphoreFlag(Thread dwhSemaphoreFlag) {
		this.dwhSemaphoreFlag = dwhSemaphoreFlag;
	}

	public String getTranslatedSql() {
		return translatedSql;
	}

	public synchronized void setTranslatedSql(String translatedSql) {
		this.translatedSql = translatedSql;
	}

	public List<Exception> getWarnings() {
		return warnings;
	}

	public void setWarnings(List<Exception> warnings) {
		this.warnings = warnings;
	}

	public void addWarnings(List<Exception> warnings) {
		if (this.warnings == null)
			this.warnings = new ArrayList<>(warnings);
		else
			this.warnings.addAll(warnings);
	}

	private static final ReentrantLock lock = new ReentrantLock();

	public static final class UpdateFlag {
		private final Cache<DQPWorkContext, Boolean> cache =
				Caffeine.newBuilder().weakKeys().build();

		/** Atomically marks all contexts as stale. */
		public void invalidateAll() {
			lock.lock();
			try {
				cache.invalidateAll();
			} finally {
				lock.unlock();
			}
		}

		/** Non-destructive: whether this context still has a pending update. */
		public boolean hasPending(DQPWorkContext ctx) {
			Boolean v = cache.getIfPresent(ctx);
			return v == null || !Boolean.TRUE.equals(v);
		}
	/** Destructive: returns true once after invalidate; marks this context as up-to-date. */
		public boolean consumeIfPending(DQPWorkContext ctx) {
			boolean pending = hasPending(ctx);
			if (pending) {
				cache.put(ctx, Boolean.TRUE);
			}
			return pending;
		}

		/** Explicitly mark a specific context as up-to-date. */
		public void mark(DQPWorkContext ctx) {
			cache.put(ctx, Boolean.TRUE);
		}
	}

	/** Public flags to be used across the codebase (no wrappers). */
	public static final UpdateFlag METADATA_UPDATE = new UpdateFlag();
	public static final UpdateFlag POLICY_UPDATE   = new UpdateFlag();

	public boolean isLoginAllowed() {
		if (METADATA_UPDATE.hasPending(this)) {
			return Boolean.TRUE.equals(DVServices.callne(DVUsersHandlerIntf.class, s -> s.isLoginAllowed(FunctionMethods.trimSecurityDomain(getUserName()))));
		}
		return true;
	}
}
