#include "ling3/flow_gate.h"
#include "ling3/chat_sampling.h"
#include <future>
#include <iostream>
#include <thread>

using namespace ling3::chat;
using namespace std::chrono_literals;
void Check(bool value, const char * message) { if (!value) throw std::runtime_error(message); }
template<class F> void Rejected(F f) {
    bool rejected = false;
    try { f(); } catch (const Error &) { rejected = true; }
    Check(rejected, "stale or invalid command accepted");
}
int main() {
    FlowGate gate;
    gate.Begin("a", true);
    Check(gate.Prepare() == 1, "sequence start");
    auto waiting = std::async(std::launch::async, [&] { return gate.Wait([] { return true; }); });
    Check(waiting.wait_for(30ms) == std::future_status::timeout, "generation advanced without ack");
    Rejected([&] { gate.Ack("b", 1); });
    Rejected([&] { gate.Ack("a", 2); });
    gate.Ack("a", 1);
    Check(waiting.get(), "ack did not resume");
    gate.Ack("a", 1); // Idempotent retry cannot grant a future credit.
    gate.Prepare();
    Rejected([&] { gate.Ack("a", 1); });
    auto canceled = std::async(std::launch::async, [&] { return gate.Wait([] { return true; }); });
    Check(!gate.Cancel("b"), "canceled another request");
    Check(canceled.wait_for(20ms) == std::future_status::timeout, "stale cancel unblocked generation");
    Check(gate.Cancel("a") && !canceled.get(), "cancel while parked failed");
    gate.End(); gate.Begin("b", true); gate.Prepare();
    Rejected([&] { gate.Ack("a", 1); });
    Check(!gate.Wait([] { return false; }), "disconnected parked client not canceled");
    gate.End(); gate.Begin("parallel", false);
    auto pause = std::async(std::launch::async, [&] { gate.Pause("parallel"); });
    while (!gate.Status()["pause_requested"].get<bool>()) std::this_thread::yield();
    auto paused = std::async(std::launch::async, [&] { return gate.Wait([] { return true; }); });
    pause.get();
    Check(gate.Status()["paused"], "pause returned before safe boundary");
    Check(paused.wait_for(30ms) == std::future_status::timeout, "pause did not park generation");
    Rejected([&] { gate.Resume("a"); });
    gate.Resume("parallel"); Check(paused.get(), "resume failed"); gate.End();

    Json body = {{"model", kModel}, {"messages", Json::array({{{"role", "user"}, {"content", "hi"}}})}};
    for (auto bad : std::vector<Json>{{{"top_k", -1}}, {{"top_k", 1.5}}, {{"repeat_penalty", 0}},
                                      {{"repeat_penalty", "1"}}, {{"enable_thinking", 1}}, {{"flow_control", "ack"}},
                                      {{"reuse_generated_state", "true"}}}) {
        auto request = body; request.update(bad); Rejected([&] { Parse(request); });
    }
    body.update({{"top_k", 2}, {"repeat_penalty", 2}, {"temperature", 0.7}, {"enable_thinking", true},
                 {"flow_control", "ack"}, {"stream", true}});
    auto request = Parse(body);
    Check(request.reuse_generated_state, "QA cache must be enabled by default");
    auto incremental = body; incremental["reuse_generated_state"] = true;
    Check(Parse(incremental).reuse_generated_state, "generated reuse flag missing");
    gate.Begin("commit", false);
    Check(gate.TryCommit(), "completion commit failed");
    Check(!gate.Cancel("commit"), "late cancellation accepted after commit");
    Check(!gate.TryCommit(), "duplicate commit accepted");
    Rejected([&] { gate.Pause("commit"); });
    gate.End();
    gate.Begin("cancel-first", false);
    Check(gate.Cancel("cancel-first") && !gate.TryCommit(), "canceled generation committed");
    gate.End();
    Check(request.prompt.ends_with("\n<think>"), "thinking template not enabled");
    Check(request.prompt.find("detailed thinking on") != std::string::npos, "thinking instruction missing");
    Check(request.flow_ack, "flow mode missing");
    auto conflicting = body;
    conflicting["enable_thinking"] = false;
    conflicting["session_id"] = "conversation-a";
    conflicting["messages"].insert(conflicting["messages"].begin(),
        Json{{"role", "system"}, {"content", "detailed thinking on"}});
    const auto normalized = Parse(conflicting);
    Check(normalized.session_id == "conversation-a" && normalized.prompt.find("detailed thinking on") == std::string::npos
        && normalized.prompt.ends_with("<think></think>"), "thinking switch did not override system");
    conflicting["enable_thinking"] = true;
    conflicting["messages"][0]["content"] = "detailed thinking off; detailed thinking on";
    const auto enabled = Parse(conflicting);
    Check(enabled.prompt.find("detailed thinking off") == std::string::npos && enabled.prompt.ends_with("<think>"),
          "thinking enable did not override mixed system instructions");
    conflicting["session_id"] = ""; Rejected([&] { Parse(conflicting); });
    std::mt19937 random(42);
    request.temperature = 0;
    Check(SampleToken(std::vector<float>{6, 4, -1}, request, {true, false, false}, random) == 1,
          "positive repetition penalty not applied before greedy selection");
    Check(SampleToken(std::vector<float>{-1, -1.5F, -3}, request, {true, false, false}, random) == 1,
          "negative repetition penalty has wrong sign");
    request.temperature = 1; request.repeat_penalty = 1; request.top_k = 2;
    for (int i=0; i<500; ++i)
        Check(SampleToken(std::vector<float>{3, 2, 1, 0}, request, {}, random) < 2, "top_k tail sampled");
    request.top_k = 1;
    Check(SampleToken(std::vector<float>{1, 2, 0}, request, {}, random) == 1, "top_k=1 not deterministic");
    request.top_k = 0; request.top_p = 0.01;
    Check(SampleToken(std::vector<float>{1, 2, 0}, request, {}, random) == 1, "top_p prefix ignored");
    std::cout << "flow ack, stale controls, parked cancellation, pause/resume, sampling and thinking passed\n";
}