"""Experimental W=1 CPU tensor contraction using the bounded MPI protocol.

All ranks, including empty ranks, call ``contract`` on the MPI owner thread.
The native Session factory authenticates a caller-supplied stable manifest;
IDs are application identities, never inferred from pointers or data hashes.
No MPI calls occur in native kernels or provider callbacks.
"""
from __future__ import annotations

import numpy as np

from .mpi import BoundedSumPhase
from .mpi_transfer import BoundedTransferPhase

__all__ = ["contract"]
_failed_cpu_transfers = []


def contract(session_factory, provide_row):
    """Return native Published on the output owner, None on other ranks.

    ``session_factory(rank, ranks)`` constructs a native ``_tensor_cpu.Session``
    with a per-rank Budget cap and declared live provider-byte inventory.
    ``provide_row(p)`` returns contiguous complex128 left/right rows, only on
    their declared owner. It must terminate, perform no MPI, and respect that
    inventory. Caller keeps the operation manifest identical on every rank.

    Recoverable refusals discard private output after requests drain. Fatal
    transport errors retain sessions and phases until job termination. There
    is no collective rollback, retry, full-tensor gather or GPU integration.
    """
    agreement = BoundedSumPhase()
    transfer = BoundedTransferPhase()
    session = None
    ticket = None
    empty = np.empty(0, dtype=np.float64)

    def agree(callback):
        def prepare():
            callback()
            return empty
        # A fixed control contract admits even mismatching native manifests.
        agreement.sum(prepare, phase_id=0, epoch=0, count=0, max_elements=0)

    def admit():
        nonlocal session
        session = session_factory(transfer.world.rank, transfer.world.size)
        if session.ranks != transfer.world.size:
            raise ValueError("tensor placement does not match communicator")

    try:
        agree(admit)
        # Compare full canonical bounded metadata, including on empty ranks,
        # before any rank can choose a differing contribution loop length.
        for destination in range(transfer.world.size):
            transfer.begin(
                lambda: session.manifest, source=0, destination=destination,
                nbytes=session.manifest.size, max_bytes=session.manifest.size,
                phase_id=0, epoch=0,
            ).finish(session.verify_manifest)
        for p in range(session.count):
            def prepare():
                left, right = provide_row(p)
                return session.produce(p, p, left, right)
            ticket = transfer.begin(
                prepare, source=session.source(p), destination=session.output_owner,
                nbytes=session.packet_bytes, max_bytes=session.packet_bytes,
                phase_id=session.phase_id, epoch=session.epoch,
            )
            # Local send completion never retires the native input leases.
            ticket.progress()
            ticket.finish(lambda packet: session.consume(p, p, packet))
            agree(lambda: session.retire(p, p))
        agree(session.ready)
        transfer.close()
        result = session.publish() if transfer.world.rank == session.output_owner else None
    except BaseException:
        if (transfer.poisoned or agreement.poisoned or transfer.active is not None):
            _failed_cpu_transfers.append((session, transfer, agreement))
        else:
            # CollectiveRefusal guarantees drained requests. The same path
            # also covers a post-close local publication allocation failure.
            try:
                if not transfer.closed:
                    transfer.close()
            except BaseException:
                _failed_cpu_transfers.append((session, transfer, agreement))
                raise
            # Drop ticket/phase arrays before releasing their imported census.
            ticket = transfer = agreement = None
            if session is not None:
                session.release_safe()
        raise
    ticket = transfer = agreement = None
    session.release_safe(True)
    return result