/*
 * Copyright 2024-2026 the original author or authors.
 *
 * Licensed 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
 *
 *      https://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 com.alibaba.cloud.ai.graph.checkpoint.savers.h2;

import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.checkpoint.Checkpoint;
import com.alibaba.cloud.ai.graph.checkpoint.savers.jdbc.AbstractJdbcCheckpointSaver;
import com.alibaba.cloud.ai.graph.serializer.StateSerializer;

import java.io.IOException;
import java.io.PrintWriter;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.sql.Statement;
import java.util.Base64;
import java.util.Collection;
import java.util.Collections;
import java.util.LinkedList;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import java.util.logging.Logger;

import javax.sql.DataSource;

import org.slf4j.LoggerFactory;

import static java.lang.String.format;
import static java.util.Objects.requireNonNull;

/**
 * H2 implementation of a JDBC checkpoint saver.
 * <p>
 * H2 is commonly used as an embedded or in-memory database, but this saver still
 * uses the same persistent JDBC saver contract as MySQL, PostgreSQL and Oracle:
 * the database is the checkpoint history source of truth, while the JVM can use
 * an opt-in bounded latest-checkpoint cache.
 * </p>
 * <p>
 * Two tables are used to store the workflow state:
 *
 * <pre>
 *     CREATE TABLE GRAPH_THREAD (
 *          thread_id VARCHAR(36) PRIMARY KEY,
 *          thread_name VARCHAR(255) NOT NULL,
 *          is_released BOOLEAN DEFAULT FALSE NOT NULL,
 *          active_thread_name VARCHAR(255) GENERATED ALWAYS AS (
 *              CASE WHEN is_released = FALSE THEN thread_name ELSE NULL END
 *          )
 *     )
 *     CREATE UNIQUE INDEX IDX_GRAPH_THREAD_ACTIVE_NAME
 *          ON GRAPH_THREAD(active_thread_name)
 *
 *     CREATE TABLE GRAPH_CHECKPOINT (
 *          checkpoint_seq BIGINT GENERATED BY DEFAULT AS IDENTITY UNIQUE,
 *          checkpoint_id VARCHAR(36) PRIMARY KEY,
 *          thread_id VARCHAR(36) NOT NULL,
 *          node_id VARCHAR(255),
 *          next_node_id VARCHAR(255),
 *          state_data CLOB NOT NULL,
 *          state_content_type VARCHAR(100) NOT NULL,
 *          saved_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
 *
 *          CONSTRAINT GRAPH_FK_THREAD
 *              FOREIGN KEY(thread_id)
 *              REFERENCES GRAPH_THREAD(thread_id)
 *              ON DELETE CASCADE
 *     )
 * </pre>
 * </p>
 * <p>
 * A builder can be used to create an instance of H2Saver. The builder allows
 * configuring a DataSource or JDBC URL, CreateOption, StateSerializer, and the
 * maximum number of latest checkpoints retained in memory.
 * </p>
 * <p>
 * Ex:
 *
 * <pre>
 * var saver = H2Saver.builder()
 *         .stateSerializer(STATE_SERIALIZER)
 *         .createOption(CreateOption.CREATE_OR_REPLACE)
 *         .dataSource(DATA_SOURCE)
 *         .build();
 * </pre>
 * </p>
 */
public class H2Saver extends AbstractJdbcCheckpointSaver {

	private static final org.slf4j.Logger log = LoggerFactory.getLogger(H2Saver.class);

	private static final String DEFAULT_JDBC_URL = "jdbc:h2:mem:langgraph4j;DB_CLOSE_DELAY=-1;DB_CLOSE_ON_EXIT=FALSE";

	private static final String DROP_TABLES = """
			DROP TABLE IF EXISTS GRAPH_CHECKPOINT;
			DROP TABLE IF EXISTS GRAPH_THREAD;
			""";

	private static final String CREATE_TABLES = """
			CREATE TABLE IF NOT EXISTS GRAPH_THREAD (
			   thread_id VARCHAR(36) PRIMARY KEY,
			   thread_name VARCHAR(255) NOT NULL,
			   is_released BOOLEAN DEFAULT FALSE NOT NULL,
			   active_thread_name VARCHAR(255) GENERATED ALWAYS AS (
			       CASE WHEN is_released = FALSE THEN thread_name ELSE NULL END
			   )
			);

			CREATE TABLE IF NOT EXISTS GRAPH_CHECKPOINT (
			   checkpoint_seq BIGINT GENERATED BY DEFAULT AS IDENTITY UNIQUE,
			   checkpoint_id VARCHAR(36) PRIMARY KEY,
			   thread_id VARCHAR(36) NOT NULL,
			   node_id VARCHAR(255),
			   next_node_id VARCHAR(255),
			   state_data CLOB NOT NULL,
			   state_content_type VARCHAR(100) NOT NULL,
			   saved_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,

			   CONSTRAINT GRAPH_FK_THREAD
			       FOREIGN KEY(thread_id)
			       REFERENCES GRAPH_THREAD(thread_id)
			       ON DELETE CASCADE
			);
			""";

	private static final String CREATE_INDEXES = """
			CREATE UNIQUE INDEX IF NOT EXISTS IDX_GRAPH_THREAD_ACTIVE_NAME
			  ON GRAPH_THREAD(active_thread_name);
			CREATE INDEX IF NOT EXISTS IDX_GRAPH_THREAD_NAME_RELEASED
			  ON GRAPH_THREAD(thread_name, is_released);
			CREATE INDEX IF NOT EXISTS IDX_GRAPH_CHECKPOINT_THREAD_SEQUENCE
			  ON GRAPH_CHECKPOINT(thread_id, checkpoint_seq);
			""";

	private static final String SELECT_ACTIVE_THREAD = """
			SELECT thread_id
			FROM GRAPH_THREAD
			WHERE thread_name = ? AND is_released = FALSE
			LIMIT 1
			""";

	private static final String INSERT_THREAD = """
			INSERT INTO GRAPH_THREAD (thread_id, thread_name, is_released)
			VALUES (?, ?, FALSE)
			""";

	private static final String INSERT_CHECKPOINT = """
			INSERT INTO GRAPH_CHECKPOINT(checkpoint_id, thread_id, node_id, next_node_id, state_data, state_content_type)
			VALUES (?, ?, ?, ?, ?, ?)
			""";

	private static final String UPDATE_CHECKPOINT = """
			UPDATE GRAPH_CHECKPOINT c
			SET
			  checkpoint_id = ?,
			  node_id = ?,
			  next_node_id = ?,
			  state_data = ?,
			  state_content_type = ?
			WHERE c.checkpoint_id = ?
			  AND EXISTS (
			    SELECT 1
			    FROM GRAPH_THREAD t
			    WHERE t.thread_id = c.thread_id
			      AND t.thread_name = ?
			      AND t.is_released = FALSE
			  )
			""";

	private static final String SELECT_CHECKPOINTS = """
			SELECT
			  c.checkpoint_id,
			  c.node_id,
			  c.next_node_id,
			  c.state_data,
			  c.state_content_type
			FROM GRAPH_CHECKPOINT c
			  INNER JOIN GRAPH_THREAD t ON c.thread_id = t.thread_id
			WHERE t.thread_name = ? AND t.is_released = FALSE
			ORDER BY c.checkpoint_seq DESC
			""";

	private static final String SELECT_LATEST_CHECKPOINT = """
			SELECT
			  c.checkpoint_id,
			  c.node_id,
			  c.next_node_id,
			  c.state_data,
			  c.state_content_type
			FROM GRAPH_CHECKPOINT c
			  INNER JOIN GRAPH_THREAD t ON c.thread_id = t.thread_id
			WHERE t.thread_name = ? AND t.is_released = FALSE
			ORDER BY c.checkpoint_seq DESC
			LIMIT 1
			""";

	private static final String SELECT_CHECKPOINT_BY_ID = """
			SELECT
			  c.checkpoint_id,
			  c.node_id,
			  c.next_node_id,
			  c.state_data,
			  c.state_content_type
			FROM GRAPH_CHECKPOINT c
			  INNER JOIN GRAPH_THREAD t ON c.thread_id = t.thread_id
			WHERE t.thread_name = ? AND t.is_released = FALSE
			  AND c.checkpoint_id = ?
			""";

	private static final String DELETE_CHECKPOINTS = """
			DELETE FROM GRAPH_CHECKPOINT
			WHERE checkpoint_id IN (%s)
			  AND thread_id IN (
			    SELECT thread_id
			    FROM GRAPH_THREAD
			    WHERE thread_name = ? AND is_released = FALSE
			  )
			""";

	private static final String RELEASE_THREAD = """
			UPDATE GRAPH_THREAD SET is_released = TRUE WHERE thread_name = ? AND is_released = FALSE
			""";

	private final DataSource dataSource;

	private final StateSerializer stateSerializer;

	private final CreateOption createOption;

	private H2Saver(Builder builder) throws SQLException {
		super(builder.maxCachedThreads);
		this.dataSource = builder.dataSource;
		this.stateSerializer = builder.stateSerializer;
		this.createOption = builder.createOption;
		initTables(createOption);
	}

	public static Builder builder() {
		return new Builder();
	}

	private void initTables(CreateOption createOption) throws SQLException {
		if (createOption == CreateOption.CREATE_NONE) {
			return;
		}
		try (Connection connection = getConnection(); Statement statement = connection.createStatement()) {
			if (createOption == CreateOption.CREATE_OR_REPLACE) {
				executeStatements(statement, DROP_TABLES);
			}
			executeStatements(statement, CREATE_TABLES);
			executeStatements(statement, CREATE_INDEXES);
		}
	}

	private void executeStatements(Statement statement, String sqlStatements) throws SQLException {
		for (String sql : sqlStatements.split(";")) {
			if (!sql.isBlank()) {
				statement.execute(sql);
			}
		}
	}

	private String encodeState(Map<String, Object> data) throws IOException {
		return Base64.getEncoder().encodeToString(stateSerializer.dataToBytes(data));
	}

	private Map<String, Object> decodeState(String binaryPayload, String contentType)
			throws IOException, ClassNotFoundException {
		if (!Objects.equals(contentType, stateSerializer.contentType())) {
			throw new IllegalStateException(
					format("Content Type used for store state '%s' is different from one '%s' used for deserialize it",
							contentType,
							stateSerializer.contentType()));
		}
		return stateSerializer.dataFromBytes(Base64.getDecoder().decode(binaryPayload));
	}

	private Checkpoint readCheckpoint(ResultSet resultSet)
			throws SQLException, IOException, ClassNotFoundException {
		return Checkpoint.builder()
				.id(resultSet.getString(1))
				.nodeId(resultSet.getString(2))
				.nextNodeId(resultSet.getString(3))
				.state(decodeState(resultSet.getString(4), resultSet.getString(5)))
				.build();
	}

	@Override
	protected LinkedList<Checkpoint> selectCheckpoints(String threadId) throws Exception {
		LinkedList<Checkpoint> checkpoints = new LinkedList<>();
		try (Connection conn = getConnection();
				PreparedStatement ps = conn.prepareStatement(SELECT_CHECKPOINTS)) {
			ps.setString(1, threadId);
			try (ResultSet rs = ps.executeQuery()) {
				while (rs.next()) {
					checkpoints.add(readCheckpoint(rs));
				}
			}
		}
		catch (SQLException | IOException | ClassNotFoundException ex) {
			throw new Exception("Unable to load checkpoints", ex);
		}
		return checkpoints;
	}

	@Override
	protected Optional<Checkpoint> selectLatestCheckpoint(String threadId) throws Exception {
		try (Connection conn = getConnection();
				PreparedStatement ps = conn.prepareStatement(SELECT_LATEST_CHECKPOINT)) {
			ps.setString(1, threadId);
			try (ResultSet rs = ps.executeQuery()) {
				if (rs.next()) {
					return Optional.of(readCheckpoint(rs));
				}
				return Optional.empty();
			}
		}
		catch (SQLException | IOException | ClassNotFoundException ex) {
			throw new Exception("Unable to load latest checkpoint", ex);
		}
	}

	@Override
	protected Optional<Checkpoint> selectCheckpointById(String threadId, String checkpointId) throws Exception {
		try (Connection conn = getConnection();
				PreparedStatement ps = conn.prepareStatement(SELECT_CHECKPOINT_BY_ID)) {
			ps.setString(1, threadId);
			ps.setString(2, checkpointId);
			try (ResultSet rs = ps.executeQuery()) {
				if (rs.next()) {
					return Optional.of(readCheckpoint(rs));
				}
				return Optional.empty();
			}
		}
		catch (SQLException | IOException | ClassNotFoundException ex) {
			throw new Exception("Unable to load checkpoint", ex);
		}
	}

	@Override
	protected void insertCheckpoint(String threadId, Checkpoint checkpoint) throws Exception {
		Connection conn = null;
		try (Connection ignored = conn = getConnection()) {
			conn.setAutoCommit(false);
			String persistedThreadId = activeThreadId(conn, threadId);
			try (PreparedStatement ps = conn.prepareStatement(INSERT_CHECKPOINT)) {
				ps.setString(1, checkpoint.getId());
				ps.setString(2, persistedThreadId);
				ps.setString(3, checkpoint.getNodeId());
				ps.setString(4, checkpoint.getNextNodeId());
				ps.setString(5, encodeState(checkpoint.getState()));
				ps.setString(6, stateSerializer.contentType());
				ps.executeUpdate();
			}
			conn.commit();
		}
		catch (SQLException | IOException ex) {
			rollback(conn, threadId);
			throw new Exception("Unable to insert checkpoint", ex);
		}
	}

	private String activeThreadId(Connection conn, String threadName) throws SQLException {
		Optional<String> activeThreadId = selectActiveThreadId(conn, threadName);
		if (activeThreadId.isPresent()) {
			return activeThreadId.get();
		}

		String persistedThreadId = UUID.randomUUID().toString();
		try (PreparedStatement ps = conn.prepareStatement(INSERT_THREAD)) {
			ps.setString(1, persistedThreadId);
			ps.setString(2, threadName);
			ps.executeUpdate();
			return persistedThreadId;
		}
		catch (SQLException ex) {
			if (isUniqueConstraintViolation(ex)) {
				return selectActiveThreadId(conn, threadName).orElseThrow(() -> ex);
			}
			throw ex;
		}
	}

	private Optional<String> selectActiveThreadId(Connection conn, String threadName) throws SQLException {
		try (PreparedStatement ps = conn.prepareStatement(SELECT_ACTIVE_THREAD)) {
			ps.setString(1, threadName);
			try (ResultSet rs = ps.executeQuery()) {
				if (rs.next()) {
					return Optional.of(rs.getString(1));
				}
			}
		}
		return Optional.empty();
	}

	private boolean isUniqueConstraintViolation(SQLException ex) {
		return "23505".equals(ex.getSQLState());
	}

	@Override
	protected void updateCheckpoint(String threadId, String checkpointId, Checkpoint checkpoint) throws Exception {
		Connection conn = null;
		try (Connection ignored = conn = getConnection()) {
			conn.setAutoCommit(false);
			try (PreparedStatement ps = conn.prepareStatement(UPDATE_CHECKPOINT)) {
				ps.setString(1, checkpoint.getId());
				ps.setString(2, checkpoint.getNodeId());
				ps.setString(3, checkpoint.getNextNodeId());
				ps.setString(4, encodeState(checkpoint.getState()));
				ps.setString(5, stateSerializer.contentType());
				ps.setString(6, checkpointId);
				ps.setString(7, threadId);
				int rowsAffected = ps.executeUpdate();
				if (rowsAffected == 0) {
					conn.rollback();
					throw new NoSuchElementException(format("Checkpoint with id %s not found!", checkpointId));
				}
			}
			conn.commit();
		}
		catch (SQLException | IOException ex) {
			rollback(conn, threadId);
			throw new Exception("Unable to update checkpoint", ex);
		}
	}

	@Override
	protected void deleteCheckpoints(String threadId, Collection<String> checkpointIds) throws Exception {
		if (checkpointIds.isEmpty()) {
			return;
		}
		String placeholders = String.join(", ", Collections.nCopies(checkpointIds.size(), "?"));
		try (Connection conn = getConnection();
				PreparedStatement ps = conn.prepareStatement(DELETE_CHECKPOINTS.formatted(placeholders))) {
			int index = 1;
			for (String checkpointId : checkpointIds) {
				ps.setString(index++, checkpointId);
			}
			ps.setString(index, threadId);
			ps.executeUpdate();
		}
		catch (SQLException ex) {
			throw new Exception("Unable to delete retained checkpoints", ex);
		}
	}

	@Override
	protected void releaseThread(String threadId) throws Exception {
		Connection conn = null;
		try (Connection ignored = conn = getConnection()) {
			conn.setAutoCommit(false);
			try (PreparedStatement ps = conn.prepareStatement(RELEASE_THREAD)) {
				ps.setString(1, threadId);
				int rowsAffected = ps.executeUpdate();
				if (rowsAffected == 0) {
					conn.rollback();
					throw new IllegalStateException(format("Thread '%s' not found or already released", threadId));
				}
			}
			conn.commit();
		}
		catch (SQLException ex) {
			rollback(conn, threadId);
			throw new Exception("Unable to release checkpoint", ex);
		}
	}

	protected Connection getConnection() throws SQLException {
		return dataSource.getConnection();
	}

	private void rollback(Connection conn, String threadId) {
		if (conn == null) {
			return;
		}
		try {
			conn.rollback();
		}
		catch (SQLException ex) {
			log.error("Failed to rollback transaction for thread {}", threadId, ex);
		}
	}

	public static class Builder {

		private StateSerializer stateSerializer;

		private DataSource dataSource;

		private String jdbcUrl = DEFAULT_JDBC_URL;

		private String user = "sa";

		private String password = "";

		private CreateOption createOption = CreateOption.CREATE_IF_NOT_EXISTS;

		private int maxCachedThreads = 1024;

		public Builder stateSerializer(StateSerializer stateSerializer) {
			this.stateSerializer = stateSerializer;
			return this;
		}

		public Builder dataSource(DataSource dataSource) {
			this.dataSource = dataSource;
			return this;
		}

		public Builder jdbcUrl(String jdbcUrl) {
			this.jdbcUrl = jdbcUrl;
			return this;
		}

		public Builder database(String database) {
			this.jdbcUrl = "jdbc:h2:mem:%s;DB_CLOSE_DELAY=-1;DB_CLOSE_ON_EXIT=FALSE".formatted(
					requireNotBlank(database, "database"));
			return this;
		}

		public Builder user(String user) {
			this.user = user;
			return this;
		}

		public Builder password(String password) {
			this.password = password;
			return this;
		}

		public Builder createOption(CreateOption createOption) {
			this.createOption = requireNonNull(createOption, "createOption cannot be null");
			return this;
		}

		public Builder maxCachedThreads(int maxCachedThreads) {
			if (maxCachedThreads < 0) {
				throw new IllegalArgumentException("maxCachedThreads must be greater than or equal to 0");
			}
			this.maxCachedThreads = maxCachedThreads;
			return this;
		}

		private String requireNotBlank(String value, String name) {
			if (requireNonNull(value, format("'%s' cannot be null", name)).isBlank()) {
				throw new IllegalArgumentException(format("'%s' cannot be blank", name));
			}
			return value;
		}

		public H2Saver build() {
			if (stateSerializer == null) {
				this.stateSerializer = StateGraph.DEFAULT_JACKSON_SERIALIZER;
			}
			if (dataSource == null) {
				dataSource = new DriverManagerDataSource(requireNotBlank(jdbcUrl, "jdbcUrl"), user, password);
			}
			try {
				return new H2Saver(this);
			}
			catch (SQLException ex) {
				throw new RuntimeException(ex);
			}
		}

	}

	private static final class DriverManagerDataSource implements DataSource {

		private final String jdbcUrl;

		private final String user;

		private final String password;

		private PrintWriter logWriter;

		private int loginTimeout;

		private DriverManagerDataSource(String jdbcUrl, String user, String password) {
			this.jdbcUrl = jdbcUrl;
			this.user = user;
			this.password = password;
		}

		@Override
		public Connection getConnection() throws SQLException {
			return DriverManager.getConnection(jdbcUrl, user, password);
		}

		@Override
		public Connection getConnection(String username, String password) throws SQLException {
			return DriverManager.getConnection(jdbcUrl, username, password);
		}

		@Override
		public PrintWriter getLogWriter() {
			return logWriter;
		}

		@Override
		public void setLogWriter(PrintWriter out) {
			this.logWriter = out;
		}

		@Override
		public void setLoginTimeout(int seconds) {
			this.loginTimeout = seconds;
		}

		@Override
		public int getLoginTimeout() {
			return loginTimeout;
		}

		@Override
		public Logger getParentLogger() throws SQLFeatureNotSupportedException {
			throw new SQLFeatureNotSupportedException();
		}

		@Override
		public <T> T unwrap(Class<T> iface) throws SQLException {
			if (iface.isInstance(this)) {
				return iface.cast(this);
			}
			throw new SQLException("Not a wrapper for " + iface.getName());
		}

		@Override
		public boolean isWrapperFor(Class<?> iface) {
			return iface.isInstance(this);
		}

	}

}
