/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.cassandra.auth;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Lists;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.cassandra.cql3.QueryOptions;
import org.apache.cassandra.cql3.QueryProcessor;
import org.apache.cassandra.cql3.UntypedResultSet;
import org.apache.cassandra.cql3.statements.SelectStatement;
import org.apache.cassandra.db.ConsistencyLevel;
import org.apache.cassandra.db.marshal.UTF8Type;
import org.apache.cassandra.exceptions.ConfigurationException;
import org.apache.cassandra.exceptions.RequestExecutionException;
import org.apache.cassandra.schema.SchemaConstants;
import org.apache.cassandra.service.ClientState;
import org.apache.cassandra.transport.Dispatcher;
import org.apache.cassandra.transport.messages.ResultMessage;
import org.apache.cassandra.utils.ByteBufferUtil;

import static org.apache.cassandra.service.QueryState.forInternalCalls;

public class CassandraNetworkAuthorizer implements INetworkAuthorizer
{
    private static final Logger logger = LoggerFactory.getLogger(CassandraNetworkAuthorizer.class);
    private SelectStatement authorizeUserStatement = null;

    public void setup()
    {
        String query = String.format("SELECT dcs FROM %s.%s WHERE role = ?",
                                     SchemaConstants.AUTH_KEYSPACE_NAME,
                                     AuthKeyspace.NETWORK_PERMISSIONS);
        authorizeUserStatement = (SelectStatement) QueryProcessor.getStatement(query, ClientState.forInternalCalls());
    }

    @VisibleForTesting
    ResultMessage.Rows select(SelectStatement statement, QueryOptions options)
    {
        return statement.execute(forInternalCalls(), options, Dispatcher.RequestTime.forImmediateExecution());
    }

    /**
     * This is exposed so we can override the consistency level for tests that are single node
     */
    @VisibleForTesting
    UntypedResultSet process(String query, ConsistencyLevel cl) throws RequestExecutionException
    {
        return QueryProcessor.process(query, cl);
    }

    private Set<String> getAuthorizedDcs(String name)
    {
        QueryOptions options = QueryOptions.forInternalCalls(CassandraAuthorizer.authReadConsistencyLevel(),
                                                             Lists.newArrayList(ByteBufferUtil.bytes(name)));

        ResultMessage.Rows rows = select(authorizeUserStatement, options);
        UntypedResultSet result = UntypedResultSet.create(rows.result);
        Set<String> dcs = null;
        if (!result.isEmpty() && result.one().has("dcs"))
        {
            dcs = result.one().getFrozenSet("dcs", UTF8Type.instance);
        }
        return dcs;
    }

    public DCPermissions authorize(RoleResource role)
    {
        if (!Roles.canLogin(role))
        {
            return DCPermissions.none();
        }
        if (Roles.hasSuperuserStatus(role))
        {
            return DCPermissions.all();
        }

        Set<String> dcs = getAuthorizedDcs(role.getName());

        if (dcs == null || dcs.isEmpty())
        {
            return DCPermissions.all();
        }
        else
        {
            return DCPermissions.subset(dcs);
        }
    }

    private static String getSetString(DCPermissions permissions)
    {
        if (permissions.restrictsAccess())
        {
            StringBuilder builder = new StringBuilder();
            builder.append('{');
            boolean first = true;
            for (String dc: permissions.allowedDCs())
            {
                if (first)
                {
                    first = false;
                }
                else
                {
                    builder.append(", ");
                }
                builder.append('\'');
                builder.append(dc);
                builder.append('\'');
            }
            builder.append('}');
            return builder.toString();
        }
        else
        {
            return "{}";
        }
    }

    public void setRoleDatacenters(RoleResource role, DCPermissions permissions)
    {
        String query = String.format("UPDATE %s.%s SET dcs = %s WHERE role = '%s'",
                                     SchemaConstants.AUTH_KEYSPACE_NAME,
                                     AuthKeyspace.NETWORK_PERMISSIONS,
                                     getSetString(permissions),
                                     role.getName());

        process(query, CassandraAuthorizer.authWriteConsistencyLevel());
    }

    public void drop(RoleResource role)
    {
        String query = String.format("DELETE FROM %s.%s WHERE role = '%s'",
                                     SchemaConstants.AUTH_KEYSPACE_NAME,
                                     AuthKeyspace.NETWORK_PERMISSIONS,
                                     role.getName());

        process(query, CassandraAuthorizer.authWriteConsistencyLevel());
    }

    public void validateConfiguration() throws ConfigurationException
    {
        // noop
    }

    @Override
    public Supplier<Map<RoleResource, DCPermissions>> bulkLoader()
    {
        return () -> {
            logger.info("Pre-warming datacenter permissions cache from network_permissions table");
            Map<RoleResource, DCPermissions> entries = new HashMap<>();
            UntypedResultSet rows = process(String.format("SELECT role, dcs FROM %s.%s",
                                                          SchemaConstants.AUTH_KEYSPACE_NAME,
                                                          AuthKeyspace.NETWORK_PERMISSIONS),
                                            CassandraAuthorizer.authReadConsistencyLevel());

            for (UntypedResultSet.Row row : rows)
            {
                RoleResource role = RoleResource.role(row.getString("role"));
                DCPermissions.Builder builder = new DCPermissions.Builder();
                Set<String> dcs = row.getFrozenSet("dcs", UTF8Type.instance);
                for (String dc : dcs)
                    builder.add(dc);
                entries.put(role, builder.build());
            }

            return entries;
        };
    }
}
