import os
import struct
import tempfile
from typing import BinaryIO, Iterator

import gradio as gr

PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"


# ────────────────────────────────────────────────────────── helpers ──
def _read_chunks(fp: BinaryIO) -> Iterator[tuple[bytes, bytes]]:
    """
    Yield (raw_chunk_bytes, chunk_type_bytes) for every chunk in *fp*.
    `raw_chunk_bytes` includes length, type, data and CRC exactly as in file,
    so we can copy them verbatim later.
    The PNG signature must already be consumed by the caller.
    """
    while True:
        length_bytes = fp.read(4)
        if not length_bytes:
            break  # EOF
        if len(length_bytes) != 4:
            raise ValueError("Corrupted PNG: unexpected EOF while reading chunk length.")

        length = struct.unpack(">I", length_bytes)[0]          # ← fixed (no stray space)
        chunk_type = fp.read(4)
        data       = fp.read(length)
        crc        = fp.read(4)

        if len(chunk_type) != 4 or len(data) != length or len(crc) != 4:
            raise ValueError("Corrupted PNG: truncated chunk.")

        yield length_bytes + chunk_type + data + crc, chunk_type


def clone_metadata(src_path: str, tgt_path: str, out_path: str) -> str:
    """
    Copy all ancillary chunks (anything except IHDR / IDAT / IEND) from *src_path*
    into *tgt_path* immediately after its IHDR, then save to *out_path*.
    Pixel data of the target stays untouched.
    """
    with open(src_path, "rb") as fs, open(tgt_path, "rb") as ft:
        if fs.read(8) != PNG_SIGNATURE or ft.read(8) != PNG_SIGNATURE:
            raise ValueError("Both files must be valid PNG images.")

        src_chunks = list(_read_chunks(fs))
        tgt_chunks = list(_read_chunks(ft))

    # take every non-critical chunk from source
    ancillary = [
        raw for raw, tp in src_chunks
        if tp not in (b"IHDR", b"IDAT", b"IEND")
    ]

    out_chunks: list[bytes] = []
    for raw, tp in tgt_chunks:
        out_chunks.append(raw)
        if tp == b"IHDR":                     # inject right after IHDR
            out_chunks.extend(ancillary)

    with open(out_path, "wb") as fo:
        fo.write(PNG_SIGNATURE)
        for chunk in out_chunks:
            fo.write(chunk)

    return out_path


# ────────────────────────────────────────────────────────── gradio ui ──
def process(src_file, tgt_file):
    """Gradio wrapper → returns path to the merged PNG."""
    if not src_file or not tgt_file:
        raise gr.Error("Please upload both images.")

    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    tmp.close()  # we only need the path
    try:
        clone_metadata(src_file, tgt_file, tmp.name)
    except Exception as e:
        os.remove(tmp.name)
        raise gr.Error(str(e))

    return tmp.name


with gr.Blocks(title="Kovaaks PNG Metadata Cloner") as demo:
    gr.Markdown(
        """
        ## Kovaaks PNG Metadata Cloner 🏹

        **Step 1.** Upload a *metadata source* PNG that already works in-game  
        **Step 2.** Upload your *target* PNG (new 50 × 50 crosshair)  
        **Step 3.** Click **Clone metadata** and download the result.

        Put the output file into  
        `%LOCALAPPDATA%\\FPSAimTrainer\\Saved\\MyImport\\crosshairs`
        """
    )

    with gr.Row():
        src = gr.File(label="Image 1 – source (metadata)", type="filepath")
        tgt = gr.File(label="Image 2 – target (pixels)",   type="filepath")

    out = gr.File(label="Output PNG", interactive=False)
    btn = gr.Button("Clone metadata →")

    btn.click(fn=process, inputs=[src, tgt], outputs=[out])

    gr.Markdown(
        "✔️ Ancillary chunks (`sRGB`, `gAMA`, `pHYs`, `eXIf`, `tEXt`, …) are copied verbatim. "
        "Pixel data from the target image is left untouched."
    )

if __name__ == "__main__":
    demo.launch(show_api=False, share=False)