import torch, struct, zlib, numpy as np, itertools, os, gc, sys, psutil
from diffusers import StableDiffusionXLPipeline

_ARTIFACT = "Marius_SDXL_V18.lzr2"
_BASE = "stabilityai/stable-diffusion-xl-base-1.0"

def _stat():
    process = psutil.Process(os.getpid())
    return process.memory_info().rss // (1024 ** 3)

def inject_marius(path, pipe):
    if not os.path.exists(path):
        raise FileNotFoundError(f"Missing artifact: {path}")
    print("Initializing streaming engine...")

    _opts, u = {}, pipe.unet
    _g_v = lambda d: np.array(list(itertools.product([-1, 0, 1], repeat=d)), dtype=np.float32)
    idx = 0

    with open(path, "rb") as f:
        if f.read(4) != b"LZR2":
            raise ValueError("Invalid signature")

        while True:
            lkb = f.read(4)
            if not lkb: break

            key = f.read(struct.unpack('I', lkb)[0]).decode('utf-8')
            ls = struct.unpack('I', f.read(4))[0]
            sh = [struct.unpack('I', f.read(4))[0] for _ in range(ls)]
            tf = struct.unpack('B', f.read(1))[0]
            _w = None

            if tf == 1:
                dp, C = struct.unpack('I', f.read(4))[0], sh[0]
                _a = np.frombuffer(f.read(C*dp*4), dtype=np.float32).reshape(C, dp)
                _mn = np.frombuffer(f.read(C*4), dtype=np.float32)
                _sc = np.frombuffer(f.read(C*4), dtype=np.float32)
                lz = struct.unpack('I', f.read(4))[0]
                _ix_flat = np.frombuffer(zlib.decompress(f.read(lz)), dtype=np.uint16)
                n_blocks = _ix_flat.size // C
                _ix = _ix_flat.reshape(C, n_blocks)
                no = struct.unpack('I', f.read(4))[0]
                N_feat = int(np.prod(sh[1:])) if len(sh) > 1 else 1

                if dp not in _opts:
                    _opts[dp] = _g_v(dp)
                rc = _opts[dp][_ix].reshape(C, -1) if n_blocks > 0 else np.zeros((C, 0), dtype=np.float32)

                fb = np.zeros((C, N_feat), dtype=np.float32)
                vw = min(rc.shape[1], N_feat)
                if vw > 0:
                    fb[:, :vw] = rc[:, :vw]
                fb = (fb + _mn[:, None]) * _sc[:, None]

                if no > 0:
                    md = max(C, n_blocks) * dp
                    fmt, fsz = ('H', 8) if md < 65536 else ('I', 12)
                    dt = np.dtype([('r', np.uint16 if fmt=='H' else np.uint32),
                                   ('c', np.uint16 if fmt=='H' else np.uint32),
                                   ('v', np.float32)])
                    batch = np.frombuffer(f.read(no * fsz), dtype=dt)
                    m = (batch['r'] < C) & (batch['c'] < N_feat)
                    vb = batch[m]
                    fb[vb['r'], vb['c']] = vb['v']

                _w = torch.from_numpy(fb.reshape(sh).astype(np.float16))

            if _w is not None:
                try:
                    t = u
                    pts = key.split('.')
                    for p in pts[:-1]:
                        t = getattr(t, p)
                    getattr(t, pts[-1]).data.copy_(_w.to(pipe.device, dtype=torch.float16))
                except:
                    pass
                del _w

            idx += 1
            if idx % 10 == 0:
                sys.stdout.write(f"\r[STREAM] Module {idx:04d} | RAM: {_stat()}GB")
                sys.stdout.flush()
            if idx % 200 == 0:
                gc.collect()

    print(f"\nStream complete ({idx} modules loaded)")

def get_pipe():
    print("Loading base architecture...")
    pipe = StableDiffusionXLPipeline.from_pretrained(
        _BASE,
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True
    )
    pipe.enable_model_cpu_offload()
    inject_marius(_ARTIFACT, pipe)
    return pipe