"""SATELLITE side. Image in, downlink bytes out. Runs with the flight bundle alone.

Imports satlink.sat and satlink.protocol. Nothing else -- no decoder, no VLM, no ground code. On
this machine that is 5.9M parameters and torch.

    python -m satlink.examples.sat_encode \
        --image scene.png --out codes.json --weights satlink/weights/sat_encoder.pt

The profile (frame size, tile size, tokens per crop, codebook width) is AGREED with the ground and
never transmitted, which is why the output is a bare stream with no header. Both ends construct
the same LinkProfile from their own configuration.

`--n-tok` is the rate knob. One-D-Piece is trained with tail-token drop, so tokens are ordered by
importance and any prefix is a valid encoding -- cutting it trades detail for bandwidth linearly,
at run time, with no retraining and no change on the ground beyond its own profile.
"""
import argparse
import sys
import time

import torch

from ..protocol import LinkProfile
from ..sat.encoder import SatEncoder


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", required=True)
    ap.add_argument("--out", required=True,
                    help="downlink file, always JSON. Writes the codes as text plus the packed "
                         "bitstream hex-encoded alongside them -- readable, diffable, and openable "
                         "in a text editor, which a raw .bin is not")
    ap.add_argument("--weights", default="satlink/weights/sat_encoder.pt")
    ap.add_argument("--n-tok", type=int, default=48)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"])

    a = ap.parse_args()

    from PIL import Image
    img = Image.open(a.image).convert("RGB")
    profile = LinkProfile(*img.size, n_tok=a.n_tok)
    sat = SatEncoder(a.weights, profile, device=a.device, dtype=getattr(torch, a.dtype))

    t0 = time.time()
    frame = sat.encode(img)          # structured form: codes + geometry
    wire = sat.pack(frame)           # packed form: what goes down the link
    dt = time.time() - t0

    import json
    import os
    # Both forms in one file: `codes` is the structured view another subsystem consumes, `wire_hex`
    # is the actual packed payload byte for byte. Keeping them together means a dump can be checked
    # against itself -- unpacking wire_hex must reproduce codes -- which is what the ground side
    # verifies before it will decode anything.
    json.dump({"image": os.path.basename(a.image),
               "width": profile.width, "height": profile.height, "tile": profile.tile,
               "n_tok": profile.n_tok, "bits": profile.bits, "dtype": a.dtype,
               "tile_origins": frame.tile_origins(), "codes": frame.tolist(),
               "wire_bytes": len(wire), "wire_hex": wire.hex()},
              open(a.out, "w"), indent=1)
    tx, ty = profile.grid
    raw = img.size[0] * img.size[1]  # bytes of raw grayscale this stands for
    print(f"image      {img.size[0]}x{img.size[1]}")
    print(f"tiling     {tx}x{ty} crops of 256x256, native resolution, "
          f"{256 - (profile.starts(0)[1] - profile.starts(0)[0]) if tx > 1 else 0} px overlap")
    print(f"model      {sum(p.numel() for p in sat.parameters())/1e6:.2f}M params ({a.dtype})")
    print(f"downlink   {len(wire)} bytes -> {a.out}   ({raw/len(wire):.0f}x vs raw grayscale)")
    print(f"encode     {1000*dt:.0f} ms")

    return 0


if __name__ == "__main__":
    sys.exit(main())