#include "../../src/TrainingArchive/TrainingArchive.h"
#include "../../src/TransformerRanker/TransformerRanker.h"

#include <algorithm>
#include <array>
#include <cctype>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <iostream>
#include <limits>
#include <optional>
#include <set>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#ifdef _WIN32
#include <fcntl.h>
#include <io.h>
#include <windows.h>
#endif

namespace {

namespace training = szilassi::training;
namespace transformer = szilassi::transformer;

constexpr std::array<char, 8> kExportMagic{{'S', 'Z', 'T', 'X', 'P', '0', '0', '1'}};
constexpr std::uint32_t kExportVersion = 1;
constexpr std::size_t kFaceCount = 12;
constexpr std::size_t kFaceFeatureCount = 16;
constexpr std::size_t kGlobalFeatureCount = 24;
constexpr std::size_t kFeatureFloatCount =
    kFaceCount * kFaceFeatureCount + kGlobalFeatureCount;

// Stable packed little-endian record layout (there is no native-struct padding):
//   0    u8[16]  canonical UUID bytes
//   16   u64     record_sequence
//   24   u32     topology
//   28   f32[216] proposal features (12x16 face, then 24 global)
//   892  f32[216] anchor features (12x16 face, then 24 global)
//   1756 i32     crossing_gain = anchor.C - result.C
//   1760 i32     intersection_gain = anchor.I - result.I
//   1764 i32     total_defect_gain
//   1768 u8      defect_improved (total_defect_gain > 0)
//   1769 u8      objective_better (TrajectoryImproved)
//   1770 u8      defect_count_tie (total_defect_gain == 0)
//   1771 u8      entered_archive
//   1772 i32[3]  plane, move, scale actions (-1 means absent)
//   1784 f32[3]  actual plane, move, scale propensities
//   1796 u32     replica_count / assigned chains
//   1800 u64     rollout_iterations
//   1808 u32     InjectedTrajectory flags
//   1812 u32     anchor Metrics flags
//   1816 u32     injected/proposal Metrics flags
//   1820 u32     result Metrics flags (contains MetricDdVerified when applicable)
constexpr std::uint32_t kExportRecordSize = 1824;

using FeatureFaceArray =
    decltype(std::declval<transformer::Features>().face);
using FeatureFaceRow = typename FeatureFaceArray::value_type;
using FeatureGlobalArray =
    decltype(std::declval<transformer::Features>().global);
static_assert(std::tuple_size<FeatureFaceArray>::value == kFaceCount,
              "Transformer face count changed; bump the export format");
static_assert(std::tuple_size<FeatureFaceRow>::value == kFaceFeatureCount,
              "Transformer face feature count changed; bump the export format");
static_assert(std::tuple_size<FeatureGlobalArray>::value == kGlobalFeatureCount,
              "Transformer global feature count changed; bump the export format");

struct Options {
    std::filesystem::path root;
    std::set<std::string> included_runs;
    bool root_supplied = false;
};

struct ShardSnapshot {
    std::filesystem::path path;
    std::string run_name;
    training::RunId expected_run_id{};
    std::optional<std::uint64_t> filename_index;
    std::uint64_t size = 0;
};

void print_usage() {
    std::cerr
        << "Usage: TransformerTrainingExport --root <results/search> "
           "[--include-run <UUID>]...\n"
        << "Writes SZTXP001 version-1 packed little-endian records to stdout.\n";
}

int hexadecimal_digit(char value) {
    if (value >= '0' && value <= '9') return value - '0';
    if (value >= 'a' && value <= 'f') return 10 + value - 'a';
    if (value >= 'A' && value <= 'F') return 10 + value - 'A';
    return -1;
}

bool parse_canonical_uuid(
    const std::string& text,
    std::string& normalized,
    training::RunId& run_id
) {
    constexpr std::array<std::size_t, 4> kHyphens{{8, 13, 18, 23}};
    if (text.size() != 36) return false;
    for (std::size_t position : kHyphens) {
        if (text[position] != '-') return false;
    }
    normalized.clear();
    normalized.reserve(text.size());
    std::array<char, 32> hexadecimal{};
    std::size_t digit_count = 0;
    for (std::size_t index = 0; index < text.size(); ++index) {
        if (std::find(kHyphens.begin(), kHyphens.end(), index) != kHyphens.end()) {
            normalized.push_back('-');
            continue;
        }
        const int digit = hexadecimal_digit(text[index]);
        if (digit < 0 || digit_count >= hexadecimal.size()) return false;
        const char lower = static_cast<char>(std::tolower(
            static_cast<unsigned char>(text[index])));
        hexadecimal[digit_count++] = lower;
        normalized.push_back(lower);
    }
    if (digit_count != hexadecimal.size()) return false;
    for (std::size_t index = 0; index < run_id.bytes.size(); ++index) {
        const int high = hexadecimal_digit(hexadecimal[index * 2]);
        const int low = hexadecimal_digit(hexadecimal[index * 2 + 1]);
        run_id.bytes[index] = static_cast<std::uint8_t>((high << 4) | low);
    }
    return true;
}

bool parse_options(int argc, char* argv[], Options& options) {
    for (int index = 1; index < argc; ++index) {
        const std::string argument = argv[index];
        if (argument == "--help" || argument == "-h") {
            print_usage();
            return false;
        }
        if (argument == "--root") {
            if (++index >= argc || options.root_supplied) {
                std::cerr << "--root requires exactly one path\n";
                return false;
            }
            options.root = std::filesystem::path(argv[index]);
            options.root_supplied = true;
            continue;
        }
        if (argument == "--include-run") {
            if (++index >= argc) {
                std::cerr << "--include-run requires a canonical UUID\n";
                return false;
            }
            std::string normalized;
            training::RunId ignored;
            if (!parse_canonical_uuid(argv[index], normalized, ignored)) {
                std::cerr << "Invalid run UUID: " << argv[index] << '\n';
                return false;
            }
            options.included_runs.insert(std::move(normalized));
            continue;
        }
        std::cerr << "Unknown argument: " << argument << '\n';
        return false;
    }
    if (!options.root_supplied || options.root.empty()) {
        std::cerr << "--root is required\n";
        return false;
    }
    return true;
}

std::optional<std::uint64_t> parse_shard_filename(
    const std::filesystem::path& path
) {
    const std::string filename = path.filename().string();
    constexpr const char* kPrefix = "shard_";
    constexpr const char* kSuffix = ".sztd";
    constexpr std::size_t kPrefixSize = 6;
    constexpr std::size_t kSuffixSize = 5;
    if (filename.size() <= kPrefixSize + kSuffixSize ||
        filename.compare(0, kPrefixSize, kPrefix) != 0 ||
        filename.compare(filename.size() - kSuffixSize, kSuffixSize, kSuffix) != 0) {
        return std::nullopt;
    }
    const std::string digits = filename.substr(
        kPrefixSize,
        filename.size() - kPrefixSize - kSuffixSize);
    std::uint64_t value = 0;
    for (char digit : digits) {
        if (digit < '0' || digit > '9' ||
            value > (std::numeric_limits<std::uint64_t>::max() - 9U) / 10U) {
            return std::nullopt;
        }
        value = value * 10U + static_cast<unsigned>(digit - '0');
    }
    return value;
}

bool snapshot_shards(
    const Options& options,
    std::vector<ShardSnapshot>& shards,
    std::string& error
) {
    std::error_code path_error;
    const std::filesystem::path absolute_root =
        std::filesystem::absolute(options.root, path_error);
    if (path_error) {
        error = "Cannot resolve archive root: " + path_error.message();
        return false;
    }
    const std::filesystem::path runs_root = absolute_root / "runs";
    if (!std::filesystem::is_directory(runs_root, path_error) || path_error) {
        error = "Archive runs directory does not exist: " + runs_root.string();
        return false;
    }

    struct RunDirectory {
        std::filesystem::path path;
        std::string normalized_name;
        training::RunId run_id{};
    };
    std::vector<RunDirectory> runs;
    std::set<std::string> found_requested_runs;
    std::filesystem::directory_iterator iterator(runs_root, path_error);
    const std::filesystem::directory_iterator end;
    if (path_error) {
        error = "Cannot enumerate archive runs: " + path_error.message();
        return false;
    }
    while (iterator != end) {
        std::error_code type_error;
        const bool is_directory = iterator->is_directory(type_error);
        if (type_error) {
            error = "Cannot inspect run directory entry: " + type_error.message();
            return false;
        }
        if (is_directory) {
            std::string normalized;
            training::RunId run_id;
            const std::string name = iterator->path().filename().string();
            if (parse_canonical_uuid(name, normalized, run_id) &&
                (options.included_runs.empty() ||
                 options.included_runs.count(normalized) != 0)) {
                runs.push_back({iterator->path(), normalized, run_id});
                found_requested_runs.insert(normalized);
            }
        }
        iterator.increment(path_error);
        if (path_error) {
            error = "Cannot continue enumerating archive runs: " + path_error.message();
            return false;
        }
    }
    for (const std::string& requested : options.included_runs) {
        if (found_requested_runs.count(requested) == 0) {
            error = "Requested run does not exist below the archive root: " + requested;
            return false;
        }
    }
    std::sort(runs.begin(), runs.end(), [](const RunDirectory& left, const RunDirectory& right) {
        return left.normalized_name < right.normalized_name;
    });

    for (const RunDirectory& run : runs) {
        const std::filesystem::path training_directory = run.path / "training";
        std::error_code exists_error;
        const bool has_training =
            std::filesystem::is_directory(training_directory, exists_error);
        if (exists_error) {
            error = "Cannot inspect training directory for run " +
                run.normalized_name + ": " + exists_error.message();
            return false;
        }
        if (!has_training) {
            if (!options.included_runs.empty()) {
                error = "Requested run has no training directory: " + run.normalized_name;
                return false;
            }
            continue;
        }
        std::filesystem::directory_iterator shard_iterator(
            training_directory,
            exists_error);
        if (exists_error) {
            error = "Cannot enumerate training directory for run " +
                run.normalized_name + ": " + exists_error.message();
            return false;
        }
        while (shard_iterator != end) {
            std::error_code file_error;
            const bool regular = shard_iterator->is_regular_file(file_error);
            if (file_error) {
                error = "Cannot inspect training file: " + file_error.message();
                return false;
            }
            const std::filesystem::path path = shard_iterator->path();
            if (regular && path.extension() == ".sztd") {
                const std::uint64_t size = shard_iterator->file_size(file_error);
                if (file_error) {
                    error = "Cannot size training shard: " + file_error.message();
                    return false;
                }
                shards.push_back({
                    path,
                    run.normalized_name,
                    run.run_id,
                    parse_shard_filename(path),
                    size});
            }
            shard_iterator.increment(exists_error);
            if (exists_error) {
                error = "Cannot continue enumerating training shards: " +
                    exists_error.message();
                return false;
            }
        }
    }

    std::sort(shards.begin(), shards.end(), [](const ShardSnapshot& left,
                                                const ShardSnapshot& right) {
        if (left.run_name != right.run_name) return left.run_name < right.run_name;
        if (left.filename_index.has_value() != right.filename_index.has_value()) {
            return left.filename_index.has_value();
        }
        if (left.filename_index && right.filename_index &&
            *left.filename_index != *right.filename_index) {
            return *left.filename_index < *right.filename_index;
        }
        return left.path.filename().string() < right.path.filename().string();
    });
    if (shards.empty()) {
        error = "No sealed .sztd shards matched the requested runs";
        return false;
    }
    return true;
}

class PackedRecord {
public:
    void bytes(const std::uint8_t* values, std::size_t count) {
        if (!room(count)) return;
        std::memcpy(data_.data() + size_, values, count);
        size_ += count;
    }

    void u8(std::uint8_t value) {
        if (!room(1)) return;
        data_[size_++] = value;
    }

    void u32(std::uint32_t value) {
        if (!room(4)) return;
        for (std::size_t byte = 0; byte < 4; ++byte) {
            data_[size_++] = static_cast<std::uint8_t>(value >> (byte * 8));
        }
    }

    void i32(std::int32_t value) {
        std::uint32_t bits = 0;
        std::memcpy(&bits, &value, sizeof(bits));
        u32(bits);
    }

    void u64(std::uint64_t value) {
        if (!room(8)) return;
        for (std::size_t byte = 0; byte < 8; ++byte) {
            data_[size_++] = static_cast<std::uint8_t>(value >> (byte * 8));
        }
    }

    void f32(float value) {
        std::uint32_t bits = 0;
        static_assert(sizeof(bits) == sizeof(value), "Unexpected float width");
        std::memcpy(&bits, &value, sizeof(bits));
        u32(bits);
    }

    const std::uint8_t* data() const { return data_.data(); }
    std::size_t size() const { return size_; }
    bool overflowed() const { return overflowed_; }

private:
    bool room(std::size_t count) {
        if (count > data_.size() - size_) {
            overflowed_ = true;
            return false;
        }
        return true;
    }

    std::array<std::uint8_t, kExportRecordSize> data_{};
    std::size_t size_ = 0;
    bool overflowed_ = false;
};

bool append_features(
    PackedRecord& packed,
    const transformer::Features& features,
    std::uint32_t topology,
    std::string& error
) {
    if (features.topology != topology) {
        error = "Transformer feature topology does not match the archive record";
        return false;
    }
    for (const auto& face : features.face) {
        for (float value : face) {
            if (!std::isfinite(value)) {
                error = "Transformer face feature is not finite";
                return false;
            }
            packed.f32(value);
        }
    }
    for (float value : features.global) {
        if (!std::isfinite(value)) {
            error = "Transformer global feature is not finite";
            return false;
        }
        packed.f32(value);
    }
    return true;
}

bool checked_i32(std::int64_t value, std::int32_t& result) {
    if (value < std::numeric_limits<std::int32_t>::min() ||
        value > std::numeric_limits<std::int32_t>::max()) {
        return false;
    }
    result = static_cast<std::int32_t>(value);
    return true;
}

bool encode_trajectory(
    const training::InjectedTrajectory& trajectory,
    const training::RunId& expected_run_id,
    PackedRecord& packed,
    std::string& error
) {
    if (trajectory.context.run_id.bytes != expected_run_id.bytes) {
        error = "Training record run UUID does not match its parent run directory";
        return false;
    }

    const transformer::SearchBudget budget{
        trajectory.rollout_iterations,
        trajectory.replica_count,
        trajectory.replica_count,
        trajectory.pool_size};
    const bool guided =
        (trajectory.flags & training::TrajectoryNeuralGuided) != 0;
    const transformer::FeatureInput proposal_input{
        trajectory.injected_state,
        trajectory.injected_metrics,
        budget,
        trajectory.context.topology,
        guided};
    const transformer::FeatureInput anchor_input{
        trajectory.anchor_state,
        trajectory.anchor_metrics,
        budget,
        trajectory.context.topology,
        guided};
    transformer::Features proposal_features;
    transformer::Features anchor_features;
    std::string feature_error;
    if (!transformer::make_features(
            proposal_input,
            proposal_features,
            &feature_error)) {
        error = "Cannot build proposal features: " + feature_error;
        return false;
    }
    if (!transformer::make_features(
            anchor_input,
            anchor_features,
            &feature_error)) {
        error = "Cannot build anchor features: " + feature_error;
        return false;
    }

    const std::int64_t crossing_gain_wide =
        static_cast<std::int64_t>(trajectory.anchor_metrics.crossings) -
        static_cast<std::int64_t>(trajectory.result_metrics.crossings);
    const std::int64_t intersection_gain_wide =
        static_cast<std::int64_t>(trajectory.anchor_metrics.intersections) -
        static_cast<std::int64_t>(trajectory.result_metrics.intersections);
    const std::int64_t defect_gain_wide =
        crossing_gain_wide + intersection_gain_wide;
    std::int32_t crossing_gain = 0;
    std::int32_t intersection_gain = 0;
    std::int32_t defect_gain = 0;
    if (!checked_i32(crossing_gain_wide, crossing_gain) ||
        !checked_i32(intersection_gain_wide, intersection_gain) ||
        !checked_i32(defect_gain_wide, defect_gain)) {
        error = "Trajectory label does not fit the export format";
        return false;
    }

    packed.bytes(trajectory.context.run_id.bytes.data(),
                 trajectory.context.run_id.bytes.size());
    packed.u64(trajectory.context.record_sequence);
    packed.u32(trajectory.context.topology);
    if (!append_features(
            packed,
            proposal_features,
            trajectory.context.topology,
            error) ||
        !append_features(
            packed,
            anchor_features,
            trajectory.context.topology,
            error)) {
        return false;
    }
    packed.i32(crossing_gain);
    packed.i32(intersection_gain);
    packed.i32(defect_gain);
    packed.u8(defect_gain > 0 ? 1U : 0U);
    packed.u8((trajectory.flags & training::TrajectoryImproved) != 0 ? 1U : 0U);
    packed.u8(defect_gain == 0 ? 1U : 0U);
    packed.u8((trajectory.flags & training::TrajectoryEnteredArchive) != 0 ? 1U : 0U);
    packed.i32(trajectory.plane_action);
    packed.i32(trajectory.move_action);
    packed.i32(trajectory.scale_action);
    packed.f32(trajectory.plane_propensity);
    packed.f32(trajectory.move_propensity);
    packed.f32(trajectory.scale_propensity);
    packed.u32(trajectory.replica_count);
    packed.u64(trajectory.rollout_iterations);
    packed.u32(trajectory.flags);
    packed.u32(trajectory.anchor_metrics.flags);
    packed.u32(trajectory.injected_metrics.flags);
    packed.u32(trajectory.result_metrics.flags);
    if (packed.overflowed() || packed.size() != kExportRecordSize) {
        error = "Internal export record-size mismatch";
        return false;
    }
    return true;
}

bool write_u32(std::ostream& output, std::uint32_t value) {
    std::array<char, 4> bytes{};
    for (std::size_t index = 0; index < bytes.size(); ++index) {
        bytes[index] = static_cast<char>(
            static_cast<unsigned char>(value >> (index * 8)));
    }
    output.write(bytes.data(), static_cast<std::streamsize>(bytes.size()));
    return static_cast<bool>(output);
}

bool write_header(std::ostream& output) {
    output.write(kExportMagic.data(), static_cast<std::streamsize>(kExportMagic.size()));
    return static_cast<bool>(output) &&
        write_u32(output, kExportVersion) &&
        write_u32(output, kExportRecordSize);
}

void lower_windows_priority() {
#ifdef _WIN32
    if (SetPriorityClass(GetCurrentProcess(), BELOW_NORMAL_PRIORITY_CLASS) == 0) {
        std::cerr << "Warning: cannot set below-normal process priority (Windows error "
                  << GetLastError() << ")\n";
    }
#endif
}

bool enable_binary_stdout() {
#ifdef _WIN32
    return _setmode(_fileno(stdout), _O_BINARY) != -1;
#else
    return true;
#endif
}

} // namespace

int main(int argc, char* argv[]) {
    Options options;
    if (!parse_options(argc, argv, options)) {
        return 2;
    }
    lower_windows_priority();
    if (!enable_binary_stdout()) {
        std::cerr << "Cannot switch stdout to binary mode\n";
        return 3;
    }

    std::vector<ShardSnapshot> shards;
    std::string error;
    if (!snapshot_shards(options, shards, error)) {
        std::cerr << "Cannot snapshot training shards: " << error << '\n';
        return 4;
    }
    std::uint64_t snapshot_bytes = 0;
    for (const ShardSnapshot& shard : shards) {
        if (shard.size > std::numeric_limits<std::uint64_t>::max() - snapshot_bytes) {
            std::cerr << "Training shard snapshot byte count overflow\n";
            return 4;
        }
        snapshot_bytes += shard.size;
    }
    std::cerr << "Training export snapshot: " << shards.size()
              << " sealed shards, " << snapshot_bytes << " bytes";
    if (!options.included_runs.empty()) {
        std::cerr << ", " << options.included_runs.size() << " selected run(s)";
    }
    std::cerr << '\n';

    std::ios::sync_with_stdio(false);
    if (!write_header(std::cout)) {
        std::cerr << "Cannot write training export header to stdout\n";
        return 5;
    }

    std::uint64_t exported_records = 0;
    for (std::size_t shard_ordinal = 0; shard_ordinal < shards.size(); ++shard_ordinal) {
        const ShardSnapshot& snapshot = shards[shard_ordinal];
        bool output_failed = false;
        std::string record_error;
        training::StreamCallbacks callbacks;
        callbacks.injected_trajectory = [&](const training::InjectedTrajectory& trajectory) {
            PackedRecord packed;
            if (!encode_trajectory(
                    trajectory,
                    snapshot.expected_run_id,
                    packed,
                    record_error)) {
                return false;
            }
            std::cout.write(
                reinterpret_cast<const char*>(packed.data()),
                static_cast<std::streamsize>(packed.size()));
            if (!std::cout) {
                output_failed = true;
                return false;
            }
            ++exported_records;
            return true;
        };
        training::ShardInfo info;
        std::string shard_error;
        if (!training::stream_read_shard(
                snapshot.path,
                callbacks,
                &info,
                &shard_error)) {
            if (output_failed) {
                std::cerr << "stdout failed while exporting "
                          << snapshot.path.string() << '\n';
                return 5;
            }
            std::cerr << "Cannot export " << snapshot.path.string() << ": "
                      << (record_error.empty() ? shard_error : record_error) << '\n';
            return 6;
        }
        if (snapshot.filename_index && info.shard_index != *snapshot.filename_index) {
            std::cerr << "Shard header index does not match filename: "
                      << snapshot.path.string() << '\n';
            return 6;
        }
        if ((shard_ordinal + 1U) % 16U == 0U ||
            shard_ordinal + 1U == shards.size()) {
            std::cerr << "Validated " << (shard_ordinal + 1U) << '/'
                      << shards.size() << " shards; exported "
                      << exported_records << " trajectories\n";
        }
    }
    std::cout.flush();
    if (!std::cout) {
        std::cerr << "Cannot finish writing the training export stream\n";
        return 5;
    }
    std::cerr << "Training export complete: " << exported_records
              << " fixed records, record_size=" << kExportRecordSize << '\n';
    return 0;
}