"""Measure the declared 2048px editing suite in a fresh process per precision."""
import argparse
import gc
import json
import os
import subprocess
import time
from pathlib import Path

import psutil
import torch
from scripts.benchmark import environment, generate, load_pipeline, memory_status
from scripts.integrity import sha256
from scripts.provenance import model_identity


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', required=True)
    ap.add_argument('--output', required=True)
    ap.add_argument('--cases', default='benchmarks/editing.json')
    ap.add_argument('--inputs', required=True, help='Directory containing the published input PNGs')
    args = ap.parse_args()
    cases = json.loads(Path(args.cases).read_text(encoding='utf-8'))
    for case in cases:
        case['input'] = str(Path(args.inputs) / case['input_file'])
        if sha256(case['input']) != case['input_sha256']:
            raise ValueError('Input hash mismatch: ' + case['id'])
    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=False)
    env = environment()
    env.update(pid=os.getpid(), before_load_memory=memory_status(),
               model_identity=model_identity(args.model), offload='model', offload_aux_fix=True,
               warmup=True, generator_device='cpu', cases_sha256=sha256(args.cases),
               source_sha256={n: sha256(Path(__file__).with_name(n)) for n in
                              ('benchmark_edits.py', 'benchmark.py', 'runtime.py')})
    for key, query in [('gpu_before', '--query-gpu=name,driver_version,memory.used,memory.free'),
                       ('gpu_processes_before', '--query-compute-apps=pid,process_name,used_memory')]:
        env[key] = subprocess.check_output(['nvidia-smi', query, '--format=csv'], text=True)
    if env['before_load_memory']['allocated_bytes'] != 0:
        raise ValueError('Expected a fresh CUDA allocation baseline')
    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    pipe = load_pipeline(args.model)
    pipe.vae.enable_tiling()
    env.update(load_seconds=time.perf_counter()-start,
               load_peak_allocated_bytes=torch.cuda.max_memory_allocated(),
               after_load_memory=memory_status(), vae_tiling=True,
               vae_tiles={n: getattr(pipe.vae, n) for n in
                          ('tile_sample_min_height', 'tile_sample_min_width',
                           'tile_sample_stride_height', 'tile_sample_stride_width')})
    (out/'environment.json').write_text(json.dumps(env, indent=2), encoding='utf-8')
    print('Full-settings 2048px warmup, seed 0 (excluded)', flush=True)
    generate(pipe, cases[0], 0, 2048, 2048, 40)
    torch.cuda.synchronize()
    for case in cases:
        print('Generating ' + case['id'], flush=True)
        gc.collect()
        torch.cuda.empty_cache()
        before = memory_status()
        torch.cuda.reset_peak_memory_stats()
        start = time.perf_counter()
        image = generate(pipe, case, 42, 2048, 2048, 40)
        torch.cuda.synchronize()
        seconds = time.perf_counter()-start
        path = out/(case['id']+'-s42.png')
        image.save(path)
        row = dict(case_id=case['id'], category=case['category'], prompt=case['prompt'],
                   seed=42, width=2048, height=2048, output_resolution=2048, steps=40,
                   cfg=1.0, kv_cache=True, offload='model', vae_tiling=True,
                   input_file=case['input_file'], input_sha256=case['input_sha256'],
                   before_memory=before, after_memory=memory_status(), seconds=seconds,
                   peak_allocated_bytes=torch.cuda.max_memory_allocated(),
                   peak_reserved_bytes=torch.cuda.max_memory_reserved(),
                   process_rss_after_bytes=psutil.Process().memory_info().rss,
                   image=path.name, image_sha256=sha256(path), image_mode=image.mode,
                   actual_size=list(image.size))
        with (out/'records.jsonl').open('a', encoding='utf-8') as stream:
            stream.write(json.dumps(row, ensure_ascii=False)+'\n')
        print(f'{case["id"]}: {seconds:.2f}s, {row["peak_allocated_bytes"]/2**30:.2f} GiB', flush=True)


if __name__ == '__main__':
    main()