// SPDX-License-Identifier: MPL-2.0
#pragma once
#include "vibeqc/tensor/distributed_wire.hpp"

namespace vibeqc::tensor::distributed {
// A CPU rank-local role. MPI remains entirely in #820's Python owner thread.
// This class imports the known transport inventory into the existing Budget;
// neither Python object overhead nor MPI-internal memory is a tracked guarantee.
class CpuTransferSession {
    Operation op_;
    std::size_t rank_;
    Budget budget_;
    Registry registry_;
    Budget::Reservation control_,external_,allowance_,result_control_;
    std::shared_ptr<Allocation> left_,right_,partial_,carry_,sum_,correction_;
    std::optional<Array<std::uint8_t>> packet_;
    std::shared_ptr<const Binding> lb_,rb_;
    std::optional<Lease> ll_,rl_;
    WireHeader manifest_;
    std::size_t next_=0;
    std::uint64_t generation_=0;
    bool produced_=false,consumed_=false,failed_=false,published_=false,released_=false;
    std::thread::id owner_=std::this_thread::get_id();
    bool producer() const { return rank_<op_.placement.active && rank_<count(); }
    bool destination() const { return rank_==op_.placement.output_owner; }
    void owner() const { require(owner_==std::this_thread::get_id(),"CPU tensor session requires its owner thread"); }
    void usable() const { owner();require(!failed_ && !published_ && !released_,"CPU tensor session is terminal"); }
    void check_step(std::size_t p,std::uint64_t generation) const {
        usable();require(p==next_ && p<count() && generation==p && generation<=INT64_MAX,"missing/duplicate tensor contribution");
    }
    void intern(Token t) {
        try { registry_.resolve(t); }
        catch(const std::invalid_argument&) {registry_.authenticate(t,Space::create());}
    }
    void release_inputs() noexcept {ll_.reset();rl_.reset();lb_.reset();rb_.reset();}
    void release_workspace() noexcept {
        release_inputs();left_.reset();right_.reset();partial_.reset();carry_.reset();
        // Python read-only packet views retain this session. Keep their storage
        // charged and alive through terminal release until the session dies.
        sum_.reset();correction_.reset();allowance_.release();result_control_.release();
    }
    static std::size_t n(const Operation& op) { return mul(op.output.axes[0].extent,op.output.axes[1].extent); }
    static void schema(const Operation& op,std::size_t rank,std::size_t provider_bytes) {
        op.id.validate(Kind::operation);op.placement.validate();
        require(rank<op.placement.ranks && op.placement.ranks<=INT32_MAX && op.placement.executor_shift==0,"CPU transfer requires colocated input/executor");
        require(op.value_epoch && op.value_epoch<=INT64_MAX && op.id.id.serial<=INT64_MAX,"MPI phase integer overflow");
        for(const auto* d:{&op.left,&op.right,&op.output}) {
            d->validate();require(d->rank==2 && d->dtype==DType::complex_fp64,"complex matrix transfer required");
            for(std::size_t i=0;i<2;++i)require(!d->axes[i].selected && !d->axes[i].dependent,"first CPU transfer requires affine independent domains");
        }
        require(op.left.axes[0]==op.right.axes[0] && op.left.axes[1]==op.output.axes[0] && op.right.axes[1]==op.output.axes[1],"transfer contraction domains mismatch");
        for(const auto* d:{&op.right,&op.output})require(d->source==op.left.source && d->metric==op.left.metric && d->frame==op.left.frame,"transfer source mismatch");
        require(op.left.axes[0].extent<=INT64_MAX && mul(op.left.axes[0].extent,add(n(op),1))<=op.work_limit,"transfer work cap exceeded");
        if(rank<op.placement.active && rank<op.left.axes[0].extent)
            require(provider_bytes>=mul(add(op.left.axes[1].extent,op.right.axes[1].extent),16),"missing provider live-byte inventory");
        const auto header=encode_header(op,0,0,0);
        require(add(header.size,mul(n(op),16))<=INT32_MAX,"transfer exceeds bounded MPI byte count");
    }
public:
    // Includes registry, inputs/copies, output, ephemeral kernel controls and
    // imported #820 controls (144 sum + 232 transfer) plus its one payload copy.
    static std::size_t required_bytes(const Operation& op,std::size_t rank,std::size_t provider_bytes) {
        schema(op,rank,provider_bytes);
        const auto header=encode_header(op,0,0,0);const auto elements=n(op);
        const bool src=rank<op.placement.active && rank<op.left.axes[0].extent,dst=rank==op.placement.output_owner;
        const auto packet=add(header.size,mul(elements,16));
        auto bytes=add(Budget::baseline_bytes(),add(Registry::bytes(),sizeof(CpuTransferSession)));
        // Empty participants still receive a bounded manifest and own all controls.
        bytes=add(bytes,add(provider_bytes,add(376,(src || dst)?packet:header.size)));
        if(src)bytes=add(bytes,add(stream_control_bytes,add(mul(4,sizeof(Allocation)),
            add(mul(add(op.left.axes[1].extent,op.right.axes[1].extent),16),add(mul(elements,32),packet)))));
        if(dst)bytes=add(bytes,add(sizeof(Published),add(mul(2,sizeof(Allocation)),mul(elements,32))));
        return bytes;
    }
    CpuTransferSession(Operation op,std::size_t rank,std::size_t cap,std::size_t provider_bytes)
        :op_(std::move(op)),rank_(rank),budget_(cap),registry_(budget_),control_(budget_.reserve(sizeof(CpuTransferSession))) {
        const auto required=required_bytes(op_,rank,provider_bytes);
        require(cap>=required,"CPU tensor rank budget refused");
        for(const auto* d:{&op_.left,&op_.right,&op_.output}) {
            intern(d->source);intern(d->metric);intern(d->frame);
            for(std::size_t i=0;i<d->rank;++i)intern(d->axes[i].token);
            registry_.register_tensor(*d);
        }
        op_.validate(registry_);manifest_=encode_header(op_,0,0,0);
        external_=budget_.reserve(add(376,add((producer() || destination())?packet_bytes():manifest_.size,provider_bytes)));
        if(producer()) {
            allowance_=budget_.reserve(stream_control_bytes);
            left_=own<Complex>(budget_,op_.left.axes[1].extent);right_=own<Complex>(budget_,op_.right.axes[1].extent);
            partial_=own<Complex>(budget_,elements());carry_=own<Complex>(budget_,elements());
            packet_.emplace(budget_,packet_bytes());
        }
        if(destination()) {
            result_control_=budget_.reserve(sizeof(Published));sum_=own<Complex>(budget_,elements());correction_=own<Complex>(budget_,elements());
            std::fill_n(static_cast<Complex*>(sum_->base),elements(),Complex{});
            std::fill_n(static_cast<Complex*>(correction_->base),elements(),Complex{});
        }
        require(budget_.live()==required,"CPU transfer census mismatch");
    }
    CpuTransferSession(const CpuTransferSession&)=delete;
    std::size_t count() const { return op_.left.axes[0].extent; }
    std::size_t elements() const { return n(op_); }
    std::size_t packet_bytes() const { return add(manifest_.size,mul(elements(),16)); }
    std::size_t source(std::size_t p) const { require(p<count(),"source row out of range");return op_.placement.owner(p); }
    std::size_t output_owner() const {return op_.placement.output_owner;}
    std::size_t ranks() const {return op_.placement.ranks;}
    std::uint64_t phase_id() const {return op_.id.id.serial;}
    std::uint64_t epoch() const {return op_.value_epoch;}
    std::size_t live() const {return budget_.live();}
    std::size_t peak() const {return budget_.peak();}
    bool pending() const {return produced_ || consumed_;}
    bool failed() const {return failed_;}
    const WireHeader& manifest() const {usable();return manifest_;}
    void verify_manifest(const std::uint8_t* data,std::size_t size) const {usable();validate_header(manifest_,data,size);}
    const std::uint8_t* produce(std::size_t p,std::uint64_t generation,const Complex* l,std::size_t nl,const Complex* r,std::size_t nr) {
        check_step(p,generation);
        try {
            require(rank_==source(p) && !pending(),"invalid producer/credit");
            require(nl==op_.left.axes[1].extent && nr==op_.right.axes[1].extent && (l || !nl) && (r || !nr),"provider row shape mismatch");
            { Lease lw(left_,true),rw(right_,true);
              std::copy_n(l,nl,static_cast<Complex*>(left_->base));std::copy_n(r,nr,static_cast<Complex*>(right_->base)); }
            allowance_.release();const auto base=budget_.live();
            {
                auto lk=full_key(op_.left),rk=full_key(op_.right);lk.ranges[0]=rk.ranges[0]={p,p+1};
                View lv(left_,{{registry_.domain(op_.left.axes[0],{p,p+1}),signed_size(nl)},
                               {registry_.domain(op_.left.axes[1],{0,nl}),1}},0,Access::read,op_.adjoint);
                View rv(right_,{{registry_.domain(op_.right.axes[0],{p,p+1}),signed_size(nr)},
                                {registry_.domain(op_.right.axes[1],{0,nr}),1}});
                lb_=registry_.bind(op_.left,{lk,epoch()},op_.placement.epoch,rank_,lv,registry_.stamp(op_.left));
                rb_=registry_.bind(op_.right,{rk,epoch()},op_.placement.epoch,rank_,rv,registry_.stamp(op_.right));
                lb_->check_owned();rb_->check_owned();ll_.emplace(left_,false);rl_.emplace(right_,false);
                auto a=registry_.domain(op_.output.axes[0],{0,nl}),b=registry_.domain(op_.output.axes[1],{0,nr});
                View sv(partial_,{{a,signed_size(nr)},{b,1}},0,Access::read_write),cv(carry_,{{a,signed_size(nr)},{b,1}},0,Access::read_write);
                std::fill_n(static_cast<Complex*>(partial_->base),elements(),Complex{});std::fill_n(static_cast<Complex*>(carry_->base),elements(),Complex{});
                CompensatedPairPlan plan({lv,registry_.stamp(op_.left)},{rv,registry_.stamp(op_.right)},sv,cv,registry_.stamp(op_.left));plan.execute();
            }
            const auto dynamic=budget_.live()-base;require(dynamic<=stream_control_bytes,"CPU control envelope exceeded");
            allowance_=budget_.reserve(stream_control_bytes-dynamic);
            auto* payload=packet_->data()+manifest_.size;
            const auto* sum=static_cast<const Complex*>(partial_->base);const auto* carry=static_cast<const Complex*>(carry_->base);
            for(std::size_t i=0;i<elements();++i) {const auto v=finite_complex(sum[i]+carry[i]);encode_double(v.real(),payload+16*i);encode_double(v.imag(),payload+16*i+8);}
            const auto h=encode_header(op_,p,generation,crc32(payload,mul(elements(),16)));
            std::copy_n(h.bytes.data(),h.size,packet_->data());generation_=generation;produced_=true;
            return packet_->data();
        } catch(...) {failed_=true;throw;}
    }
    void consume(std::size_t p,std::uint64_t generation,const std::uint8_t* data,std::size_t size) {
        check_step(p,generation);
        try {
            require(destination() && !consumed_,"invalid consumer/duplicate contribution");
            validate_packet(op_,p,generation,data,size);
            auto* s=static_cast<Complex*>(sum_->base);auto* c=static_cast<Complex*>(correction_->base);
            const auto* payload=data+manifest_.size;
            for(std::size_t i=0;i<elements();++i)compensated_add({decode_double(payload+16*i),decode_double(payload+16*i+8)},s[i],c[i]);
            generation_=generation;consumed_=true;
        } catch(...) {failed_=true;throw;}
    }
    // Call only after #820's successful terminal agreement/release receipt.
    void retire(std::size_t p,std::uint64_t generation) {
        check_step(p,generation);
        require((rank_!=source(p) || produced_) && (!destination() || consumed_) && (!pending() || generation==generation_),"incomplete native ticket retirement");
        release_inputs();produced_=consumed_=false;++next_;
        if(producer()) {allowance_.release();allowance_=budget_.reserve(stream_control_bytes);}
    }
    void ready() const {
        usable();require(!pending() && next_==count(),"missing contribution at publication");
        if(destination()) {
            const auto* s=static_cast<const Complex*>(sum_->base);const auto* c=static_cast<const Complex*>(correction_->base);
            for(std::size_t i=0;i<elements();++i)finite_complex(s[i]+c[i]);
        }
    }
    Published publish() {
        ready();require(destination(),"only output owner publishes");
        auto* s=static_cast<Complex*>(sum_->base);auto* c=static_cast<const Complex*>(correction_->base);
        for(std::size_t i=0;i<elements();++i)s[i]=finite_complex(s[i]+c[i]);
        sum_->writable=false;published_=true;
        Published result(std::move(result_control_),std::move(sum_),{full_key(op_.output),epoch()},registry_,op_.output);
        release_workspace();return result;
    }
    // Only after a known collective refusal or normal communicator shutdown.
    // Fatal transport paths must retain the entire session instead of calling this.
    void release_safe(bool success=false) {
        owner();if(success && !published_)ready();
        failed_=failed_ || (!success && !published_);released_=true;
        release_workspace();produced_=consumed_=false;external_.release();
    }
};
} // namespace vibeqc::tensor::distributed