from __future__ import annotations

from pickle import dump
from struct import pack
from pathlib import Path

from ..version import __version__


class NMRWriter:
    def __init__(self, graphs: dict, data: dict):
        self._states = {'graphs': [], 'sets': []}

        for g in graphs.values():
            self._states['graphs'].append(g.get_state())

            for s in g.sets:
                s_state = data[s].get_state()
                self._states['sets'].append(s_state)

    def export(self, fname: (str | Path)):
        if isinstance(fname, str):
            fname = Path(fname)

        try:
            with fname.open(mode='wb') as fh:
                fh.write(pack('16s', b'NMREVAL' + bytes(__version__, 'utf8')))
                dump(self._states, fh)
            return True
        except Exception as e:
            print(f'Saving {fname} failed with exception {e.args}')
            return False
