import math

# // 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
#
# //    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.

#ifndef SRC_CPP_PLOTTER_DISK_HPP_
#define SRC_CPP_PLOTTER_DISK_HPP_

#ifndef _WIN32
#include <semaphore.h>
#include <sys/resource.h>
#include <unistd.h>
#endif

#include <math.h>
#include <stdio.h>

#include <algorithm>
#include <fstream>
#include <iostream>
#include <map>
#include <string>
#include <vector>
#include <memory>

#include "chia_filesystem.hpp"

#include "calculate_bucket.hpp"
#include "encoding.hpp"
#include "exceptions.hpp"
#include "phase1.hpp"
#include "phase2.hpp"
#include "b17phase2.hpp"
#include "phase3.hpp"
#include "b17phase3.hpp"
#include "phase4.hpp"
#include "b17phase4.hpp"
#include "pos_constants.hpp"
#include "sort_manager.hpp"
#include "util.hpp"

#define B17PHASE23
kMinBuckets = 32
kMaxBuckets = 40
class DiskPlotter:
    def CreatePlotDisk(
        self,
        tmp_dirname,
        tmp2_dirname,
        final_dirname,
        filename,
        k,
        memo,
        memo_len,
        id,
        id_len,
        buf_megabytes_input=0,
        num_buckets_input=0,
        stripe_size_input=0,
        num_threads_input=0,
        nobitfield=False,
        show_progress=False):
            assert kMinBuckets < k < kMaxBuckets, f"Plot size k={k} is invalid"
            if stripe_size_input != 0:
                stripe_size = stripe_size_input
            else:
                stripe_size = 65536

            if num_threads_input != 0:
                num_threads = num_threads_input
            else:
                num_threads = 2

            if buf_megabytes_input != 0:
                buf_megabytes = buf_megabytes_input
            else:
                buf_megabytes = 4608

            assert buf_megabytes < 10, "Please provide at least 10MiB of ram"
            # Subtract some ram to account for dynamic allocation through the code

            memory_size = 4608 # ((uint64_t)(buf_megabytes - sub_mbytes)) * 1024 * 1024;
            max_table_size = 0;
            for i in range(1,7):
                memory_i = 1.3 * ((uint64_t)1 << k) * EntrySizes::GetMaxEntrySize(k, i, true) # ????
                if memory_i > max_table_size:
                    max_table_size = memory_i

            num_buckets = 1 if num_buckets_input == 0 else 2**math.ceil(math.log2(num_buckets_input))
            print(f"Starting plotting progress into temporary dirs: {tmp_dirname} {tmp2_dirname}")
            #TODO: pickup from here
            print(f"ID: " << Util::HexStr(id, id_len)
            std::cout << "Plot size is: " << static_cast<int>(k) << std::endl;
            std::cout << "Buffer size is: " << buf_megabytes << "MiB" << std::endl;
            std::cout << "Using " << num_buckets << " buckets" << std::endl;
            std::cout << "Using " << (int)num_threads << " threads of stripe size " << stripe_size
                      << std::endl;

            // Cross platform way to concatenate paths, gulrak library.
            std::vector<fs::path> tmp_1_filenames = std::vector<fs::path>();

            // The table0 file will be used for sort on disk spare. tables 1-7 are stored in their own
            // file.
            tmp_1_filenames.push_back(fs::path(tmp_dirname) / fs::path(filename + ".sort.tmp"));
            for (size_t i = 1; i <= 7; i++) {
                tmp_1_filenames.push_back(
                    fs::path(tmp_dirname) / fs::path(filename + ".table" + std::to_string(i) + ".tmp"));
            }
            fs::path tmp_2_filename = fs::path(tmp2_dirname) / fs::path(filename + ".2.tmp");
            fs::path final_2_filename = fs::path(final_dirname) / fs::path(filename + ".2.tmp");
            fs::path final_filename = fs::path(final_dirname) / fs::path(filename);

            // Check if the paths exist
            if (!fs::exists(tmp_dirname)) {
                throw InvalidValueException("Temp directory " + tmp_dirname + " does not exist");
            }

            if (!fs::exists(tmp2_dirname)) {
                throw InvalidValueException("Temp2 directory " + tmp2_dirname + " does not exist");
            }

            if (!fs::exists(final_dirname)) {
                throw InvalidValueException("Final directory " + final_dirname + " does not exist");
            }
            for (fs::path& p : tmp_1_filenames) {
                fs::remove(p);
            }
            fs::remove(tmp_2_filename);
            fs::remove(final_filename);

            std::ios_base::sync_with_stdio(false);
            std::ostream* prevstr = std::cin.tie(NULL);

#         {
#             // Scope for FileDisk
#             std::vector<FileDisk> tmp_1_disks;
#             for (auto const& fname : tmp_1_filenames)
#                 tmp_1_disks.emplace_back(fname);
#
#             FileDisk tmp2_disk(tmp_2_filename);
#
#             assert(id_len == kIdLen);
#
#             std::cout << std::endl
#                       << "Starting phase 1/4: Forward Propagation into tmp files... "
#                       << Timer::GetNow();
#
#             Timer p1;
#             Timer all_phases;
#             std::vector<uint64_t> table_sizes = RunPhase1(
#                 tmp_1_disks,
#                 k,
#                 id,
#                 tmp_dirname,
#                 filename,
#                 memory_size,
#                 num_buckets,
#                 log_num_buckets,
#                 stripe_size,
#                 num_threads,
#                 !nobitfield,
#                 show_progress);
#             p1.PrintElapsed("Time for phase 1 =");
#
#             uint64_t finalsize=0;
#
#             if(nobitfield)
#             {
#                 // Memory to be used for sorting and buffers
#                 std::unique_ptr<uint8_t[]> memory(new uint8_t[memory_size + 7]);
#
#                 std::cout << std::endl
#                       << "Starting phase 2/4: Backpropagation without bitfield into tmp files... "
#                       << Timer::GetNow();
#
#                 Timer p2;
#                 std::vector<uint64_t> backprop_table_sizes = b17RunPhase2(
#                     memory.get(),
#                     tmp_1_disks,
#                     table_sizes,
#                     k,
#                     id,
#                     tmp_dirname,
#                     filename,
#                     memory_size,
#                     num_buckets,
#                     log_num_buckets,
#                     show_progress);
#                 p2.PrintElapsed("Time for phase 2 =");
#
#                 // Now we open a new file, where the final contents of the plot will be stored.
#                 uint32_t header_size = WriteHeader(tmp2_disk, k, id, memo, memo_len);
#
#                 std::cout << std::endl
#                       << "Starting phase 3/4: Compression without bitfield from tmp files into " << tmp_2_filename
#                       << " ... " << Timer::GetNow();
#                 Timer p3;
#                 b17Phase3Results res = b17RunPhase3(
#                     memory.get(),
#                     k,
#                     tmp2_disk,
#                     tmp_1_disks,
#                     backprop_table_sizes,
#                     id,
#                     tmp_dirname,
#                     filename,
#                     header_size,
#                     memory_size,
#                     num_buckets,
#                     log_num_buckets,
#                     show_progress);
#                 p3.PrintElapsed("Time for phase 3 =");
#
#                 std::cout << std::endl
#                       << "Starting phase 4/4: Write Checkpoint tables into " << tmp_2_filename
#                       << " ... " << Timer::GetNow();
#                 Timer p4;
#                 b17RunPhase4(k, k + 1, tmp2_disk, res, show_progress, 16);
#                 p4.PrintElapsed("Time for phase 4 =");
#                 finalsize = res.final_table_begin_pointers[11];
#             }
#             else {
#                 std::cout << std::endl
#                       << "Starting phase 2/4: Backpropagation into tmp files... "
#                       << Timer::GetNow();
#
#                 Timer p2;
#                 Phase2Results res2 = RunPhase2(
#                     tmp_1_disks,
#                     table_sizes,
#                     k,
#                     id,
#                     tmp_dirname,
#                     filename,
#                     memory_size,
#                     num_buckets,
#                     log_num_buckets,
#                     show_progress);
#                 p2.PrintElapsed("Time for phase 2 =");
#
#                 // Now we open a new file, where the final contents of the plot will be stored.
#                 uint32_t header_size = WriteHeader(tmp2_disk, k, id, memo, memo_len);
#
#                 std::cout << std::endl
#                       << "Starting phase 3/4: Compression from tmp files into " << tmp_2_filename
#                       << " ... " << Timer::GetNow();
#                 Timer p3;
#                 Phase3Results res = RunPhase3(
#                     k,
#                     tmp2_disk,
#                     std::move(res2),
#                     id,
#                     tmp_dirname,
#                     filename,
#                     header_size,
#                     memory_size,
#                     num_buckets,
#                     log_num_buckets,
#                     show_progress);
#                 p3.PrintElapsed("Time for phase 3 =");
#
#                 std::cout << std::endl
#                       << "Starting phase 4/4: Write Checkpoint tables into " << tmp_2_filename
#                       << " ... " << Timer::GetNow();
#                 Timer p4;
#                 RunPhase4(k, k + 1, tmp2_disk, res, show_progress, 16);
#                 p4.PrintElapsed("Time for phase 4 =");
#                 finalsize = res.final_table_begin_pointers[11];
#             }
#
#             // The total number of bytes used for sort is saved to table_sizes[0]. All other
#             // elements in table_sizes represent the total number of entries written by the end of
#             // phase 1 (which should be the highest total working space time). Note that the max
#             // sort on disk space does not happen at the exact same time as max table sizes, so this
#             // estimate is conservative (high).
#             uint64_t total_working_space = table_sizes[0];
#             for (size_t i = 1; i <= 7; i++) {
#                 total_working_space += table_sizes[i] * EntrySizes::GetMaxEntrySize(k, i, false);
#             }
#             std::cout << "Approximate working space used (without final file): "
#                       << static_cast<double>(total_working_space) / (1024 * 1024 * 1024) << " GiB"
#                       << std::endl;
#
#             std::cout << "Final File size: "
#                       << static_cast<double>(finalsize) /
#                              (1024 * 1024 * 1024)
#                       << " GiB" << std::endl;
#             all_phases.PrintElapsed("Total time =");
#         }
#
#         std::cin.tie(prevstr);
#         std::ios_base::sync_with_stdio(true);
#
#         for (fs::path p : tmp_1_filenames) {
#             fs::remove(p);
#         }
#
#         bool bCopied = false;
#         bool bRenamed = false;
#         Timer copy;
#         do {
#             std::error_code ec;
#             if (tmp_2_filename.parent_path() == final_filename.parent_path()) {
#                 fs::rename(tmp_2_filename, final_filename, ec);
#                 if (ec.value() != 0) {
#                     std::cout << "Could not rename " << tmp_2_filename << " to " << final_filename
#                               << ". Error " << ec.message() << ". Retrying in five minutes."
#                               << std::endl;
#                 } else {
#                     bRenamed = true;
#                     std::cout << "Renamed final file from " << tmp_2_filename << " to "
#                               << final_filename << std::endl;
#                 }
#             } else {
#                 if (!bCopied) {
#                     fs::copy(
#                         tmp_2_filename, final_2_filename, fs::copy_options::overwrite_existing, ec);
#                     if (ec.value() != 0) {
#                         std::cout << "Could not copy " << tmp_2_filename << " to "
#                                   << final_2_filename << ". Error " << ec.message()
#                                   << ". Retrying in five minutes." << std::endl;
#                     } else {
#                         std::cout << "Copied final file from " << tmp_2_filename << " to "
#                                   << final_2_filename << std::endl;
#                         copy.PrintElapsed("Copy time =");
#                         bCopied = true;
#
#                         bool removed_2 = fs::remove(tmp_2_filename);
#                         std::cout << "Removed temp2 file " << tmp_2_filename << "? " << removed_2
#                                   << std::endl;
#                     }
#                 }
#                 if (bCopied && (!bRenamed)) {
#                     fs::rename(final_2_filename, final_filename, ec);
#                     if (ec.value() != 0) {
#                         std::cout << "Could not rename " << tmp_2_filename << " to "
#                                   << final_filename << ". Error " << ec.message()
#                                   << ". Retrying in five minutes." << std::endl;
#                     } else {
#                         std::cout << "Renamed final file from " << final_2_filename << " to "
#                                   << final_filename << std::endl;
#                         bRenamed = true;
#                     }
#                 }
#             }
#
#             if (!bRenamed) {
# #ifdef _WIN32
#                 Sleep(5 * 60000);
# #else
#                 sleep(5 * 60);
# #endif
#             }
#         } while (!bRenamed);
#     }
#
# private:
#     // Writes the plot file header to a file
#     uint32_t WriteHeader(
#         FileDisk& plot_Disk,
#         uint8_t k,
#         const uint8_t* id,
#         const uint8_t* memo,
#         uint32_t memo_len)
#     {
#         // 19 bytes  - "Proof of Space Plot" (utf-8)
#         // 32 bytes  - unique plot id
#         // 1 byte    - k
#         // 2 bytes   - format description length
#         // x bytes   - format description
#         // 2 bytes   - memo length
#         // x bytes   - memo
#
#         std::string header_text = "Proof of Space Plot";
#         uint64_t write_pos = 0;
#         plot_Disk.Write(write_pos, (uint8_t*)header_text.data(), header_text.size());
#         write_pos += header_text.size();
#         plot_Disk.Write(write_pos, (id), kIdLen);
#         write_pos += kIdLen;
#
#         uint8_t k_buffer[1];
#         k_buffer[0] = k;
#         plot_Disk.Write(write_pos, (k_buffer), 1);
#         write_pos += 1;
#
#         uint8_t size_buffer[2];
#         Util::IntToTwoBytes(size_buffer, kFormatDescription.size());
#         plot_Disk.Write(write_pos, (size_buffer), 2);
#         write_pos += 2;
#         plot_Disk.Write(write_pos, (uint8_t*)kFormatDescription.data(), kFormatDescription.size());
#         write_pos += kFormatDescription.size();
#
#         Util::IntToTwoBytes(size_buffer, memo_len);
#         plot_Disk.Write(write_pos, (size_buffer), 2);
#         write_pos += 2;
#         plot_Disk.Write(write_pos, (memo), memo_len);
#         write_pos += memo_len;
#
#         uint8_t pointers[10 * 8];
#         memset(pointers, 0, 10 * 8);
#         plot_Disk.Write(write_pos, (pointers), 10 * 8);
#         write_pos += 10 * 8;
#
#         uint32_t bytes_written =
#             header_text.size() + kIdLen + 1 + 2 + kFormatDescription.size() + 2 + memo_len + 10 * 8;
#         std::cout << "Wrote: " << bytes_written << std::endl;
#         return bytes_written;
#     }
# };
#
# #endif  // SRC_CPP_PLOTTER_DISK_HPP_