"""GROUND side. Downlink bytes in, answer out. Never sees an image, never imports satellite code.

Imports satlink.ground / protocol / frame. The satellite bundle does not need to exist on this
machine; a capture file is enough, which is also how a replay or a regression test works.

Pixel path -- the codes are decoded to pixels INSIDE the model, then embedded:

    python -m satlink.examples.ground_answer \
        --wire codes.json \
        --model model/RSThinker --decoder satlink/weights/ground_decoder.pt \
        --question "where is the ship"

Latent path -- no pixels are ever produced; a partial decode is projected straight into the
vision tower's patch-embedding space:

    python -m satlink.examples.ground_answer \
        --wire codes.json --model model/RSThinker \
        --latent-ground satlink/weights/latent_ground_dec12.pt \
        --lora outputs/latent_sft_dec12/checkpoint-2241 \
        --question "where is the ship"

`--width/--height` are the agreed profile, not metadata read off the stream: nothing in the
downlink says how big the frame is, because in flight that is a deployment constant.
"""
import argparse
import os
import sys

import torch

from ..frame import VQFrame
from ..protocol import LinkProfile, unpack_frame


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--wire", required=True)
    ap.add_argument("--width", type=int, default=0, help="agreed frame width (not needed for .json)")
    ap.add_argument("--height", type=int, default=0)
    ap.add_argument("--n-tok", type=int, default=48)
    ap.add_argument("--model", required=True, help="GLM-4.1V / RSThinker directory")
    ap.add_argument("--decoder", default=None, help="pixel path: ground_decoder.pt")
    ap.add_argument("--latent-ground", default=None, help="latent path: latent_ground*.pt")
    ap.add_argument("--lora", default=None, help="latent path: the SFT adapter directory")
    ap.add_argument("--question", default="describe the image")
    ap.add_argument("--max-new", type=int, default=256)
    ap.add_argument("--device", default="cuda")
    a = ap.parse_args()
    if bool(a.decoder) == bool(a.latent_ground):
        ap.error("give exactly one of --decoder or --latent-ground")

    from tokenizers import Tokenizer
    from ..ground.infer import greedy
    from ..ground.modeling_vqvlm import mrope_positions

    # Both forms are accepted. JSON is the readable one -- what the tests and other subsystems
    # consume; the packed stream is what a radio actually carries. Either way the LINK layer ends
    # HERE: the model takes a VQFrame, so it stays callable from a test or a training loop with no
    # bitstream in sight.
    if a.wire.endswith(".json"):
        import json
        j = json.load(open(a.wire))
        profile = LinkProfile(j["width"], j["height"], n_tok=j["n_tok"])
        codes = j["codes"]
        print(f"downlink   {profile.frame_bytes} bytes equivalent ({a.wire}, {j.get('dtype','?')}), "
              f"{profile.n_tiles} crops x {profile.n_tok} tokens")
    else:
        profile = LinkProfile(a.width, a.height, n_tok=a.n_tok)
        wire = open(a.wire, "rb").read()
        codes = unpack_frame(profile, wire)
        print(f"downlink   {len(wire)} bytes, {profile.n_tiles} crops x {profile.n_tok} tokens")
    frame = VQFrame(torch.tensor(codes, dtype=torch.long), profile)

    if a.latent_ground:
        from ..ground.modeling_latentvlm import LatentGlm4vForConditionalGeneration as M
        model = M.from_pretrained(a.model, a.latent_ground, lora=a.lora, device=a.device)
        n_img, (gh, gw) = model.n_image_tokens(profile), model.image_grid(profile)
        tok_params = sum(p.numel() for p in model.bridge.parameters())
        tail, path = "<|assistant|>", f"latent ({len(model.bridge.decoder.transformer)}/24 layers)"
    else:
        from ..ground.modeling_vqvlm import VQGlm4vForConditionalGeneration as M
        model = M.from_pretrained(a.model, a.decoder, profile, device=a.device)
        n_img = model.n_image_tokens()
        f = model.patch_size * model.merge
        gh = max(f, round(profile.height / f) * f) // model.patch_size
        gw = max(f, round(profile.width / f) * f) // model.patch_size
        tok_params = sum(p.numel() for p in model.opd.parameters())
        tail, path = "<|assistant|>\n", "pixel (full decode)"
    print(f"ground     {path}: tokenizer {tok_params/1e6:.1f}M, then vision tower + LLM, one module")

    tok = Tokenizer.from_file(os.path.join(a.model, "tokenizer.json"))
    IMG = model.image_token_id
    ids = torch.tensor(
        tok.encode("[gMASK]<sop><|user|>\n<|begin_of_image|>", add_special_tokens=False).ids
        + [IMG] * n_img
        + tok.encode(f"<|end_of_image|>{a.question}{tail}", add_special_tokens=False).ids,
        device=a.device)
    pos = mrope_positions(ids, IMG, gh, gw, model.merge).to(a.device)
    print(f"prompt     {ids.numel()} tokens ({n_img} image, patch grid {gh}x{gw})")

    with model.reference_attention():
        out = greedy(model, ids, pos, frame.to(a.device), max_new=a.max_new)
    print(f"\nQ: {a.question}\nA: {tok.decode(out, skip_special_tokens=False)}")
    return 0


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