package com.lingniu.ingest.tdenginehistory;

import javax.sql.DataSource;
import java.io.IOException;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.sql.Timestamp;
import java.time.Instant;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;

public final class TdengineJdbcHistoryWriter implements TdengineHistoryWriter {

    private static final int MAX_LITERAL_ROWS_PER_STATEMENT = 200;

    private final DataSource dataSource;
    private final TdengineHistorySchema schema;
    private final TdengineHistoryStatements statements;
    private final Object schemaInitializationMonitor = new Object();
    private final Map<String, Instant> lastTimestampByInsertSql = new ConcurrentHashMap<>();
    private volatile boolean schemaInitialized;

    public TdengineJdbcHistoryWriter(DataSource dataSource, TdengineHistorySchema schema) {
        if (dataSource == null) {
            throw new IllegalArgumentException("dataSource must not be null");
        }
        if (schema == null) {
            throw new IllegalArgumentException("schema must not be null");
        }
        this.dataSource = dataSource;
        this.schema = schema;
        this.statements = new TdengineHistoryStatements(schema);
    }

    public void initializeSchema() throws IOException {
        inTransaction(connection -> {
            try (Statement statement = connection.createStatement()) {
                for (String sql : schema.bootstrapSql()) {
                    statement.execute(sql);
                }
            }
        });
        schemaInitialized = true;
    }

    @Override
    public void appendRawFrames(List<TdengineRawFrameRow> rows) throws IOException {
        append(rows, statements::rawFrame);
    }

    @Override
    public void appendLocations(List<TdengineLocationRow> rows) throws IOException {
        append(rows, statements::location);
    }

    private <T> void append(List<T> rows, Function<T, TdengineBatchStatement> mapper) throws IOException {
        if (rows == null || rows.isEmpty()) {
            return;
        }
        if (!schemaInitialized) {
            synchronized (schemaInitializationMonitor) {
                if (!schemaInitialized) {
                    append(rows, mapper, true);
                    schemaInitialized = true;
                    return;
                }
            }
        }
        append(rows, mapper, false);
    }

    private <T> void append(List<T> rows,
                            Function<T, TdengineBatchStatement> mapper,
                            boolean initializeSchema) throws IOException {
        inTransaction(connection -> {
            if (initializeSchema) {
                try (Statement statement = connection.createStatement()) {
                    for (String sql : schema.bootstrapSql()) {
                        statement.execute(sql);
                    }
                }
            }
            Set<String> createdChildTables = new LinkedHashSet<>();
            Map<String, LiteralBatch> literalBatches = new LinkedHashMap<>();
            for (T row : rows) {
                if (row == null) {
                    continue;
                }
                TdengineBatchStatement batch = mapper.apply(row);
                if (createdChildTables.add(batch.createChildTableSql())) {
                    try (Statement statement = connection.createStatement()) {
                        statement.execute(batch.createChildTableSql());
                    }
                }
                literalBatches.computeIfAbsent(batch.insertSql(), LiteralBatch::new)
                        .add(uniqueTimestampValues(batch.insertSql(), batch.values()));
            }
            for (LiteralBatch batch : literalBatches.values()) {
                executeLiteralBatch(connection, batch);
            }
        });
    }

    private List<Object> uniqueTimestampValues(String insertSql, List<Object> values) {
        if (values == null || values.isEmpty() || !(values.getFirst() instanceof Instant candidate)) {
            return values;
        }
        Instant uniqueTimestamp = lastTimestampByInsertSql.compute(insertSql, (key, lastTimestamp) -> {
            if (lastTimestamp == null || candidate.isAfter(lastTimestamp)) {
                return candidate;
            }
            return lastTimestamp.plusMillis(1);
        });
        if (candidate.equals(uniqueTimestamp)) {
            return values;
        }
        List<Object> copy = new java.util.ArrayList<>(values);
        copy.set(0, uniqueTimestamp);
        return copy;
    }

    private static void executeLiteralBatch(Connection connection, LiteralBatch batch) throws SQLException {
        try (Statement statement = connection.createStatement()) {
            List<List<Object>> rows = batch.rows();
            for (int start = 0; start < rows.size(); start += MAX_LITERAL_ROWS_PER_STATEMENT) {
                int end = Math.min(start + MAX_LITERAL_ROWS_PER_STATEMENT, rows.size());
                statement.execute(toLiteralInsertSql(batch.insertSql(), rows.subList(start, end)));
            }
        }
    }

    private static String toLiteralInsertSql(String insertSql, List<List<Object>> rows) {
        int valuesIndex = insertSql.lastIndexOf("VALUES");
        if (valuesIndex < 0) {
            throw new IllegalArgumentException("insert SQL must contain VALUES: " + insertSql);
        }
        if (rows == null || rows.isEmpty()) {
            throw new IllegalArgumentException("literal rows must not be empty");
        }
        StringBuilder sql = new StringBuilder(insertSql.substring(0, valuesIndex))
                .append("VALUES ");
        for (int rowIndex = 0; rowIndex < rows.size(); rowIndex++) {
            if (rowIndex > 0) {
                sql.append(' ');
            }
            List<Object> values = rows.get(rowIndex);
            sql.append('(');
            for (int i = 0; i < values.size(); i++) {
                if (i > 0) {
                    sql.append(", ");
                }
                sql.append(literal(values.get(i)));
            }
            sql.append(')');
        }
        return sql.toString();
    }

    private static String literal(Object value) {
        if (value == null) {
            return "NULL";
        }
        if (value instanceof Instant instant) {
            return Long.toString(instant.toEpochMilli());
        }
        if (value instanceof Timestamp timestamp) {
            return Long.toString(timestamp.toInstant().toEpochMilli());
        }
        if (value instanceof Double doubleValue) {
            return Double.isFinite(doubleValue) ? doubleValue.toString() : "NULL";
        }
        if (value instanceof Float floatValue) {
            return Float.isFinite(floatValue) ? floatValue.toString() : "NULL";
        }
        if (value instanceof Number number) {
            return number.toString();
        }
        return "'" + value.toString().replace("'", "''") + "'";
    }

    private void inTransaction(SqlWork work) throws IOException {
        try (Connection connection = dataSource.getConnection()) {
            boolean autoCommit = connection.getAutoCommit();
            connection.setAutoCommit(false);
            try {
                work.execute(connection);
                connection.commit();
            } catch (JdbcRuntimeException e) {
                rollback(connection);
                throw e.getCause();
            } catch (SQLException e) {
                rollback(connection);
                throw e;
            } finally {
                connection.setAutoCommit(autoCommit);
            }
        } catch (SQLException e) {
            throw new IOException("tdengine history write failed", e);
        }
    }

    private static void rollback(Connection connection) throws SQLException {
        connection.rollback();
    }

    @FunctionalInterface
    private interface SqlWork {
        void execute(Connection connection) throws SQLException;
    }

    private static final class JdbcRuntimeException extends RuntimeException {
        private JdbcRuntimeException(SQLException cause) {
            super(cause);
        }

        @Override
        public synchronized SQLException getCause() {
            return (SQLException) super.getCause();
        }
    }

    private static final class LiteralBatch {
        private final String insertSql;
        private final List<List<Object>> rows = new java.util.ArrayList<>();

        private LiteralBatch(String insertSql) {
            this.insertSql = insertSql;
        }

        private void add(List<Object> row) {
            rows.add(new java.util.ArrayList<>(row));
        }

        private String insertSql() {
            return insertSql;
        }

        private List<List<Object>> rows() {
            return rows;
        }
    }
}
