#include "ling3/decoder_state.h"
#include "ling3/state_cache.h"
#include <filesystem>
#include <fstream>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <unistd.h>

void Check(bool ok, const char * why) { if (!ok) throw std::runtime_error(why); }
template<class F> void Reject(F f) {
    try { f(); } catch (const std::invalid_argument &) { return; }
    throw std::runtime_error("invalid state accepted");
}
int main() {
    char name[] = "/tmp/ling3-state-test-XXXXXX";
    if (!mkdtemp(name)) return 1;
    const auto directory = std::filesystem::path(name);
    try {
        ling3::DecoderState state; state.signature = "test-model/quant-v1"; state.position = 2;
        state.layers.resize(24);
        for (std::size_t i=0; i<24; ++i) {
            auto & l = state.layers[i];
            if ((i+1)%4 == 0) {
                l.keys.assign(2*16*192, 0x3f80); l.values.assign(2*16*128, 0xbf80);
            } else {
                l.gdn.fp16.assign(16*128*128, 0x3c00);
                l.gdn.fp32.assign(16*128*128, .375F); l.gdn.fp32_valid = true;
                for (auto & c : l.conv) c.assign(16*128*3, -.125F);
            }
        }
        const auto path = directory / "state.bin";
        ling3::WriteDecoderState(state, path);
        auto restored = ling3::ReadDecoderState(path, state.signature, 2);
        Check(restored.position == 2 && restored.bytes() == state.bytes(), "state shape changed");
        for (std::size_t i=0; i<24; ++i) {
            const auto & a = state.layers[i]; const auto & b = restored.layers[i];
            Check(a.keys==b.keys && a.values==b.values && a.conv==b.conv &&
                  a.gdn.fp16==b.gdn.fp16 && a.gdn.fp32==b.gdn.fp32 &&
                  a.gdn.fp32_valid==b.gdn.fp32_valid, "round trip changed state bits");
        }
        Reject([&] { ling3::ReadDecoderState(path, "different-quantization", 2); });
        Reject([&] { ling3::ReadDecoderState(path, state.signature, 1); });
        restored.layers[3].keys.pop_back();
        Reject([&] { ling3::ValidateDecoderState(restored, state.signature, 2); });
        state.layers[0].gdn.fp32[0] = std::numeric_limits<float>::quiet_NaN();
        Reject([&] { ling3::WriteDecoderState(state, directory / "bad.bin"); });
        Check(!std::filesystem::exists(directory / "bad.bin"), "invalid state published");
        state.layers[0].gdn.fp32[0] = .375F;
        // FP32 is authoritative: an overflowing FP16 shadow is permitted.
        state.layers[0].gdn.fp16[0] = 0x7c00;
        ling3::ValidateDecoderState(state, state.signature, 2);
        state.layers[0].gdn.fp32_valid = false;
        Reject([&] { ling3::ValidateDecoderState(state, state.signature, 2); });
        state.layers[0].gdn.fp32_valid = true;
        { std::fstream out(path, std::ios::binary|std::ios::in|std::ios::out); out.put('\0'); }
        Reject([&] { ling3::ReadDecoderState(path, state.signature, 2); });
        ling3::WriteDecoderState(state, path);
        { std::ofstream out(path, std::ios::binary|std::ios::app); out.put('x'); }
        Reject([&] { ling3::ReadDecoderState(path, state.signature, 2); });
        ling3::WriteDecoderState(state, path);
        { std::fstream out(path, std::ios::binary|std::ios::in|std::ios::out);
          out.seekp(-1, std::ios::end); out.put('\x12'); }
        Reject([&] { ling3::ReadDecoderState(path, state.signature, 2); });
        std::filesystem::resize_file(path, 12);
        Reject([&] { ling3::ReadDecoderState(path, state.signature, 2); });

        ling3::StateCache<int> cache(100);
        Check(cache.Put("a",1,40) && cache.Put("b",2,40), "cache insert");
        Check(*cache.Get("a")==1, "cache touch");
        cache.Put("c",3,40);
        Check(!cache.Get("b") && cache.Get("a") && cache.evictions()==1, "LRU eviction");
        Check(!cache.Put("a",4,101) && !cache.Get("a") && cache.bytes()==40, "oversize stale entry");
        ling3::StateCache<int> moved(std::move(cache));
        Check(*moved.Get("c")==3, "cache move iterators");
        moved.Erase("c"); Check(moved.bytes()==0 && moved.size()==0, "cache release");
        ling3::StateCache<int> disabled(0); Check(!disabled.Put("x",1,0), "disabled cache");
        std::filesystem::remove_all(directory);
        std::cout << "PASS complete state round-trip, corruption/version/capacity rejection, budget and LRU\n";
    } catch (const std::exception & e) {
        std::filesystem::remove_all(directory); std::cerr << e.what() << '\n'; return 1;
    }
}