// Licensed to the Software Freedom Conservancy (SFC) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The SFC 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.openqa.selenium.grid.sessionmap.jdbc;

import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES;
import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES_EVENT;
import static org.openqa.selenium.remote.RemoteTags.SESSION_ID;
import static org.openqa.selenium.remote.RemoteTags.SESSION_ID_EVENT;
import static org.openqa.selenium.remote.tracing.Tags.EXCEPTION;

import java.io.Closeable;
import java.net.URI;
import java.net.URISyntaxException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.time.Instant;
import java.util.HashMap;
import java.util.Map;
import java.util.logging.Logger;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.NoSuchSessionException;
import org.openqa.selenium.events.EventBus;
import org.openqa.selenium.grid.config.Config;
import org.openqa.selenium.grid.config.ConfigException;
import org.openqa.selenium.grid.data.NodeRemovedEvent;
import org.openqa.selenium.grid.data.NodeRestartedEvent;
import org.openqa.selenium.grid.data.Session;
import org.openqa.selenium.grid.data.SessionClosedEvent;
import org.openqa.selenium.grid.log.LoggingOptions;
import org.openqa.selenium.grid.server.EventBusOptions;
import org.openqa.selenium.grid.sessionmap.SessionMap;
import org.openqa.selenium.internal.Require;
import org.openqa.selenium.json.Json;
import org.openqa.selenium.remote.SessionId;
import org.openqa.selenium.remote.tracing.AttributeKey;
import org.openqa.selenium.remote.tracing.EventAttribute;
import org.openqa.selenium.remote.tracing.EventAttributeValue;
import org.openqa.selenium.remote.tracing.Span;
import org.openqa.selenium.remote.tracing.Status;
import org.openqa.selenium.remote.tracing.Tracer;

public class JdbcBackedSessionMap extends SessionMap implements Closeable {

  private static final Json JSON = new Json();
  private static final Logger LOG = Logger.getLogger(JdbcBackedSessionMap.class.getName());
  private static final String TABLE_NAME = "sessions_map";
  private static final String SESSION_ID_COL = "session_ids";
  private static final String SESSION_CAPS_COL = "session_caps";
  private static final String SESSION_STEREOTYPE_COL = "session_stereotype";
  private static final String SESSION_URI_COL = "session_uri";
  private static final String SESSION_START_COL = "session_start";
  private static final String DATABASE_STATEMENT = AttributeKey.DATABASE_STATEMENT.getKey();
  private static final String DATABASE_OPERATION = AttributeKey.DATABASE_OPERATION.getKey();
  private static final String DATABASE_USER = AttributeKey.DATABASE_USER.getKey();
  private static final String DATABASE_CONNECTION_STRING =
      AttributeKey.DATABASE_CONNECTION_STRING.getKey();
  private static String jdbcUser;
  private static String jdbcUrl;
  private final EventBus bus;
  private final Connection connection;

  public JdbcBackedSessionMap(Tracer tracer, Connection jdbcConnection, EventBus bus) {
    super(tracer);

    Require.nonNull("JDBC Connection Object", jdbcConnection);
    this.bus = Require.nonNull("Event bus", bus);

    this.connection = jdbcConnection;
    this.bus.addListener(SessionClosedEvent.listener(this::remove));

    this.bus.addListener(
        NodeRemovedEvent.listener(
            nodeStatus ->
                nodeStatus.getSlots().stream()
                    .filter(slot -> slot.getSession() != null)
                    .map(slot -> slot.getSession().getId())
                    .forEach(this::remove)));

    bus.addListener(
        NodeRestartedEvent.listener(nodeStatus -> this.removeByUri(nodeStatus.getExternalUri())));
  }

  public static SessionMap create(Config config) {
    Tracer tracer = new LoggingOptions(config).getTracer();
    EventBus bus = new EventBusOptions(config).getEventBus();

    JdbcSessionMapOptions sessionMapOptions = new JdbcSessionMapOptions(config);

    Connection connection;

    try {
      jdbcUser = sessionMapOptions.getJdbcUser();
      jdbcUrl = sessionMapOptions.getJdbcUrl();
      connection = sessionMapOptions.getJdbcConnection();
    } catch (SQLException e) {
      throw new ConfigException(e);
    }

    return new JdbcBackedSessionMap(tracer, connection, bus);
  }

  @Override
  public boolean isReady() {
    try {
      return !connection.isClosed();
    } catch (SQLException throwable) {
      return false;
    }
  }

  @Override
  public boolean add(Session session) {
    Require.nonNull("Session to add", session);

    try (Span span =
        tracer
            .getCurrentContext()
            .createSpan(
                "INSERT into  sessions_map (session_ids, session_uri, session_caps, session_start)"
                    + " values (?, ?, ?, ?) ")) {
      Map<String, EventAttributeValue> attributeMap = new HashMap<>();
      SESSION_ID.accept(span, session.getId());
      SESSION_ID_EVENT.accept(attributeMap, session.getId());
      CAPABILITIES.accept(span, session.getCapabilities());
      CAPABILITIES_EVENT.accept(attributeMap, session.getCapabilities());
      setCommonSpanAttributes(span);
      setCommonEventAttributes(attributeMap);
      attributeMap.put(
          AttributeKey.SESSION_URI.getKey(), EventAttribute.setValue(session.getUri().toString()));

      try (PreparedStatement statement =
          connection.prepareStatement(
              String.format(
                  "insert into %1$s (%2$s, %3$s, %4$s, %5$s, %6$s) values (?, ?, ?, ?, ?)",
                  TABLE_NAME,
                  SESSION_ID_COL,
                  SESSION_URI_COL,
                  SESSION_STEREOTYPE_COL,
                  SESSION_CAPS_COL,
                  SESSION_START_COL))) {

        statement.setString(1, session.getId().toString());
        statement.setString(2, session.getUri().toString());
        statement.setString(3, JSON.toJson(session.getStereotype()));
        statement.setString(4, JSON.toJson(session.getCapabilities()));
        statement.setString(5, JSON.toJson(session.getStartTime()));

        String statementStr = statement.toString();
        span.setAttribute(DATABASE_STATEMENT, statementStr);
        span.setAttribute(DATABASE_OPERATION, "insert");
        attributeMap.put(DATABASE_STATEMENT, EventAttribute.setValue(statementStr));
        attributeMap.put(DATABASE_OPERATION, EventAttribute.setValue("insert"));

        int rowCount = statement.executeUpdate();
        attributeMap.put("rows.added", EventAttribute.setValue(rowCount));
        span.addEvent("Inserted into the database", attributeMap);
        return rowCount >= 1;
      } catch (SQLException e) {
        span.setAttribute("error", true);
        span.setStatus(Status.CANCELLED);
        EXCEPTION.accept(attributeMap, e);
        attributeMap.put(
            AttributeKey.EXCEPTION_MESSAGE.getKey(),
            EventAttribute.setValue(
                "Unable to add session information to the database: " + e.getMessage()));
        span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);

        throw new JdbcException(e);
      }
    }
  }

  @Override
  public Session get(SessionId id) throws NoSuchSessionException {
    Require.nonNull("Session ID", id);

    URI uri = null;
    Capabilities stereotype = null;
    Capabilities caps = null;
    Instant start = null;
    String rawUri = null;
    Map<String, EventAttributeValue> attributeMap = new HashMap<>();

    try (Span span =
        tracer
            .getCurrentContext()
            .createSpan("SELECT * from  sessions_map where session_ids = ?")) {
      SESSION_ID.accept(span, id);
      SESSION_ID_EVENT.accept(attributeMap, id);
      setCommonSpanAttributes(span);
      setCommonEventAttributes(attributeMap);

      try (PreparedStatement statement =
          connection.prepareStatement(
              String.format("select * from %1$s where %2$s = ?", TABLE_NAME, SESSION_ID_COL))) {

        statement.setMaxRows(1);
        statement.setString(1, id.toString());

        String statementStr = statement.toString();
        span.setAttribute(DATABASE_STATEMENT, statementStr);
        span.setAttribute(DATABASE_OPERATION, "select");
        attributeMap.put(DATABASE_OPERATION, EventAttribute.setValue("select"));
        attributeMap.put(DATABASE_STATEMENT, EventAttribute.setValue(statementStr));

        try (ResultSet sessions = statement.executeQuery()) {
          if (!sessions.next()) {
            NoSuchSessionException exception =
                new NoSuchSessionException("Unable to find session.");
            span.setAttribute("error", true);
            span.setStatus(Status.NOT_FOUND);
            EXCEPTION.accept(attributeMap, exception);
            attributeMap.put(
                AttributeKey.EXCEPTION_MESSAGE.getKey(),
                EventAttribute.setValue(
                    "Session id does not exist in the database :" + exception.getMessage()));
            span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);

            throw exception;
          }

          rawUri = sessions.getString(SESSION_URI_COL);

          String rawStereotype = sessions.getString(SESSION_STEREOTYPE_COL);

          stereotype =
              rawStereotype == null
                  ? new ImmutableCapabilities()
                  : JSON.toType(rawStereotype, Capabilities.class);

          String rawCapabilities = sessions.getString(SESSION_CAPS_COL);

          caps =
              rawCapabilities == null
                  ? new ImmutableCapabilities()
                  : JSON.toType(rawCapabilities, Capabilities.class);

          String rawStart = sessions.getString(SESSION_START_COL);
          start = JSON.toType(rawStart, Instant.class);
        }
        CAPABILITIES_EVENT.accept(attributeMap, caps);

        try {
          attributeMap.put(AttributeKey.SESSION_URI.getKey(), EventAttribute.setValue(rawUri));
          uri = new URI(rawUri);
        } catch (URISyntaxException e) {
          span.setAttribute("error", true);
          span.setStatus(Status.INTERNAL);
          EXCEPTION.accept(attributeMap, e);
          attributeMap.put(AttributeKey.SESSION_URI.getKey(), EventAttribute.setValue(rawUri));
          attributeMap.put(
              AttributeKey.EXCEPTION_MESSAGE.getKey(),
              EventAttribute.setValue("Unable to convert session id to uri: " + e.getMessage()));
          span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);

          throw new NoSuchSessionException(
              String.format("Unable to convert session id (%s) to uri: %s", id, rawUri), e);
        }

        span.addEvent("Retrieved session from the database", attributeMap);
        return new Session(id, uri, stereotype, caps, start);
      } catch (SQLException e) {
        span.setAttribute("error", true);
        span.setStatus(Status.CANCELLED);
        EXCEPTION.accept(attributeMap, e);
        attributeMap.put(
            AttributeKey.EXCEPTION_MESSAGE.getKey(),
            EventAttribute.setValue(
                "Unable to get session information from the database: " + e.getMessage()));
        span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);
        throw new JdbcException(e);
      }
    }
  }

  @Override
  public void remove(SessionId id) {
    Require.nonNull("Session ID", id);
    try (Span span =
        tracer.getCurrentContext().createSpan("DELETE from  sessions_map where session_ids = ?")) {
      Map<String, EventAttributeValue> attributeMap = new HashMap<>();
      SESSION_ID.accept(span, id);
      SESSION_ID_EVENT.accept(attributeMap, id);
      setCommonSpanAttributes(span);
      setCommonEventAttributes(attributeMap);

      try (PreparedStatement statement =
          connection.prepareStatement(
              String.format("delete from %1$s where %2$s = ?", TABLE_NAME, SESSION_ID_COL))) {

        statement.setString(1, id.toString());
        String statementStr = statement.toString();
        span.setAttribute(DATABASE_STATEMENT, statementStr);
        span.setAttribute(DATABASE_OPERATION, "delete");
        attributeMap.put(DATABASE_STATEMENT, EventAttribute.setValue(statementStr));
        attributeMap.put(DATABASE_OPERATION, EventAttribute.setValue("delete"));

        int rowCount = statement.executeUpdate();
        attributeMap.put("rows.deleted", EventAttribute.setValue(rowCount));
        span.addEvent("Deleted session from the database", attributeMap);

      } catch (SQLException e) {
        span.setAttribute("error", true);
        span.setStatus(Status.CANCELLED);
        EXCEPTION.accept(attributeMap, e);
        attributeMap.put(
            AttributeKey.EXCEPTION_MESSAGE.getKey(),
            EventAttribute.setValue(
                "Unable to delete session information from the database: " + e.getMessage()));
        span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);
        throw new JdbcException(e.getMessage());
      }
    }
  }

  public void removeByUri(URI sessionUri) {
    Require.nonNull("Session URI", sessionUri);
    try (Span span =
        tracer.getCurrentContext().createSpan("DELETE from  sessions_map where session_uri = ?")) {
      Map<String, EventAttributeValue> attributeMap = new HashMap<>();

      try (PreparedStatement statement =
          connection.prepareStatement(
              String.format("delete from %1$s where %2$s = ?", TABLE_NAME, SESSION_URI_COL))) {

        statement.setString(1, sessionUri.toString());
        String statementStr = statement.toString();
        span.setAttribute(DATABASE_STATEMENT, statementStr);
        span.setAttribute(DATABASE_OPERATION, "delete");
        attributeMap.put(DATABASE_STATEMENT, EventAttribute.setValue(statementStr));
        attributeMap.put(DATABASE_OPERATION, EventAttribute.setValue("delete"));

        int rowCount = statement.executeUpdate();
        attributeMap.put("rows.deleted", EventAttribute.setValue(rowCount));
        span.addEvent("Deleted session from the database", attributeMap);

      } catch (SQLException e) {
        span.setAttribute("error", true);
        span.setStatus(Status.CANCELLED);
        EXCEPTION.accept(attributeMap, e);
        attributeMap.put(
            AttributeKey.EXCEPTION_MESSAGE.getKey(),
            EventAttribute.setValue(
                "Unable to delete session information from the database: " + e.getMessage()));
        span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);
        throw new JdbcException(e.getMessage());
      }
    }
  }

  @Override
  public void close() {
    try {
      connection.close();
    } catch (SQLException e) {
      LOG.warning("SQL exception while closing JDBC Connection:" + e.getMessage());
    }
  }

  private void setCommonSpanAttributes(Span span) {
    span.setAttribute("span.kind", Span.Kind.CLIENT.toString());
    if (jdbcUser != null) {
      span.setAttribute(DATABASE_USER, jdbcUser);
    }
    if (jdbcUrl != null) {
      span.setAttribute(DATABASE_CONNECTION_STRING, jdbcUrl);
    }
  }

  private void setCommonEventAttributes(Map<String, EventAttributeValue> attributeMap) {
    if (jdbcUser != null) {
      attributeMap.put(DATABASE_USER, EventAttribute.setValue(jdbcUser));
    }
    if (jdbcUrl != null) {
      attributeMap.put(DATABASE_CONNECTION_STRING, EventAttribute.setValue(jdbcUrl));
    }
  }
}
