"""Fresh-process memory residency and editing reproduction audit."""
import argparse
import gc
import json
import os
import subprocess
import time
from pathlib import Path

import torch
from PIL import Image
from scripts.benchmark import load_pipeline, generate, environment
from scripts.integrity import sha256


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', required=True)
    ap.add_argument('--output', required=True)
    ap.add_argument('--fix-offload', action='store_true')
    ap.add_argument('--native-edit', action='store_true')
    ap.add_argument('--vae-tiling', action='store_true')
    ap.add_argument('--edit-only', action='store_true')
    args = ap.parse_args()
    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=False)
    pipe = None
    def snapshot(label):
        torch.cuda.synchronize()
        row = dict(label=label, pid=os.getpid(), allocated=torch.cuda.memory_allocated(),
                   reserved=torch.cuda.memory_reserved(), peak=torch.cuda.max_memory_allocated(),
                   free_total=list(torch.cuda.mem_get_info()))
        if pipe is not None:
            row['components'] = {}
            for name in ['text_encoder', 'transformer', 'vae']:
                component = getattr(pipe, name)
                counts = {}
                states = {}
                for p in component.parameters():
                    key = f'{p.device}/{p.dtype}'
                    counts[key] = counts.get(key, 0) + p.numel()*p.element_size()
                    for attr in ('CB', 'SCB'):
                        v = getattr(p, attr, None)
                        if isinstance(v, torch.Tensor):
                            key = f'weight.{attr}/{v.device}/{v.dtype}'
                            states[key] = states.get(key, 0) + v.numel()*v.element_size()
                for m in component.modules():
                    state = getattr(m, 'state', None)
                    if state:
                        for k, v in vars(state).items():
                            if isinstance(v, torch.Tensor):
                                key = f'{k}/{v.device}/{v.dtype}'
                                states[key] = states.get(key, 0) + v.numel()*v.element_size()
                row['components'][name] = {'parameters': counts, 'bnb_state': states}
        with (out/'memory.jsonl').open('a', encoding='utf-8') as f:
            f.write(json.dumps(row)+'\n')
        print(label, 'allocated GiB', round(row['allocated']/2**30, 3), flush=True)

    (out/'environment.json').write_text(json.dumps(environment(), indent=2), encoding='utf-8')
    (out/'gpu-before.txt').write_text(subprocess.check_output(['nvidia-smi'], text=True), encoding='utf-8')
    snapshot('before_load')
    pipe = load_pipeline(args.model, offload_fix=False)
    snapshot('after_load_and_offload')
    if args.fix_offload:
        from scripts.runtime import enable_int8_cpu_offload
        enable_int8_cpu_offload(pipe)
        snapshot('after_fix')
    if args.vae_tiling:
        pipe.vae.enable_tiling()
    # Wrapping pipeline boundaries keeps model hooks intact and avoids per-step profiling overhead.
    for method in ['encode_prompt', 'prepare_latents']:
        original = getattr(pipe, method)
        def wrapper(*a, _fn=original, _name=method, **kw):
            snapshot('before_'+_name)
            result = _fn(*a, **kw)
            snapshot('after_'+_name)
            return result
        setattr(pipe, method, wrapper)
    orig_decode = pipe.vae.decode
    def decode(*a, **kw):
        snapshot('before_decode')
        result = orig_decode(*a, **kw)
        snapshot('after_decode')
        return result
    pipe.vae.decode = decode
    cases = json.loads(Path('benchmarks/cases.json').read_text(encoding='utf-8'))
    portrait, edit = cases[0], cases[-1]
    runs = [('edit_fresh', edit, 42, None, True),
            ('portrait', portrait, 42, None, True),
            ('edit_repeat', edit, 42, None, True),
            ('edit_rgb', edit, 42, 'RGB', True),
            ('edit_no_cache', edit, 42, None, False)]
    if args.native_edit:
        runs = [('edit_native', edit, 42, None, True)]
    elif args.edit_only:
        runs = [('edit_fresh', edit, 42, None, True)]
    for name, case, seed, mode, cache in runs:
        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        snapshot(name+'_start')
        size = 2048 if args.native_edit else 1024
        kwargs = dict(prompt=case['prompt'], width=size, height=size, output_resolution=size,
                      num_inference_steps=40, true_cfg_scale=1.0, use_kv_cache=cache,
                      generator=torch.Generator('cpu').manual_seed(seed))
        if case.get('input'):
            with Image.open(case['input']) as im:
                kwargs['image'] = im.convert(mode) if mode else im.copy()
        started = time.perf_counter()
        with torch.inference_mode():
            image = pipe(**kwargs).images[0]
        torch.cuda.synchronize()
        elapsed = time.perf_counter()-started
        snapshot(name+'_end')
        image.save(out/(name+'.png'))
        row = dict(name=name, seed=seed, input_mode=mode or 'original', kv_cache=cache,
                   width=size, height=size, steps=40, cfg=1.0,
                   vae_tiling=args.vae_tiling,
                   seconds=elapsed, peak_allocated_bytes=torch.cuda.max_memory_allocated(),
                   image_sha256=sha256(out/(name+'.png')),
                   input_sha256=sha256(case['input']) if case.get('input') else None)
        with (out/'records.jsonl').open('a', encoding='utf-8') as f:
            f.write(json.dumps(row)+'\n')
    if not args.native_edit and not args.edit_only:
        # Isolate the common VAE from denoising and quantization.
        with Image.open(edit['input']) as im:
            pixels = pipe.image_processor.preprocess(im.copy()).unsqueeze(2)
        with torch.inference_mode():
            latents = pipe.vae.encode(pixels.to('cuda', dtype=pipe.vae.dtype)).latent_dist.mode()
            reconstruction = pipe.vae.decode(latents, return_dict=False)[0][:, :, 0]
            pipe.image_processor.postprocess(reconstruction, output_type='pil')[0].save(out/'vae_roundtrip.png')
    # Instrumentation retains bound methods; process exit is the isolation boundary.
    snapshot('before_process_exit')


if __name__ == '__main__':
    main()