import matplotlib.pyplot as plt
import pickle
import zipfile
import tempfile
import os


def embed_figure_into_png(png_path: str, figure):
    """
    Embeds a Matplotlib Figure object into a PNG file.

    The PNG remains a valid image and can be opened by any viewer,
    but it also contains a serialized copy of the Figure for later recovery.
    """
    fig_bytes = pickle.dumps(figure)

    # Create a temporary ZIP file that stores the pickled Figure
    with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as temp_zip_file:
        temp_zip_path = temp_zip_file.name
        with zipfile.ZipFile(temp_zip_file, "w", zipfile.ZIP_DEFLATED) as zf:
            zf.writestr("figure.pkl", fig_bytes)

    try:
        # Read the original PNG and the ZIP content
        with open(png_path, "rb") as f_png, open(temp_zip_path, "rb") as f_zip:
            png_data = f_png.read()
            zip_data = f_zip.read()

        # Re-write PNG with ZIP data appended
        with open(png_path, "wb") as f_out:
            f_out.write(png_data)
            f_out.write(zip_data)

        print(f"[OK] Figure embedded into file: {png_path}")

    finally:
        # Always clean up the temporary ZIP file
        os.remove(temp_zip_path)


# --- Usage example ---

png_path = "plot.png"

# 1. Create and save a plot
fig, ax = plt.subplots()
line, = ax.plot([1, 2, 3], [4, 9, 2], label="Original line")
fig.savefig(png_path)

# Close the figure to simulate "lost source code"
plt.close(fig)

# 2. Embed Figure object inside the PNG
embed_figure_into_png(png_path, fig)


def extract_figure_from_png(png_path: str):
    """
    Extracts a Matplotlib Figure object from a PNG file,
    if it was previously embedded there.
    """
    with zipfile.ZipFile(png_path, "r") as zf:
        fig_bytes = zf.read("figure.pkl")
    fig = pickle.loads(fig_bytes)

    # Some backends require explicit registration of the restored figure
    plt._backend_mod.new_figure_manager_given_figure(fig.number, fig)

    return fig


# 3. Restore the Figure object from PNG
fig_restored = extract_figure_from_png(png_path)

# 4. Modify the plot as if it was never lost
ax_restored = fig_restored.axes[0]

# Change line properties
lines = ax_restored.get_lines()
if lines:
    lines[0].set_color("red")
    lines[0].set_label("Modified line")

# Update labels and title
ax_restored.set_xlabel("X axis (modified)")
ax_restored.set_ylabel("Y axis (modified)")
ax_restored.set_title("Updated plot")
ax_restored.legend()

# Display updated figure
fig_restored.show()