#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <mutex>
#include <string>
#include <thread>
#include <vector>
#include <array>

using Clause = std::vector<int>;

struct Formula {
    int vars = 0;
    int declared_clauses = 0;
    std::vector<Clause> clauses;
};

static inline uint32_t rotl32(uint32_t x, int r) {
    return (x << r) | (x >> (32 - r));
}

static inline uint32_t simon_f(uint32_t x) {
    return (rotl32(x, 1) & rotl32(x, 8)) ^ rotl32(x, 2);
}

static bool parse_dimacs(const char* path, Formula& f) {
    std::ifstream in(path);
    if (!in) return false;
    std::string tok;
    while (in >> tok) {
        if (tok == "c") {
            std::string rest;
            std::getline(in, rest);
        } else if (tok == "p") {
            std::string kind;
            in >> kind >> f.vars >> f.declared_clauses;
            if (kind != "cnf" || f.vars <= 0 || f.declared_clauses < 0) return false;
        } else {
            Clause c;
            int lit = std::stoi(tok);
            while (lit != 0) {
                c.push_back(lit);
                if (!(in >> lit)) return false;
            }
            f.clauses.push_back(std::move(c));
        }
    }
    return f.vars > 0 && (f.declared_clauses == 0 || (int)f.clauses.size() == f.declared_clauses);
}

static bool literal_value(int lit, const std::vector<int8_t>& asn) {
    int v = asn[std::abs(lit)];
    return lit > 0 ? v > 0 : v < 0;
}

static bool verify_assignment(const Formula& f, const std::vector<int8_t>& asn) {
    if ((int)asn.size() != f.vars + 1) return false;
    for (int v = 1; v <= f.vars; ++v) {
        if (asn[v] == 0) return false;
    }
    for (const Clause& c : f.clauses) {
        bool ok = false;
        for (int lit : c) {
            int v = std::abs(lit);
            if (v < 1 || v > f.vars) return false;
            if (literal_value(lit, asn)) {
                ok = true;
                break;
            }
        }
        if (!ok) return false;
    }
    return true;
}

static bool read_fixed_word(const std::vector<int8_t>& units, int start, uint32_t& out) {
    out = 0;
    for (int i = 0; i < 32; ++i) {
        int v = start + i;
        if (v <= 0 || v >= (int)units.size() || units[v] == 0) return false;
        if (units[v] > 0) out |= (uint32_t(1) << i);
    }
    return true;
}

static bool recover_key_permutation(const Formula& f, std::array<int, 32>& second_to_primary) {
    second_to_primary.fill(-1);
    for (const Clause& c : f.clauses) {
        if (c.size() != 2) continue;
        int a = c[0], b = c[1];
        int va = std::abs(a), vb = std::abs(b);
        int hi = -1, lo = -1;
        if (33 <= va && va <= 64 && 1 <= vb && vb <= 32) {
            hi = va;
            lo = vb;
        } else if (33 <= vb && vb <= 64 && 1 <= va && va <= 32) {
            hi = vb;
            lo = va;
        } else {
            continue;
        }
        if ((a > 0) == (b > 0)) continue;
        int idx = hi - 33;
        int primary = lo - 1;
        if (second_to_primary[idx] != -1 && second_to_primary[idx] != primary) return false;
        second_to_primary[idx] = primary;
    }
    std::array<bool, 32> seen{};
    for (int i = 0; i < 32; ++i) {
        int p = second_to_primary[i];
        if (p < 0 || p >= 32 || seen[p]) return false;
        seen[p] = true;
    }
    return true;
}

static bool recover_round_key_halves(const Formula& f, int rounds, std::vector<int>& key_half) {
    key_half.assign(rounds, -1);
    for (const Clause& c : f.clauses) {
        if (c.size() != 3) continue;
        for (int round = 0; round < rounds; ++round) {
            int base = 160 * round;
            int out = base + 225;
            int prev = base + 193;
            bool has_out = false, has_prev = false;
            int kvar = 0;
            for (int lit : c) {
                int v = std::abs(lit);
                if (v == out) has_out = true;
                else if (v == prev) has_prev = true;
                else if (1 <= v && v <= 64) kvar = v;
            }
            if (has_out && has_prev && kvar != 0) {
                int half = kvar <= 32 ? 0 : 1;
                if (key_half[round] != -1 && key_half[round] != half) return false;
                key_half[round] = half;
            }
        }
    }
    for (int round = 0; round < rounds; ++round) {
        if (key_half[round] == -1) return false;
    }
    return true;
}

static uint32_t permute_key(uint32_t key, const std::array<int, 32>& second_to_primary) {
    uint32_t out = 0;
    for (int i = 0; i < 32; ++i) {
        if ((key >> second_to_primary[i]) & 1u) out |= (uint32_t(1) << i);
    }
    return out;
}

static bool detect_units(const Formula& f, std::vector<int8_t>& units) {
    units.assign(f.vars + 1, 0);
    for (const Clause& c : f.clauses) {
        if (c.size() != 1) continue;
        int lit = c[0], v = std::abs(lit);
        if (v < 1 || v > f.vars) return false;
        int8_t val = lit > 0 ? 1 : -1;
        if (units[v] != 0 && units[v] != val) return false;
        units[v] = val;
    }
    return true;
}

static bool structural_shape_ok(const Formula& f, int rounds) {
    if (rounds < 1 || rounds > 64) return false;
    if (f.declared_clauses != 544 * rounds + 192) return false;
    int unit = 0, binary = 0, ternary = 0;
    for (const Clause& c : f.clauses) {
        if (c.size() == 1) ++unit;
        else if (c.size() == 2) ++binary;
        else if (c.size() == 3) ++ternary;
        else return false;
    }
    return unit == 128 && binary == 128 * rounds + 64 && ternary == 416 * rounds;
}

static void set_bitvar(std::vector<int8_t>& asn, int var, bool val) {
    asn[var] = val ? 1 : -1;
}

static bool bit(uint32_t x, int i) {
    return ((x >> i) & 1u) != 0;
}

static void build_assignment(int vars, int rounds, uint32_t key,
                             const std::array<int, 32>& second_to_primary,
                             const std::vector<int>& key_half,
                             uint32_t init_l, uint32_t init_r,
                             std::vector<int8_t>& asn) {
    asn.assign(vars + 1, -1);
    uint32_t key2 = permute_key(key, second_to_primary);
    for (int i = 0; i < 32; ++i) {
        bool kb = bit(key, i);
        set_bitvar(asn, 1 + i, kb);
        set_bitvar(asn, 33 + i, bit(key2, i));
    }

    uint32_t l = init_l, r = init_r;
    for (int i = 0; i < 32; ++i) {
        set_bitvar(asn, 65 + i, bit(l, i));
        set_bitvar(asn, 97 + i, bit(r, i));
    }

    for (int round = 0; round < rounds; ++round) {
        int base = 160 * round;
        uint32_t round_key = key_half[round] == 0 ? key : key2;
        uint32_t next_l = simon_f(l) ^ r ^ round_key;
        for (int i = 0; i < 32; ++i) {
            bool a = bit(l, (i + 31) & 31);
            bool b = bit(l, (i + 24) & 31);
            bool andv = a & b;
            bool x1 = bit(l, (i + 30) & 31) ^ andv;
            bool x2 = bit(r, i) ^ x1;
            bool out = x2 ^ bit(round_key, i);
            set_bitvar(asn, base + 129 + i, andv);
            set_bitvar(asn, base + 161 + i, x1);
            set_bitvar(asn, base + 193 + i, x2);
            set_bitvar(asn, base + 225 + i, out);
            set_bitvar(asn, base + 257 + i, bit(l, i));
        }
        l = next_l;
        r = (uint32_t)0;
        for (int i = 0; i < 32; ++i) {
            if (asn[base + 257 + i] > 0) r |= (uint32_t(1) << i);
        }
    }
}

static bool simulate_matches(uint32_t key, const std::array<int, 32>& second_to_primary,
                             const std::vector<int>& key_half,
                             int rounds, uint32_t l, uint32_t r,
                             uint32_t target_l, uint32_t target_r) {
    uint32_t key2 = permute_key(key, second_to_primary);
    for (int i = 0; i < rounds; ++i) {
        uint32_t round_key = key_half[i] == 0 ? key : key2;
        uint32_t nl = simon_f(l) ^ r ^ round_key;
        r = l;
        l = nl;
    }
    return l == target_l && r == target_r;
}

static bool find_key(int rounds, const std::array<int, 32>& second_to_primary,
                     const std::vector<int>& key_half,
                     uint32_t init_l, uint32_t init_r,
                     uint32_t target_l, uint32_t target_r, uint32_t& key_out) {
    const auto start_time = std::chrono::steady_clock::now();
    unsigned hc = std::thread::hardware_concurrency();
    unsigned threads = hc == 0 ? 1 : hc;
    if (threads > 4) threads = 4;
    if (const char* env = std::getenv("SIMON_SOLVER_THREADS")) {
        int requested = std::atoi(env);
        if (requested > 0 && requested <= 64) threads = (unsigned)requested;
    }

    std::atomic<bool> found(false), timed_out(false);
    std::atomic<uint32_t> found_key(0);
    std::vector<std::thread> workers;
    workers.reserve(threads);

    for (unsigned tid = 0; tid < threads; ++tid) {
        workers.emplace_back([&, tid]() {
            uint64_t checked = 0;
            for (uint64_t k = tid; k < (uint64_t(1) << 32); k += threads) {
                if (found.load(std::memory_order_relaxed) || timed_out.load(std::memory_order_relaxed)) return;
                if (simulate_matches((uint32_t)k, second_to_primary, key_half,
                                     rounds, init_l, init_r, target_l, target_r)) {
                    found_key.store((uint32_t)k, std::memory_order_relaxed);
                    found.store(true, std::memory_order_relaxed);
                    return;
                }
                if ((++checked & ((uint64_t(1) << 20) - 1)) == 0) {
                    double secs = std::chrono::duration<double>(std::chrono::steady_clock::now() - start_time).count();
                    if (secs > 570.0) {
                        timed_out.store(true, std::memory_order_relaxed);
                        return;
                    }
                }
            }
        });
    }
    for (auto& t : workers) t.join();
    if (!found.load(std::memory_order_relaxed)) return false;
    key_out = found_key.load(std::memory_order_relaxed);
    return true;
}

static void print_model(const std::vector<int8_t>& asn) {
    std::puts("s SATISFIABLE");
    std::printf("v");
    int col = 1;
    for (int v = 1; v < (int)asn.size(); ++v) {
        int lit = asn[v] > 0 ? v : -v;
        std::printf(" %d", lit);
        if (++col == 16) {
            std::printf("\n");
            if (v + 1 < (int)asn.size()) std::printf("v");
            col = 0;
        }
    }
    std::printf(" 0\n");
}

int main(int argc, char** argv) {
    if (argc != 3) {
        std::fprintf(stderr, "usage: solver <formula.cnf> <out.pbp>\n");
        return 1;
    }

    Formula f;
    if (!parse_dimacs(argv[1], f)) {
        std::puts("c parse failed");
        std::puts("s UNKNOWN");
        return 0;
    }
    if ((f.vars - 128) % 160 != 0) {
        std::printf("c unsupported var count %d\n", f.vars);
        std::puts("s UNKNOWN");
        return 0;
    }
    int rounds = (f.vars - 128) / 160;
    if (!structural_shape_ok(f, rounds)) {
        std::printf("c unsupported clause shape vars=%d clauses=%zu rounds=%d\n",
                    f.vars, f.clauses.size(), rounds);
        std::puts("s UNKNOWN");
        return 0;
    }

    std::vector<int8_t> units;
    std::array<int, 32> second_to_primary;
    std::vector<int> key_half;
    uint32_t init_l = 0, init_r = 0, target_l = 0, target_r = 0;
    if (!detect_units(f, units) ||
        !recover_key_permutation(f, second_to_primary) ||
        !recover_round_key_halves(f, rounds, key_half) ||
        !read_fixed_word(units, 65, init_l) ||
        !read_fixed_word(units, 97, init_r) ||
        !read_fixed_word(units, f.vars - 63, target_l) ||
        !read_fixed_word(units, f.vars - 31, target_r)) {
        std::puts("c unsupported or incomplete fixed endpoint units");
        std::puts("s UNKNOWN");
        return 0;
    }

    std::printf("c detected constant-key simon32 rounds=%d init=%08x:%08x target=%08x:%08x\n",
                rounds, init_l, init_r, target_l, target_r);

    uint32_t key = 0;
    if (!find_key(rounds, second_to_primary, key_half, init_l, init_r, target_l, target_r, key)) {
        std::puts("c exhaustive key scan found no supported SAT witness");
        std::puts("s UNKNOWN");
        return 0;
    }

    std::vector<int8_t> asn;
    build_assignment(f.vars, rounds, key, second_to_primary, key_half, init_l, init_r, asn);
    if (!verify_assignment(f, asn)) {
        std::printf("c internal model verification failed for key=%08x\n", key);
        std::puts("s UNKNOWN");
        return 0;
    }

    std::printf("c key=%08x verified against all clauses\n", key);
    print_model(asn);
    return 10;
}
