package org.dbsp.sqlCompiler.compiler.sql.streaming;

import com.fasterxml.jackson.databind.JsonNode;
import org.dbsp.sqlCompiler.CompilerMain;
import org.dbsp.sqlCompiler.compiler.errors.CompilerMessages;
import org.dbsp.sqlCompiler.compiler.sql.StreamingTestBase;
import org.dbsp.sqlCompiler.compiler.sql.tools.BaseSQLTests;
import org.dbsp.util.Linq;
import org.dbsp.util.Utilities;
import org.junit.Assert;
import org.junit.Test;

import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.List;

/** Some tests that use the profiling API to measure perf at runtime */
public class ProfilingTests extends StreamingTestBase {
    Long[] measure(String program, String main) throws IOException, InterruptedException, SQLException {
        File script = createInputScript(program);
        CompilerMessages messages = CompilerMain.execute(
                "-o", BaseSQLTests.TEST_FILE_PATH, "--handles", "-i",
                script.getPath());
        messages.print();
        Assert.assertEquals(0, messages.errorCount());

        String mainFilePath = RUST_DIRECTORY + "/main.rs";
        File file = new File(mainFilePath);
        try (PrintWriter mainFile = new PrintWriter(file, StandardCharsets.UTF_8)) {
            mainFile.print(main);
        }
        try {
            BaseSQLTests.setupCargoLock();
            Utilities.compileAndTestRust(RUST_DIRECTORY, true, "--release");
            // After executing this Rust program the output is in file "mem.txt"
            // It contains three numbers: time taken (ms), memory used (bytes), and late records.
            String outFile = "mem.txt";
            Path outFilePath = Paths.get(RUST_DIRECTORY, "..", outFile);
            List<String> strings = Files.readAllLines(outFilePath);
            // System.out.println(strings);
            Assert.assertEquals(1, strings.size());
            String[] split = strings.get(0).split(",");
            Assert.assertEquals(3, split.length);
            Utilities.deleteFile(outFilePath.toFile(), true);
            return Linq.map(split, Long::parseLong, Long.class);
        } finally {
            File mainFile = new File(mainFilePath);
            Utilities.deleteFile(mainFile, true);
        }
    }

    String createMain(String rustDataGenerator) {
        String preamble = """
                #![allow(unused_imports)]

                use dbsp::{
                    algebra::F64,
                    circuit::{
                        CircuitConfig,
                        metrics::TOTAL_LATE_RECORDS,
                    },
                    utils::{Tup1, Tup2, Tup3},
                    zset,
                };

                use feldera_sqllib::{
                    append_to_collection_handle,
                    append_to_map_handle,
                    read_output_spine,
                    casts::{cast_to_Timestamp_s,handle_error,unwrap_value},
                    string::SqlString,
                    Timestamp,
                };

                use std::time::SystemTime;
                use std::sync::atomic::Ordering;
                use std::fs::File;
                use std::io::Write;
                use std::ops::Add;
                use temp::circuit;

                #[test]
                // Run the circuit generated by 'circuit' for a while then measure the
                // memory consumption.  Write the time taken and the memory used into
                // a file called "mem.txt".
                pub fn test() {
                    let (mut circuit, streams) = circuit(
                        CircuitConfig::with_workers(2)).expect("could not build circuit");
                    // uncomment if you want a CPU profile
                    // let _ = circuit.enable_cpu_profiler();
                    let start = SystemTime::now();
                """;
        String postamble = """
                    let late = TOTAL_LATE_RECORDS.load(Ordering::Relaxed);

                    let profile = circuit.retrieve_profile().expect("could not get profile");
                    let end = SystemTime::now();
                    let duration = end.duration_since(start).expect("could not get time");

                    // println!("{:?}", profile);
                    let mut data = String::new();
                    data.push_str(&format!("{},{},{}\\n",
                                           duration.as_millis(),
                                           profile.total_used_bytes().unwrap().bytes,
                                           late));
                    let mut file = File::create("mem.txt").expect("Could not create file");
                    file.write_all(data.as_bytes()).expect("Could not write data");
                    // println!("{:?},{:?},{:?}", duration, profile.total_used_bytes().unwrap(), late);
                }""";
        return preamble + rustDataGenerator + postamble;
    }

    static String stripLateness(String query) {
        String[] lines = query.split("\n");
        for (int i = 0; i < lines.length; i++) {
            String line = lines[i];
            int lateness = line.indexOf("LATENESS");
            if (lateness > 0) {
                boolean comma = line.endsWith(",");
                line = line.substring(0, lateness);
                if (comma) line += ",";
                lines[i] = line;
            }
        }
        return String.join("\n", lines);
    }

    @Test
    public void getJsonProfile() throws IOException, SQLException, InterruptedException {
        // Test the capability to get a profile in JSON from the circuit
        String sql = """
                CREATE TABLE T (x INT);
                CREATE VIEW V AS SELECT MAX(x) OVER () FROM T;""";
        String main = """
                use dbsp::{algebra::F64, circuit::CircuitConfig, utils::Tup1, zset};

                use feldera_sqllib::{
                    append_to_collection_handle, append_to_map_handle,
                    casts::{cast_to_Timestamp_s, handle_error},
                    read_output_spine,
                    string::SqlString,
                };

                use std::fs::File;
                use std::io::Write;
                use std::ops::Add;
                use std::sync::atomic::Ordering;
                use std::time::SystemTime;
                use temp::circuit;

                #[test]
                pub fn test() {
                    let (mut circuit, streams) = circuit(
                        CircuitConfig::with_workers(2)).expect("could not build circuit");
                    let _ = circuit.enable_cpu_profiler();
                    for i in 0..10000 {
                        let data = zset!(Tup1::new(Some(i)) => 1);
                        append_to_collection_handle(&data, &streams.0);
                        if i % 100 == 0 {
                            let _ = circuit.transaction().expect("could not run circuit");
                            let _ = &read_output_spine(&streams.2);
                        }
                    }

                    let profile = circuit.retrieve_profile().expect("could not get profile");
                    let mut file = File::create("profile.zip").expect("Could not create file");
                    file.write_all(&profile.as_json_zip()).expect("Could not write data");
                }""";
        File script = createInputScript(sql);
        CompilerMessages messages = CompilerMain.execute(
                "-o", BaseSQLTests.TEST_FILE_PATH, "--handles", "-i", "--ignoreOrder",
                script.getPath());
        messages.print();
        Assert.assertEquals(0, messages.errorCount());

        String mainFilePath = RUST_DIRECTORY + "/main.rs";
        File file = new File(mainFilePath);
        try (PrintWriter mainFile = new PrintWriter(file, StandardCharsets.UTF_8)) {
            mainFile.print(main);
        }

        try {
            BaseSQLTests.compileAndTestRust(true);
            if (!BaseSQLTests.skipRust) {
                String outFile = "profile.zip";
                Path outFilePath = Paths.get(RUST_DIRECTORY, "..", outFile);
                byte[] data = Utilities.extractFileFromZip(outFilePath.toString(), "profile.json");
                String str = new String(data, StandardCharsets.UTF_8);
                JsonNode jsonNode = Utilities.deterministicObjectMapper().readTree(str);
                Assert.assertNotNull(jsonNode);
                Utilities.deleteFile(outFilePath.toFile(), true);
            }
        } finally {
            File mainFile = new File(mainFilePath);
            Utilities.deleteFile(mainFile, true);
        }
    }

    public String getEmptyJsonProfile(File script) throws IOException, SQLException, InterruptedException {
        // Profile a SQL program from a specified file.
        // We don't know how to supply inputs, so we do this just to get the graph.
        String main = """
                use dbsp::circuit::CircuitConfig;

                use std::fs::File;
                use std::io::Write;
                use temp::circuit;

                #[test]
                pub fn test() {
                    let (mut circuit, _streams) = circuit(
                        CircuitConfig::with_workers(16)).expect("could not build circuit");
                    let _ = circuit.enable_cpu_profiler();
                    let profile = circuit.retrieve_profile().expect("could not get profile");
                    let mut file = File::create("profile.zip").expect("Could not create file");
                    file.write_all(&profile.as_json_zip()).expect("Could not write data");
                }""";

        File dataflow = new File("dataflow-" + script.getName().replace(".sql", ".json"));
        CompilerMessages messages = CompilerMain.execute(
                "-o", BaseSQLTests.TEST_FILE_PATH, "--handles", "-i", "--ignoreOrder",
                "--dataflow", dataflow.getPath(), script.getPath());
        messages.print();
        Assert.assertEquals(0, messages.errorCount());

        String mainFilePath = RUST_DIRECTORY + "/main.rs";
        File file = new File(mainFilePath);
        try (PrintWriter mainFile = new PrintWriter(file, StandardCharsets.UTF_8)) {
            mainFile.print(main);
        }

        try {
            BaseSQLTests.compileAndTestRust(true);
            String outFile = "profile.zip";
            Path outFilePath = Paths.get(RUST_DIRECTORY, "..", outFile);
            byte[] data = Utilities.extractFileFromZip(outFilePath.toString(), "profile.json");
            Utilities.deleteFile(outFilePath.toFile(), true);
            return new String(data, StandardCharsets.UTF_8);
        } finally {
            File mainFile = new File(mainFilePath);
            Utilities.deleteFile(mainFile, true);
        }
    }

    @Test
    public void profileLateness() throws IOException, InterruptedException, SQLException {
        String sql = """
                CREATE TABLE series (
                        distance DOUBLE,
                        pickup TIMESTAMP NOT NULL LATENESS INTERVAL '1:00' HOURS TO MINUTES
                );
                CREATE VIEW V AS
                SELECT MAX(distance), CAST(pickup AS DATE)
                FROM series GROUP BY CAST(pickup AS DATE);
                """;
        // Rust program which profiles the circuit.
        String main = this.createMain("""
                    // Initial data value for timestamp
                    let mut timestamp = handle_error(cast_to_Timestamp_s(SqlString::from_ref("2024-01-10 10:10:10")));
                    for i in 0..1000000 {
                        let value = Some(F64::new(i.into()));
                        timestamp = Timestamp::from_milliseconds(timestamp.milliseconds() + 20000);
                        let input = zset!(Tup2::new(value, timestamp) => 1);
                        append_to_collection_handle(&input, &streams.0);
                        if i % 1000 == 0 {
                            let _ = circuit.transaction().expect("could not run circuit");
                            let _ = &read_output_spine(&streams.1);
                        }
                    }""");
        this.profile(sql, main);
    }

    void profile(String sql, String main) throws SQLException, IOException, InterruptedException {
        if (BaseSQLTests.skipRust)
            return;
        Long[] p0 = this.measure(stripLateness(sql), main);
        Long[] p1 = this.measure(sql, main);
        // Memory consumption of program without lateness is expected to be higher
        // (because it cannot discard old data).
        if (p0[1] < 1.5 * p1[1]) {
            System.err.println("Profile statistics without and with lateness:");
            System.err.println(Arrays.toString(p0));
            System.err.println(Arrays.toString(p1));
            Assert.fail("Profiling test failed");
        }
        // No late records
        Assert.assertEquals(0L, (long)p0[2]);
        Assert.assertEquals(0L, (long)p1[2]);
    }

    @Test
    public void profileRetainValues() throws IOException, InterruptedException, SQLException {
        // Based on Q9.  Check whether integrate_trace_retain_values works as expected.
        String sql = """
                CREATE TABLE auction (
                   date_time TIMESTAMP NOT NULL LATENESS INTERVAL 1 MINUTE,
                   expires   TIMESTAMP NOT NULL,
                   id        INT
                );

                CREATE TABLE bid (
                   date_time TIMESTAMP NOT NULL LATENESS INTERVAL 1 MINUTE,
                   price INT,
                   auction INT
                );

                CREATE VIEW Q9 AS
                SELECT A.*, B.price, B.date_time AS bid_dateTime
                FROM auction A, bid B
                WHERE A.id = B.auction AND B.date_time BETWEEN A.date_time AND A.expires;
                """;
        // Rust program which profiles the circuit.
        String main = this.createMain("""
                    // Initial data value for timestamp
                    let mut timestamp = handle_error(cast_to_Timestamp_s(SqlString::from_ref("2024-01-10 10:10:10")));
                    for i in 0..1000000 {
                        let expire = Timestamp::from_milliseconds(timestamp.milliseconds() + 1000000);
                        timestamp = Timestamp::from_milliseconds(timestamp.milliseconds() + 20000);
                        let auction = zset!(Tup3::new(timestamp, expire, Some(i)) => 1);
                        append_to_collection_handle(&auction, &streams.0);
                        let bid = zset!(Tup3::new(Timestamp::from_milliseconds(timestamp.milliseconds() + 100), Some(i), Some(i)) => 1);
                        append_to_collection_handle(&bid, &streams.1);
                        if i % 100 == 0 {
                            let _ = circuit.transaction().expect("could not run circuit");
                            let _ = &read_output_spine(&streams.2);
                        }
                    }""");
        this.profile(sql, main);
    }

    @Test
    public void profileInputWithWaterline() throws IOException, InterruptedException, SQLException {
        // Check that the GC for InputMapWithWaterline works.
        String sql = """
                CREATE TABLE auction (
                   date_time TIMESTAMP NOT NULL LATENESS INTERVAL 1 MINUTE,
                   expires   TIMESTAMP NOT NULL,
                   id        INT NOT NULL PRIMARY KEY
                );

                CREATE VIEW V AS
                SELECT * FROM auction;
                """;
        // Rust program which profiles the circuit.
        String main = this.createMain("""
                    // Initial data value for timestamp
                    let mut timestamp = handle_error(cast_to_Timestamp_s(SqlString::from_ref("2024-01-10 10:10:10")));
                    for i in 0..1000000 {
                        let expire = Timestamp::from_milliseconds(timestamp.milliseconds() + 1000000);
                        timestamp = Timestamp::from_milliseconds(timestamp.milliseconds() + 20000);
                        let auction = zset!(Tup3::new(timestamp, expire, i) => 1);
                        append_to_map_handle(&auction, &streams.0, |x| Tup1::new(x.2));
                        if i % 100 == 0 {
                            let _ = circuit.transaction().expect("could not run circuit");
                            let _ = &read_output_spine(&streams.2);
                        }
                    }""");
        this.profile(sql, main);
    }

    @Test
    public void issue2228() throws SQLException, IOException, InterruptedException {
        String sql = """
                CREATE TABLE transaction (
                    id int NOT NULL,
                    unix_time BIGINT NOT NULL
                );

                CREATE TABLE feedback (
                    id int,
                    unix_time bigint LATENESS 3600 * 24
                );

                CREATE VIEW TRANSACTIONS AS
                SELECT t.*
                FROM transaction as t JOIN feedback as f
                ON t.id = f.id
                WHERE t.unix_time >= f.unix_time - 3600 * 24 * 7;
                """;

        String main = this.createMain("""
                let _ = circuit.transaction().expect("could not run circuit");
                let _ = circuit.transaction().expect("could not run circuit");
                """);
        if (BaseSQLTests.skipRust)
            return;
        this.measure(sql, main);
    }
}
