// Executable semantic laboratory for the book's mutable capstone.
// SPDX-License-Identifier: MIT
// This is NOT the proposed slab/SIMD production implementation.
#pragma once
#include "flex_wire.hpp"
#include <atomic>
#include <memory>
#include <unordered_map>
#include <unordered_set>
#include <variant>
#include <type_traits>

namespace flexlab::live {
enum class Kind : std::uint8_t { Null, Bool, Int, UInt, Float, String, Key, Blob, Vector, Map, F64 };
struct alignas(8) Cell {
 std::uint64_t payload=0;
 std::uint32_t aux=0;
 Kind kind=Kind::Null;
 std::uint8_t flags=0;
 std::uint16_t reserved=0;
};
static_assert(sizeof(Cell)==16 && alignof(Cell)==8);
static_assert(offsetof(Cell,aux)==8 && offsetof(Cell,kind)==12);
static_assert(std::is_trivially_copyable_v<Cell>);
inline bool small_text(const Cell& c) { return (c.kind==Kind::String || c.kind==Kind::Key) && (c.flags&0x80); }
inline bool external(const Cell& c) {
 return c.kind>=Kind::String && !small_text(c);
}
inline Cell null() { return {}; }
inline Cell boolean(bool x) { return {x ? 1ULL : 0ULL,0,Kind::Bool,0,0}; }
inline Cell integer(std::int64_t x) { return {std::bit_cast<std::uint64_t>(x),0,Kind::Int,0,0}; }
inline Cell uint(std::uint64_t x) { return {x,0,Kind::UInt,0,0}; }
inline Cell real(double x) { return {std::bit_cast<std::uint64_t>(x),0,Kind::Float,0,0}; }
inline std::int64_t as_int(Cell c) { require(c.kind==Kind::Int,"not an integer cell"); return std::bit_cast<std::int64_t>(c.payload); }
inline double as_float(Cell c) { require(c.kind==Kind::Float,"not a float cell"); return std::bit_cast<double>(c.payload); }
struct TransparentHash {
 using is_transparent=void;
 std::size_t operator()(std::string_view x) const { return std::hash<std::string_view>{}(x); }
};
class Document {
 struct Text { std::string data; };
 struct Blob { Bytes data; };
 struct Vector { std::vector<Cell> data; };
 struct Map { std::vector<std::pair<std::uint32_t,Cell>> data; };
 struct F64 { std::vector<double> data; };
 using Node=std::variant<Text,Blob,Vector,Map,F64>;
 struct Slot {
   std::uint32_t generation=1;
   std::unique_ptr<Node> node;
 };
 inline static std::atomic<std::uint64_t> next_owner_{1};
 const std::uint32_t owner_;
 std::vector<Slot> slots_;
 std::vector<std::uint32_t> free_;
 std::unordered_map<std::string,std::uint32_t,TransparentHash,std::equal_to<>> symbols_;
 std::vector<std::string> names_;
 static std::uint32_t new_owner() {
   const auto x=next_owner_.fetch_add(1,std::memory_order_relaxed);
   require(x<=UINT32_MAX,"document identity space exhausted"); return static_cast<std::uint32_t>(x);
 }
 void own(Cell c) const {
   if (!external(c)) return;
   require(c.aux==owner_,"cross-document handle");
   const auto i=static_cast<std::uint32_t>(c.payload), g=static_cast<std::uint32_t>(c.payload>>32);
   require(i<slots_.size() && slots_[i].node && slots_[i].generation==g,"stale handle");
 }
 Node& node(Cell c) { own(c); require(external(c),"not an external cell"); return *slots_[static_cast<std::uint32_t>(c.payload)].node; }
 const Node& node(Cell c) const { own(c); require(external(c),"not an external cell"); return *slots_[static_cast<std::uint32_t>(c.payload)].node; }
 template<class T> T& get(Cell c, Kind kind) {
   require(c.kind==kind,"cell kind mismatch"); return std::get<T>(node(c));
 }
 template<class T> const T& get(Cell c, Kind kind) const {
   require(c.kind==kind,"cell kind mismatch"); return std::get<T>(node(c));
 }
 Cell allocate(Kind kind, Node n) {
   auto ptr=std::make_unique<Node>(std::move(n)); // allocate before publishing
   std::uint32_t i;
   if (free_.empty()) {
     require(slots_.size()<UINT32_MAX,"slot space exhausted");
     i=static_cast<std::uint32_t>(slots_.size()); slots_.push_back({});
   } else { i=free_.back(); free_.pop_back(); }
   slots_[i].node=std::move(ptr);
   return { (std::uint64_t{slots_[i].generation}<<32)|i,owner_,kind,0,0 };
 }
 std::uint32_t intern(std::string_view key) {
   require(key.find('\0')==key.npos && utf8(key),"invalid map key");
   if (const auto it=symbols_.find(key); it!=symbols_.end()) return it->second;
   require(names_.size()<UINT32_MAX,"key identity space exhausted");
   const auto id=static_cast<std::uint32_t>(names_.size());
   names_.emplace_back(key);
   try { symbols_.emplace(names_.back(),id); }
   catch (...) { names_.pop_back(); throw; }
   return id;
 }
 template<class F> void children(Cell c, F&& f) const {
   if (c.kind==Kind::Vector) for (auto v:get<Vector>(c,Kind::Vector).data) f(v);
   if (c.kind==Kind::Map) for (const auto& kv:get<Map>(c,Kind::Map).data) f(kv.second);
 }
 Encoded freeze_one(Writer& w, Cell c, std::unordered_map<std::uint64_t,Encoded>& done,
                    std::unordered_set<std::uint64_t>& active, std::size_t depth) const {
   require(depth<=64,"freeze nesting limit exceeded"); own(c);
   if (external(c)) {
     if (const auto it=done.find(c.payload); it!=done.end()) return it->second;
     require(active.insert(c.payload).second,"cycle cannot be exported to FlexBuffers");
   }
   Encoded result;
   switch (c.kind) {
     case Kind::Null: result=Writer::null(); break;
     case Kind::Bool: result=Writer::boolean(c.payload!=0); break;
     case Kind::Int: result=Writer::integer(as_int(c)); break;
     case Kind::UInt: result=Writer::uint(c.payload); break;
     case Kind::Float: result=Writer::real(as_float(c)); break;
     case Kind::String: result=w.string(text(c)); break;
     case Kind::Key: result=w.key(text(c)); break;
     case Kind::Blob: result=w.blob(get<Blob>(c,Kind::Blob).data); break;
     case Kind::Vector: {
       std::vector<Encoded> v;
       for (auto item:get<Vector>(c,Kind::Vector).data) v.push_back(freeze_one(w,item,done,active,depth+1));
       result=w.vector(v); break;
     }
     case Kind::Map: {
       std::vector<std::pair<std::string,Encoded>> fields;
       for (const auto& [k,value]:get<Map>(c,Kind::Map).data)
         fields.emplace_back(names_[k],freeze_one(w,value,done,active,depth+1));
       result=w.map(std::move(fields)); break;
     }
     case Kind::F64: {
       std::vector<Encoded> v;
       for (auto x:get<F64>(c,Kind::F64).data) v.push_back(Writer::real(x));
       result=w.typed_vector(v,Type::Float); break;
     }
   }
   if (external(c)) { active.erase(c.payload); done.emplace(c.payload,result); }
   return result;
 }
 Cell thaw_one(const Ref& r, std::size_t depth) {
   require(depth<=64,"thaw nesting limit exceeded");
   switch (r.type) {
     case Type::Null: return null();
     case Type::Bool: return boolean(r.as_bool());
     case Type::Int: case Type::IndirectInt: return integer(r.as_int());
     case Type::UInt: case Type::IndirectUInt: return uint(r.as_uint());
     case Type::Float: case Type::IndirectFloat: return real(r.as_float());
     case Type::String: return string(r.text());
     case Type::Key: return key(r.text());
     case Type::Blob: return blob(r.blob());
     case Type::Map: {
       auto out=map(); const auto k=r.keys();
       for (std::size_t i=0; i<r.length(); ++i) set(out,k.at(i).text(),thaw_one(r.at(i),depth+1));
       return out;
     }
     default: {
       require(r.type==Type::Vector || typed(r.type) || fixed(r.type),"unsupported import");
       if (r.type==Type::VectorFloat) {
         std::vector<double> values; values.reserve(r.length());
         for (std::size_t i=0; i<r.length(); ++i) values.push_back(r.at(i).as_float());
         return f64(std::move(values));
       }
       auto out=vector();
       for (std::size_t i=0; i<r.length(); ++i) push(out,thaw_one(r.at(i),depth+1));
       return out;
     }
   }
 }
public:
 Document() : owner_(new_owner()) {}
 Document(const Document&)=delete;
 Document& operator=(const Document&)=delete;
 Cell string(std::string_view s, bool is_key=false) {
   require(utf8(s) && (!is_key || s.find('\0')==s.npos),"invalid text");
   const auto kind=is_key ? Kind::Key : Kind::String;
   if (s.size()<=12) {
     Cell c; c.kind=kind; c.flags=static_cast<std::uint8_t>(0x80|s.size());
     if (!s.empty()) std::memcpy(&c,s.data(),s.size());
     return c;
   }
   return allocate(kind,Text{std::string(s)});
 }
 Cell key(std::string_view s) { return string(s,true); }
 // Returns an owning string, so even an inline temporary is safe.
 std::string text(Cell c) const {
   require(c.kind==Kind::String || c.kind==Kind::Key,"not a text cell");
   if (small_text(c)) return {reinterpret_cast<const char*>(&c),static_cast<std::size_t>(c.flags&15)};
   return std::get<Text>(node(c)).data;
 }
 Cell blob(ByteSpan b) { return allocate(Kind::Blob,Blob{Bytes(b.begin(),b.end())}); }
 Cell vector() { return allocate(Kind::Vector,Vector{}); }
 Cell map() { return allocate(Kind::Map,Map{}); }
 Cell f64(std::vector<double> data={}) { return allocate(Kind::F64,F64{std::move(data)}); }
 void push(Cell c, Cell value) { own(value); get<Vector>(c,Kind::Vector).data.push_back(value); }
 Cell at(Cell c, std::size_t i) const { return get<Vector>(c,Kind::Vector).data.at(i); }
 void replace(Cell c, std::size_t i, Cell v) { own(v); get<Vector>(c,Kind::Vector).data.at(i)=v; }
 void erase_at(Cell c, std::size_t i) {
   auto& v=get<Vector>(c,Kind::Vector).data; require(i<v.size(),"bad erase index");
   v.erase(v.begin()+static_cast<std::ptrdiff_t>(i));
 }
 void set(Cell map_cell, std::string_view k, Cell value) {
   own(value); auto& entries=get<Map>(map_cell,Kind::Map).data;
   const auto key=intern(k);
   for (auto& [id,v]:entries) if (id==key) { v=value; return; }
   entries.emplace_back(key,value);
 }
 std::optional<Cell> find(Cell map_cell, std::string_view k) const {
   const auto& entries=get<Map>(map_cell,Kind::Map).data;
   const auto it=symbols_.find(k); if (it==symbols_.end()) return std::nullopt;
   for (const auto& [id,v]:entries) if (id==it->second) return v;
   return std::nullopt;
 }
 bool erase(Cell map_cell, std::string_view k) {
   auto& entries=get<Map>(map_cell,Kind::Map).data;
   const auto key=symbols_.find(k); if (key==symbols_.end()) return false;
   for (std::size_t i=0; i<entries.size(); ++i) if (entries[i].first==key->second) {
     entries.erase(entries.begin()+static_cast<std::ptrdiff_t>(i)); return true;
   }
   return false;
 }
 void push_number(Cell c, double v) { get<F64>(c,Kind::F64).data.push_back(v); }
 std::span<double> numbers(Cell c) { return get<F64>(c,Kind::F64).data; }
 Bytes freeze(Cell root) const {
   Writer w; std::unordered_map<std::uint64_t,Encoded> done;
   std::unordered_set<std::uint64_t> active;
   const auto r=freeze_one(w,root,done,active,0); return w.finish(r);
 }
 Cell thaw(ByteSpan bytes) { const auto r=verify(bytes); return thaw_one(r,0); }
 std::size_t collect(std::span<const Cell> roots) {
   std::vector<bool> mark(slots_.size(),false);
   std::vector<Cell> work(roots.begin(),roots.end());
   while (!work.empty()) {
     const auto c=work.back(); work.pop_back(); own(c); if (!external(c)) continue;
     const auto i=static_cast<std::uint32_t>(c.payload); if (mark[i]) continue;
     mark[i]=true; children(c,[&](Cell child) { work.push_back(child); });
   }
   free_.reserve(add(free_.size(),slots_.size()));
   std::size_t freed=0;
   for (std::size_t i=0; i<slots_.size(); ++i) if (slots_[i].node && !mark[i]) {
     slots_[i].node.reset(); ++freed;
     if (slots_[i].generation!=UINT32_MAX) {
       ++slots_[i].generation; free_.push_back(static_cast<std::uint32_t>(i));
     } // Retire exhausted slots rather than wrap the generation.
   }
   return freed;
 }
};
} // namespace flexlab::live