/*
 * 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.texera.web

import com.typesafe.scalalogging.LazyLogging
import org.apache.http.client.utils.URLEncodedUtils
import org.apache.texera.auth.JwtAuth.jwtConsumer
import org.apache.texera.auth.util.HeaderField
import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum
import org.apache.texera.dao.jooq.generated.tables.pojos.User

import java.net.URI
import java.nio.charset.Charset
import javax.websocket.HandshakeResponse
import javax.websocket.server.{HandshakeRequest, ServerEndpointConfig}
import scala.jdk.CollectionConverters.{ListHasAsScala, _}

/**
  * This configurator extracts user identity from the HTTP handshake request
  * and associates it with the ServerEndpointConfig, allowing it to be
  * accessed by WebSocket connections.
  */
class ServletAwareConfigurator extends ServerEndpointConfig.Configurator with LazyLogging {

  override def modifyHandshake(
      config: ServerEndpointConfig,
      request: HandshakeRequest,
      response: HandshakeResponse
  ): Unit = {
    try {
      val headers = request.getHeaders.asScala.view.mapValues(_.asScala.headOption).toMap
      if (
        headers.contains(HeaderField.UserComputingUnitAccess) &&
        headers.contains(HeaderField.UserId) &&
        headers.contains(HeaderField.UserName) &&
        headers.contains(HeaderField.UserEmail)
      ) {
        // KUBERNETES MODE: Construct the User object from trusted headers
        // coming from envoy and generated by access control service.

        val userId = headers.get(HeaderField.UserId).flatten.map(_.toInt).get
        val userName = headers.get(HeaderField.UserName).flatten.get
        val userEmail = headers.get(HeaderField.UserEmail).flatten.get
        val cuAccess = headers.get(HeaderField.UserComputingUnitAccess).flatten.getOrElse("")
        config.getUserProperties.put(HeaderField.UserComputingUnitAccess, cuAccess)
        logger.info(
          s"User ID: $userId, User Name: $userName, User Email: $userEmail with CU Access: $cuAccess"
        )

        config.getUserProperties.put(
          classOf[User].getName,
          new User(
            userId,
            userName,
            userEmail,
            null,
            null,
            null,
            null,
            null,
            null,
            null,
            null
          )
        )
        logger.debug(s"User created from headers: ID=$userId, Name=$userName")
      } else {
        // SINGLE NODE MODE: Construct the User object from JWT in query parameters.
        val params =
          URLEncodedUtils.parse(new URI("?" + request.getQueryString), Charset.defaultCharset())
        config.getUserProperties.put(
          HeaderField.UserComputingUnitAccess,
          PrivilegeEnum.WRITE.name()
        )
        params.asScala
          .map(pair => pair.getName -> pair.getValue)
          .toMap
          .get("access-token")
          .map(token => {
            val claims = jwtConsumer.process(token).getJwtClaims
            config.getUserProperties.put(
              classOf[User].getName,
              new User(
                claims.getClaimValue("userId").asInstanceOf[Long].toInt,
                claims.getSubject,
                String.valueOf(claims.getClaimValue("email").asInstanceOf[String]),
                null,
                null,
                null,
                null,
                null,
                null,
                null,
                null
              )
            )
          })
      }
    } catch {
      case e: Exception =>
        logger.error("Failed to retrieve the User during websocket handshake", e)
    }

  }
}
