// SPDX-License-Identifier: MPL-2.0
#include "vibeqc/tensor/distributed.hpp"
#include <iostream>
#include <string>
#include <optional>
using namespace vibeqc::tensor;
namespace d=vibeqc::tensor::distributed;
void check(bool b,const char* m) { if(!b)throw std::runtime_error(m); }
template<class F> void rejects(F f) { bool caught=false;try { f(); }catch(const std::exception&) {caught=true;}check(caught,"expected rejection"); }
d::Id id(d::Kind k,std::uint64_t n) {return {0,1,n,k};}
d::Token token(d::Kind k,std::uint64_t n) { return {id(k,n),1}; }
struct Fixture {
    Budget budget;
    d::Registry registry;
    Space p=Space::create(),a=Space::create(),b=Space::create();
    SourceStamp source{Space::create(),Space::create(),Space::create(),0};
    d::Operation op;
    std::shared_ptr<Allocation> la,ra;
    std::array<std::shared_ptr<const d::Binding>,3> l,r;
    Fixture(std::size_t cap=2000000,bool adjoint=false,std::size_t ranks=1,std::size_t active=0,std::size_t shift=0,std::size_t total=3)
        :budget(cap),registry(budget) {
        registry.authenticate(token(d::Kind::domain,1),p);registry.authenticate(token(d::Kind::domain,2),a);registry.authenticate(token(d::Kind::domain,3),b);
        registry.authenticate(token(d::Kind::source,1),source.source);registry.authenticate(token(d::Kind::metric,1),source.metric);registry.authenticate(token(d::Kind::frame,1),source.frame);
        d::DomainRef pd,ad,bd;
        pd.token=token(d::Kind::domain,1);pd.extent=total;
        ad.token=token(d::Kind::domain,2);ad.extent=2;
        bd.token=token(d::Kind::domain,3);bd.extent=3;
        auto descriptor=[&](unsigned n,d::DomainRef x,d::DomainRef y) {
            d::Descriptor v;v.id=id(d::Kind::tensor,n);v.source=token(d::Kind::source,1);v.metric=token(d::Kind::metric,1);v.frame=token(d::Kind::frame,1);
            v.rank=2;v.axes[0]=x;v.axes[1]=y;registry.register_tensor(v);return v;
        };
        op.left=descriptor(1,pd,ad);op.right=descriptor(2,pd,bd);op.output=descriptor(adjoint?4:3,ad,bd);
        op.id=token(d::Kind::operation,adjoint?2:1);op.value_epoch=7;op.adjoint=adjoint;op.work_limit=mul(total,7);
        op.placement={ranks,active?active:std::min(ranks,std::size_t(3)),shift,0,1};
        la=own<Complex>(budget,6);ra=own<Complex>(budget,9);
        const Complex lv[]={{1,2},{3,-1},{-2,1},{4,3},{2,-3},{-1,2}};
        const Complex rv[]={{2,-1},{-3,2},{1,4},{1,3},{2,-2},{-4,1},{3,2},{-1,-4},{2,1}};
        std::copy_n(lv,6,static_cast<Complex*>(la->base));std::copy_n(rv,9,static_cast<Complex*>(ra->base));
        for(std::size_t i=0;i<std::min(total,std::size_t(3));++i) {
            auto lk=d::full_key(op.left),rk=d::full_key(op.right);lk.ranges[0]=rk.ranges[0]={i,i+1};
            View left(la,{{Domain::range(p,1,i),2},{Domain::range(a,2),1}},2*i,Access::read,adjoint);
            View right(ra,{{Domain::range(p,1,i),3},{Domain::range(b,3),1}},3*i);
            l[i]=registry.bind(op.left,{lk,7},1,op.placement.owner(i),left,source);
            r[i]=registry.bind(op.right,{rk,7},1,op.placement.owner(i),right,source);
        }
    }
};
void submit(d::FakeExecutor& e,Fixture& f,std::size_t p) { e.submit(e.expected(p),f.l[p],f.r[p],f.op.placement.executor(p)); }
void complete(d::FakeExecutor& e) {
    e.producer_ready();auto packet=e.envelope();e.receive_complete(packet);e.send_complete();e.consumer_complete();e.release_receipt(packet.generation);
}
// Independent integer-pair oracle; neither tensor code nor complex multiplication.
Complex oracle(const Fixture& f,std::size_t a,std::size_t b) {
    long long real=0,imag=0;
    const auto* l=static_cast<const Complex*>(f.la->base);const auto* r=static_cast<const Complex*>(f.ra->base);
    for(std::size_t p=0;p<3;++p) {
        long long x=l[2*p+a].real(),y=l[2*p+a].imag(),u=r[3*p+b].real(),v=r[3*p+b].imag();
        if(f.op.adjoint)y=-y;
        real+=x*u-y*v;imag+=x*v+y*u;
    }
    return {double(real),double(imag)};
}
void numerics() {
    for(bool adjoint:{false,true})for(std::size_t ranks:{1,2,3,4})for(bool shifted:{false,true}) {
        Fixture f(2000000,adjoint,ranks,0,shifted && ranks>1?1:0);
        const auto live=f.budget.live();std::optional<d::Published> result;
        {
            d::FakeExecutor e(f.registry,f.op);auto future=e.future();
            for(std::size_t p=0;p<3;++p) {submit(e,f,p);complete(e);}
            result.emplace(e.publish());check(e.stage()==d::Stage::published,"not published");
            check(result->value()==d::ValueRef{d::full_key(f.op.output),7},"publication identity");
            for(std::size_t a=0;a<2;++a)for(std::size_t b=0;b<3;++b)check(result->at(3*a+b)==oracle(f,a,b),"oracle mismatch");
            const Complex t[]={{11,-7},{-23,-3},{7,-4},{-7,14},{16,9},{-16,6}};
            const Complex h[]={{1,1},{5,-1},{19,12},{21,0},{-16,-5},{-14,24}};
            for(std::size_t i=0;i<6;++i)check(result->at(i)==(adjoint?h[i]:t[i]),"ADR mismatch");
        }
        check(f.budget.live()==live+sizeof(d::Published)+sizeof(Allocation)+96,"publication ownership census");
        result.reset();check(f.budget.live()==live,"result leaked");
    }
    Fixture empty_rank(2000000,true,3,2);d::FakeExecutor e(empty_rank.registry,empty_rank.op);
    for(std::size_t p=0;p<3;++p) {submit(e,empty_rank,p);complete(e);}auto out=e.publish();check(out.at(0)==Complex(1,1),"empty-rank map");
}
void result_lifetime() {
    std::optional<d::Published> retained;
    {Fixture f;d::FakeExecutor e(f.registry,f.op);for(std::size_t p=0;p<3;++p){submit(e,f,p);complete(e);}retained.emplace(e.publish());}
    check(retained->at(0)==Complex(11,-7) && retained->descriptor().source==token(d::Kind::source,1),"result lost value/source lifetime");
}
void descriptors() {
    Fixture f;
    auto bad=f.op.left;bad.axes[0].extent=4;rejects([&]{f.registry.register_tensor(bad);});
    bad=f.op.left;bad.id=id(d::Kind::tensor,5);bad.axes[0].begin=1;rejects([&]{f.registry.register_tensor(bad);});
    bad=f.op.left;bad.rank=9;rejects([&]{bad.validate();});check(!(bad==bad),"oversize equality");
    bad=f.op.left;bad.dtype=static_cast<DType>(99);rejects([&]{bad.validate();});
    bad=f.op.left;bad.axes[0].begin=SIZE_MAX;rejects([&]{bad.validate();});
    bad=f.op.left;bad.source.revision=2;rejects([&]{f.registry.register_tensor(bad);});
    rejects([&]{f.registry.authenticate(token(d::Kind::domain,1),Space::create());});
    auto key=d::full_key(f.op.left);key.ranges[0]={2,4};rejects([&]{key.validate(f.op.left);});
    key.rank=99;check(!(key==key),"oversize key equality");
    auto source=f.source;source.frame=Space::create();rejects([&]{f.registry.bind(f.op.left,f.l[0]->value(),1,0,f.l[0]->view(),source);});
    View wrong(f.la,{{Domain::range(f.p,1,1),2},{Domain::range(f.a,2),1}});
    rejects([&]{f.registry.bind(f.op.left,f.l[0]->value(),1,0,wrong,f.source);});
    // Ragged dependent selected domain: parents 0/1/2, extents 2/0/3.
    Space child=Space::create();auto ct=token(d::Kind::domain,4);f.registry.authenticate(ct,child);
    for(std::size_t parent=0;parent<3;++parent) {
        auto dsc=f.op.output;dsc.id=id(d::Kind::tensor,5+parent);dsc.rank=1;
        d::DomainRef dr;dr.token=ct;dr.parent=token(d::Kind::domain,1);dr.dependent=true;dr.parent_coordinate=parent;
        dr.selected=true;dr.extent=parent==0?2:(parent==1?0:3);dr.coordinates={8,3,11};dsc.axes[0]=dr;
        f.registry.register_tensor(dsc);auto k=d::full_key(dsc);auto allocation=own<Complex>(f.budget,dr.extent);
        auto domain=f.registry.domain(dr,k.ranges[0]);View view(allocation,{{domain,1}});
        auto binding=f.registry.bind(dsc,{k,7},1,0,view,f.source);check(binding->value().key==k,"ragged key changed");
        auto wrong_key=k;wrong_key.parents[0]=(parent+1)%3;rejects([&]{wrong_key.validate(dsc);});
        if(dr.extent) { auto wrong=dsc;std::swap(wrong.axes[0].coordinates[0],wrong.axes[0].coordinates[1]);rejects([&]{f.registry.register_tensor(wrong);}); }
    }
    d::DomainRef huge;huge.token=ct;huge.selected=true;huge.extent=17;rejects([&]{huge.validate();});check(!(huge==huge),"oversize selected comparison");
    // Explicit bounds refuse before allocating a global directory.
    for(unsigned i=5;i<=29;++i)f.registry.authenticate(token(d::Kind::domain,i),Space::create());
    rejects([&]{f.registry.authenticate(token(d::Kind::domain,30),Space::create());});
}
void lifetimes() {
    Fixture f;auto baseline=f.budget.live();
    {
        d::FakeExecutor e(f.registry,f.op);auto future=e.future();submit(e,f,0);
        check(e.stage()==d::Stage::producing,"producer gate");rejects([&]{Lease w(f.la,true);});
        auto charged=f.budget.live();e.producer_ready();check(f.budget.live()==charged,"producer control census changed");
        auto packet=e.envelope();e.send_complete();check(e.pending(),"send released remote consumer");
        e.receive_complete(packet);check(e.stage()==d::Stage::consuming,"receive gate");rejects([&]{Lease w(f.la,true);});
        check(e.accepted()==0,"receive counted contribution");e.consumer_complete();check(e.accepted()==1 && e.pending(),"missing receipt gate");
        rejects([&]{Lease w(f.la,true);});e.release_receipt(packet.generation);check(!e.pending(),"receipt did not retire");
        // Completion order can be receive -> consume -> receipt -> send.
        submit(e,f,1);e.producer_ready();packet=e.envelope();e.receive_complete(packet);e.consumer_complete();e.release_receipt(packet.generation);
        check(e.pending(),"send buffer released early");e.send_complete();check(!e.pending(),"send did not retire");
        submit(e,f,2);complete(e);auto result=e.publish();check(result.at(0)==oracle(f,0,0),"delayed result");
    }
    check(f.budget.live()==baseline,"completion leak");
    {
        d::FakeExecutor e(f.registry,f.op);
        {auto future=e.future();submit(e,f,0);e.producer_ready();}
        check(e.stage()==d::Stage::failed && e.pending(),"dropped future freed ticket");rejects([&]{Lease w(f.la,true);});
        auto packet=e.envelope();e.receive_complete(packet);e.send_complete();e.consumer_complete();e.release_receipt(packet.generation);
        check(!e.pending() && e.accepted()==0,"cancel counted contribution");rejects([&]{e.publish();});
    }
    check(f.budget.live()==baseline,"cancel leak");
    // Binding retains native owners even when the importer releases its handles.
    {
        d::FakeExecutor e(f.registry,f.op);submit(e,f,0);std::weak_ptr<Allocation> weak=f.la;
        f.la.reset();for(auto& x:f.l)x.reset();check(!weak.expired(),"owner not retained");complete(e);check(weak.expired(),"owner did not release after receipt");
    }
}
void failures() {
    for(unsigned fault=0;fault<14;++fault) {
        Fixture f;const auto baseline=f.budget.live();
        {
            d::FakeExecutor e(f.registry,f.op);
            if(fault==0) {auto c=e.expected(1);rejects([&]{e.submit(c,f.l[1],f.r[1],0);});}
            else if(fault==1) {submit(e,f,0);complete(e);rejects([&]{submit(e,f,0);});}
            else if(fault==2) {submit(e,f,0);complete(e);rejects([&]{e.publish();});}
            else if(fault==3) { {Lease w(f.la,true);} rejects([&]{submit(e,f,0);}); }
            else if(fault==4) {e.fail_provider();rejects([&]{submit(e,f,0);});}
            else if(fault==5) {submit(e,f,0);rejects([&]{submit(e,f,1);});}
            else {
                submit(e,f,0);e.producer_ready();auto packet=e.envelope();
                if(fault==6)--packet.contribution.output.epoch;
                if(fault==7)++packet.placement_epoch;
                if(fault==8)++packet.generation;
                if(fault==9)++packet.sender;
                if(fault==10)++packet.bytes;
                if(fault==11)e.corrupt_byte(0);
                if(fault==12)packet.contribution.output.key.rank=99;
                if(fault==13)++packet.contribution.operation.revision;
                rejects([&]{e.receive_complete(packet);});
                check(e.accepted()==0 && e.pending(),"bad payload applied/released");
            }
            check(e.stage()==d::Stage::failed,"failure not sticky");rejects([&]{e.publish();});e.drain_fake();
        }
        check(f.budget.live()==baseline,"failure reservation leak");
    }
    Fixture f;
    BorrowScope scope(f.budget);Complex values[2]{};
    View borrowed(borrow(f.budget,scope,values,2),{{Domain::range(f.p,1),2},{Domain::range(f.a,2),1}});
    auto bound=f.registry.bind(f.op.left,f.l[0]->value(),1,0,borrowed,f.source);
    d::FakeExecutor e(f.registry,f.op);rejects([&]{e.submit(e.expected(0),bound,f.r[0],0);});
    // Numerical failure leaves only private partially updated storage.
    Fixture nonfinite;static_cast<Complex*>(nonfinite.la->base)[0]={std::numeric_limits<double>::infinity(),0};
    d::FakeExecutor bad(nonfinite.registry,nonfinite.op);submit(bad,nonfinite,0);rejects([&]{bad.producer_ready();});rejects([&]{bad.publish();});bad.drain_fake();
}
void budgets() {
    std::size_t exact=0,base=0;
    { Fixture f; base=f.budget.live();exact=add(base,d::FakeExecutor::admission_bytes(f.registry,f.op)); }
    {
        Fixture f(exact);check(f.budget.live()==base,"unstable preflight census");
        {d::FakeExecutor e(f.registry,f.op);check(f.budget.live()==exact,"admission census mismatch");
         for(std::size_t p=0;p<3;++p){submit(e,f,p);complete(e);}auto result=e.publish();check(result.at(0)==Complex(11,-7),"exact cap result");}
        check(f.budget.live()==base,"exact cap leak");check(f.budget.peak()<=exact,"tracked peak over cap");
    }
    {Fixture f(exact-1);auto live=f.budget.live();rejects([&]{d::FakeExecutor e(f.registry,f.op);});check(f.budget.live()==live,"below cap leak");}
    // Globally empty reduction still owns control and private-output storage.
    Fixture f;auto op=f.op;op.left.axes[0].extent=op.right.axes[0].extent=0;
    // Independent empty-domain fixture needs a new registry generation.
    Budget budget(2000000);d::Registry registry(budget);
    auto p=Space::create(),a=Space::create(),b=Space::create();
    registry.authenticate(op.left.axes[0].token,p);registry.authenticate(op.left.axes[1].token,a);registry.authenticate(op.right.axes[1].token,b);
    registry.authenticate(op.left.source,f.source.source);registry.authenticate(op.left.metric,f.source.metric);registry.authenticate(op.left.frame,f.source.frame);
    registry.register_tensor(op.left);registry.register_tensor(op.right);registry.register_tensor(op.output);
    d::FakeExecutor empty(registry,op);auto zero=empty.publish();for(std::size_t i=0;i<zero.size();++i)check(zero.at(i)==Complex{},"empty reduction not zero");
}
std::size_t bounded_run(std::size_t total) {
    Fixture f(2000000,false,3,0,0,total);
    for(auto& x:f.l)x.reset();for(auto& x:f.r)x.reset();
    d::FakeExecutor e(f.registry,f.op);
    for(std::size_t p=0;p<total;++p) {
        auto lk=d::full_key(f.op.left),rk=d::full_key(f.op.right);lk.ranges[0]=rk.ranges[0]={p,p+1};
        View lv(f.la,{{Domain::range(f.p,1,p),2},{Domain::range(f.a,2),1}});
        View rv(f.ra,{{Domain::range(f.p,1,p),3},{Domain::range(f.b,3),1}});
        auto left=f.registry.bind(f.op.left,{lk,7},1,f.op.placement.owner(p),lv,f.source);
        auto right=f.registry.bind(f.op.right,{rk,7},1,f.op.placement.owner(p),rv,f.source);
        e.submit(e.expected(p),left,right,f.op.placement.executor(p));complete(e);
    }
    auto out=e.publish();check(out.at(0)==Complex(4*double(total),3*double(total)),"long-stream sum");
    return f.budget.peak();
}
void bounded() {
    const auto small=bounded_run(3),large=bounded_run(10003);
    check(small==large,"metadata grows with logical blocks");
    check(10003*80>large*4,"logical data not larger than tracked state");
}
void admission_failures() {
    Fixture f;
    auto op=f.op;op.work_limit=17;rejects([&]{d::FakeExecutor e(f.registry,op);});
    auto source=f.source;source.source=Space::create();
    // Same metadata in a foreign registry cannot launder a native source token.
    d::Registry foreign(f.budget);
    foreign.authenticate(f.op.left.axes[0].token,f.p);foreign.authenticate(f.op.left.axes[1].token,f.a);
    foreign.authenticate(f.op.left.source,source.source);foreign.authenticate(f.op.left.metric,source.metric);foreign.authenticate(f.op.left.frame,source.frame);
    foreign.register_tensor(f.op.left);
    auto forged=foreign.bind(f.op.left,f.l[0]->value(),1,0,f.l[0]->view(),source);
    {d::FakeExecutor e(f.registry,f.op);rejects([&]{e.submit(e.expected(0),forged,f.r[0],0);});}
    for(unsigned variant=0;variant<4;++variant) {
        auto value=f.l[0]->value();if(variant==0)--value.epoch;
        auto binding=f.registry.bind(f.op.left,value,variant==1?2:1,variant==2?1:0,f.l[0]->view(),f.source);
        d::FakeExecutor e(f.registry,f.op);
        rejects([&]{e.submit(e.expected(0),binding,f.r[0],variant==3?1:0);});
    }
    // Thread checks refuse before touching operation state; no MPI is linked.
    {d::FakeExecutor e(f.registry,f.op);bool refused=false;
     std::thread worker([&]{try {e.stage();}catch(const std::exception&){refused=true;}});worker.join();check(refused,"worker accepted");}
    // A wrong receipt cannot permit slot reuse, even after consumer completion.
    {d::FakeExecutor e(f.registry,f.op);submit(e,f,0);e.producer_ready();auto packet=e.envelope();
     e.receive_complete(packet);e.send_complete();e.consumer_complete();
     rejects([&]{e.release_receipt(packet.generation+1);});check(e.pending(),"stale receipt released slot");e.drain_fake();}
    // Failure after admission must preserve any previously published value.
    {d::FakeExecutor first(f.registry,f.op);for(std::size_t p=0;p<3;++p){submit(first,f,p);complete(first);}
     auto old=first.publish();const auto before=old.at(0);
     d::FakeExecutor second(f.registry,f.op);submit(second,f,0);complete(second);second.fail_provider();
     rejects([&]{second.publish();});check(old.at(0)==before,"failed output mutated old epoch");}
}
void codec() {
    const std::uint8_t text[]={'1','2','3','4','5','6','7','8','9'};
    check(d::crc32(text,9)==0xcbf43926u,"CRC known answer");
    std::uint8_t bytes[8];d::encode_double(1.,bytes);check(bytes[0]==0x3f && bytes[1]==0xf0,"binary64 endian");
    for(double x:{0.,-0.,1.,-7.,0.125,std::numeric_limits<double>::max()}) {d::encode_double(x,bytes);double y=d::decode_double(bytes);check(std::memcmp(&x,&y,8)==0,"binary64 roundtrip");}
}
int main(int argc,char** argv) {
    try {const std::string c=argc>1?argv[1]:"";
        if(c=="numerics"){numerics();result_lifetime();}else if(c=="descriptors")descriptors();else if(c=="lifetimes")lifetimes();
        else if(c=="failures")failures();else if(c=="budgets")budgets();else if(c=="codec")codec();else if(c=="bounded")bounded();else if(c=="admission")admission_failures();else throw std::runtime_error("unknown case");
        std::cout<<c<<" PASS\n";
    } catch(const std::exception& e) {std::cerr<<e.what()<<'\n';return 1;}
}