#define WIN32_LEAN_AND_MEAN
#define NOMINMAX
#include <windows.h>
#include <intrin.h>
#include "llama.h"
#include <algorithm>
#include <array>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <iterator>
#include <memory>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

namespace {
constexpr uint32_t request_magic = 0x31514346; // FCQ1
constexpr uint32_t response_magic = 0x31524346; // FCR1
constexpr uint32_t max_question = 768;
constexpr uint32_t max_answer = 8192;
constexpr int context_size = 4096;
constexpr int max_tokens = 900;

bool read_exact(void *data, DWORD length) {
    auto *bytes = static_cast<char *>(data);
    while (length) {
        DWORD read = 0;
        if (!ReadFile(GetStdHandle(STD_INPUT_HANDLE), bytes, length, &read, nullptr) || !read)
            return false;
        bytes += read;
        length -= read;
    }
    return true;
}

bool write_exact(const void *data, DWORD length) {
    auto *bytes = static_cast<const char *>(data);
    while (length) {
        DWORD wrote = 0;
        if (!WriteFile(GetStdHandle(STD_OUTPUT_HANDLE), bytes, length, &wrote, nullptr) || !wrote)
            return false;
        bytes += wrote;
        length -= wrote;
    }
    return true;
}

std::string convert(const std::string &input, UINT from, UINT to) {
    int size = MultiByteToWideChar(from, 0, input.data(), static_cast<int>(input.size()), nullptr, 0);
    if (size <= 0) return {};
    std::wstring wide(size, L'\0');
    MultiByteToWideChar(from, 0, input.data(), static_cast<int>(input.size()), wide.data(), size);
    if (to == 1251) {
        for (wchar_t &c : wide) {
            if (c == 0x2014 || c == 0x2013 || c == 0x2011) c = L'-';
            if (c == 0x201c || c == 0x201d || c == 0x00ab || c == 0x00bb) c = L'"';
            if (c == 0x2018 || c == 0x2019) c = L'\'';
            if (c == 0x00a0) c = L' ';
        }
    }
    size = WideCharToMultiByte(to, 0, wide.data(), static_cast<int>(wide.size()), nullptr, 0, nullptr, nullptr);
    std::string result(size, '\0');
    WideCharToMultiByte(to, 0, wide.data(), static_cast<int>(wide.size()), result.data(), size, nullptr, nullptr);
    return result;
}

bool reply(uint32_t status, const std::string &utf8 = {}) {
    std::string text = convert(utf8, CP_UTF8, 1251);
    if (text.size() > max_answer) text.resize(max_answer);
    for (char &c : text) {
        if (static_cast<unsigned char>(c) < 32 && c != '\n') c = ' ';
    }
    const uint32_t header[] = {response_magic, status, static_cast<uint32_t>(text.size())};
    return write_exact(header, sizeof(header)) && write_exact(text.data(), static_cast<DWORD>(text.size()));
}

std::string load_text(const std::filesystem::path &path) {
    std::ifstream input(path, std::ios::binary);
    if (!input) throw std::runtime_error("Missing local prompt data.");
    std::string result{std::istreambuf_iterator<char>(input), {}};
    if (result.size() > 32768) throw std::runtime_error("Prompt data is too large.");
    return result;
}

using Model = std::unique_ptr<llama_model, decltype(&llama_model_free)>;
using Context = std::unique_ptr<llama_context, decltype(&llama_free)>;
using Batch = std::unique_ptr<llama_batch_ext, decltype(&llama_batch_ext_free)>;
using Sampler = std::unique_ptr<llama_sampler, decltype(&llama_sampler_free)>;

struct Message { std::string role, text; };

std::string grammar_literal(const std::string &text) {
    std::string result = "\"";
    for (char c : text) {
        if (c == '\\' || c == '"') result += '\\';
        result += c;
    }
    return result + "\"";
}

std::string card_name(const std::string &entry) {
    return entry.substr(0, entry.find(" — "));
}

std::string safe_message(std::string text) {
    // Untrusted user/model text cannot introduce ChatML role delimiters.
    for (const std::string marker : {"<|", "|>"}) {
        for (size_t at = 0; (at = text.find(marker, at)) != std::string::npos; at += 2)
            text.replace(at, marker.size(), "  ");
    }
    return text;
}

class Engine {
    Model model{nullptr, llama_model_free};
    Context ctx{nullptr, llama_free};
    Batch batch{nullptr, llama_batch_ext_free};
    std::string system;
    std::vector<std::string> cards;
    std::vector<llama_token> cached_tokens;

    void decode(const llama_token *tokens, int count, int pos) {
        llama_batch_ext_clear(batch.get());
        for (int i = 0; i < count; ++i) {
            int idx = llama_batch_ext_add_token(batch.get(), 0, tokens[i]);
            llama_pos location = pos + i;
            llama_batch_ext_set_pos(batch.get(), idx, &location);
        }
        llama_batch_ext_set_output_logits(batch.get(), count - 1, true);
        if (llama_process(ctx.get(), LLAMA_PROCESS_TYPE_DECODE, batch.get()))
            throw std::runtime_error("Local inference failed.");
    }

public:
    void reset() {
        llama_memory_clear(llama_get_memory(ctx.get()), true);
        cached_tokens.clear();
    }

    explicit Engine(const std::filesystem::path &root) {
        system = load_text(root / "system-prompt.txt");
        std::string catalogue = load_text(root / "cards.txt");
        size_t at = 0;
        while (at < catalogue.size()) {
            size_t end = catalogue.find('\n', at);
            std::string card = catalogue.substr(at, end == std::string::npos ? end : end - at);
            if (!card.empty() && card.back() == '\r') card.pop_back();
            if (!card.empty()) cards.push_back(card);
            if (end == std::string::npos) break;
            at = end + 1;
        }
        if (cards.size() != 22) throw std::runtime_error("Invalid tarot catalogue.");
        llama_backend_init();
        auto mp = llama_model_default_params();
        mp.n_gpu_layers = 0;
        const auto path = root / "models/fortune-qwen3-4b-q4_k_m.gguf";
        model.reset(llama_model_load_from_file(path.u8string().c_str(), mp));
        if (!model) throw std::runtime_error("Local model is missing or cannot be loaded.");
        auto cp = llama_context_default_params();
        cp.n_ctx = context_size;
        cp.n_batch = 256;
        cp.n_ubatch = 256;
        cp.n_threads = static_cast<int>(std::clamp(std::thread::hardware_concurrency() / 2, 1u, 8u));
        cp.n_threads_batch = cp.n_threads;
        ctx.reset(llama_init_from_model(model.get(), cp));
        if (!ctx) throw std::runtime_error("Not enough memory for the local model.");
        batch.reset(llama_batch_ext_init(ctx.get()));
        if (!batch) throw std::runtime_error("Cannot allocate the local inference batch.");
    }

    std::pair<std::string, bool> answer(const std::string &question, const std::array<uint32_t, 3> &draw) {
        const std::array<std::string, 4> roles = {"Ситуация", "Препятствие", "Направление", "Общий совет"};
        std::array<std::string, 4> headings;
        std::string user = "Вопрос: " + safe_message(question) + "\nТочный расклад игры:\n";
        for (size_t i = 0; i < draw.size(); ++i) {
            const auto &entry = cards[draw[i] % 22];
            headings[i] = roles[i] + " - " + card_name(entry) +
                (draw[i] >= 22 ? " (перевёрнутая): " : " (прямая): ");
            user += headings[i] + entry.substr(entry.find(" — ") + std::string(" — ").size()) + "\n";
        }
        headings[3] = roles[3] + ": ";
        user += "Ответь четырьмя строками. В каждой 1-2 коротких законченных предложения, "
                "80-220 символов после заголовка. Свяжи символ каждой карты с вопросом, "
                "не перечисляй только общие значения. Для перевёрнутой карты объясни "
                "затруднение. В последней строке свяжи все три карты в один практичный, "
                "осторожный совет. Не выдумывай факты о человеке. Если вопрос просит "
                "определить реальную оценку по личности или картам, объясни, что карты "
                "не определяют оценку, и предложи критерии выполненной работы; не "
                "подменяй вопрос советом об оформлении документов. Начала строк строго такие:\n";
        for (const auto &heading : headings) user += heading + "\n";
        std::vector<Message> messages{{"system", system}};
        messages.push_back({"user", user});
        std::vector<llama_chat_message> chat;
        for (const auto &message : messages) chat.push_back({message.role.c_str(), message.text.c_str()});
        const char *tmpl = llama_model_chat_template(model.get(), nullptr);
        if (!tmpl) throw std::runtime_error("The model has no chat template.");
        std::vector<char> formatted(32768);
        int length = llama_chat_apply_template(tmpl, chat.data(), chat.size(), true, formatted.data(), static_cast<int>(formatted.size()));
        if (length < 0 || length >= static_cast<int>(formatted.size())) throw std::runtime_error("Conversation is too long.");
        const auto *vocab = llama_model_get_vocab(model.get());
        int count = -llama_tokenize(vocab, formatted.data(), length, nullptr, 0, true, true);
        if (count <= 0 || count > context_size - max_tokens) throw std::runtime_error("Question is too long.");
        std::vector<llama_token> tokens(count);
        if (llama_tokenize(vocab, formatted.data(), length, tokens.data(), count, true, true) < 0)
            throw std::runtime_error("Cannot tokenize the question.");
        int prefix = 0;
        while (prefix < count && prefix < static_cast<int>(cached_tokens.size()) &&
               tokens[prefix] == cached_tokens[prefix]) ++prefix;
        if (prefix == count) --prefix;
        // Reuse only an exact token prefix; never reuse unrelated answer state.
        if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, prefix, -1)) {
            llama_memory_clear(llama_get_memory(ctx.get()), true);
            prefix = 0;
        }
        for (int pos = prefix; pos < count; pos += 256) decode(tokens.data() + pos, std::min(256, count - pos), pos);
        cached_tokens = tokens;
        Sampler sampler(llama_sampler_chain_init(llama_sampler_chain_default_params()), llama_sampler_free);
        // Exact headings bind names, positions and orientations to the actual draw.
        // Four bounded Russian lines map directly to the four native scene stages.
        std::string grammar = "root ::= ";
        for (size_t i = 0; i < headings.size(); ++i) {
            if (i) grammar += " \"\\n\" ";
            grammar += grammar_literal(headings[i]) + " body";
        }
        grammar += "\nbody ::= [А-ЯЁ] [а-яёА-ЯЁ0-9 ,;:()\\-]{78,218} [.!?]\n";
        auto *constraint = llama_sampler_init_grammar(vocab, grammar.c_str(), "root");
        if (!constraint) throw std::runtime_error("Invalid reading format grammar.");
        llama_sampler_chain_add(sampler.get(), constraint);
        llama_sampler_chain_add(sampler.get(), llama_sampler_init_top_k(40));
        llama_sampler_chain_add(sampler.get(), llama_sampler_init_top_p(0.8f, 1));
        llama_sampler_chain_add(sampler.get(), llama_sampler_init_temp(0.4f));
        llama_sampler_chain_add(sampler.get(), llama_sampler_init_dist(42));
        std::string output;
        bool complete = false;
        for (int i = 0; i < max_tokens; ++i) {
            llama_token token = llama_sampler_sample(sampler.get(), ctx.get(), -1);
            if (llama_vocab_is_eog(vocab, token)) { complete = true; break; }
            char piece[256];
            int bytes = llama_token_to_piece(vocab, token, piece, sizeof(piece), 0, false);
            if (bytes < 0) throw std::runtime_error("Cannot decode the answer.");
            output.append(piece, bytes);
            decode(&token, 1, count + i);
            cached_tokens.push_back(token);
        }
        if (output.empty() || !complete) throw std::runtime_error("The model did not finish the reading.");
        size_t at = 0;
        for (size_t i = 0; i < headings.size(); ++i) {
            if (output.compare(at, headings[i].size(), headings[i]) != 0)
                throw std::runtime_error("The reading changed the supplied cards.");
            const size_t end = output.find('\n', at);
            if ((i < 3) != (end != std::string::npos))
                throw std::runtime_error("The reading has an invalid number of sections.");
            at = end == std::string::npos ? output.size() : end + 1;
        }
        return {output, complete};
    }
};
}

int wmain() {
    SetErrorMode(SEM_FAILCRITICALERRORS | SEM_NOGPFAULTERRORBOX | SEM_NOOPENFILEERRORBOX);
    int cpu[4]{};
    __cpuid(cpu, 1);
    constexpr int avx_osxsave = (1 << 27) | (1 << 28);
    if ((cpu[2] & avx_osxsave) != avx_osxsave || (_xgetbv(0) & 6) != 6) {
        reply(2, "This local model requires an AVX-capable CPU and OS.");
        return 1;
    }
    try {
        std::array<wchar_t, 32768> executable{};
        if (!GetModuleFileNameW(nullptr, executable.data(), static_cast<DWORD>(executable.size()))) return 1;
        Engine engine(std::filesystem::path(executable.data()).parent_path());
        if (!reply(0)) return 0; // Ready: no model or prompt content on stdout.
        for (;;) {
            std::array<uint32_t, 5> header{};
            if (!read_exact(header.data(), sizeof(header))) return 0;
            if (header[0] != request_magic || !header[1] || header[1] > max_question) return 2;
            std::array<uint32_t, 3> cards{header[2], header[3], header[4]};
            if (std::any_of(cards.begin(), cards.end(), [](uint32_t v) { return v >= 44; })) return 2;
            if (cards[0] % 22 == cards[1] % 22 || cards[0] % 22 == cards[2] % 22 || cards[1] % 22 == cards[2] % 22) return 2;
            std::string input(header[1], '\0');
            if (!read_exact(input.data(), static_cast<DWORD>(input.size()))) return 0;
            if (input.find('\0') != std::string::npos) return 2;
            try {
                const auto result = engine.answer(convert(input, 1251, CP_UTF8), cards);
                if (!reply(result.second ? 1 : 3, result.first)) return 0;
            } catch (const std::exception &error) {
                engine.reset();
                if (!reply(2, error.what())) return 0;
            }
        }
    } catch (const std::exception &error) {
        reply(2, error.what());
        return 1;
    }
}