"""The public API. What an application outside this repository should import.

Three objects cover the whole pipeline:

    satlink.Satellite   pixels  -> codes        (encoder + VQ; 5.9M params)
    satlink.Ground      codes   -> pixels       (decoder; 337.8M params)
    satlink.GroundVLM   codes   -> answer       (decoder + vision tower + LLM, one module)

and one value type carries data between them:

    satlink.VQFrame     the codes plus the geometry that places them
    satlink.LinkProfile the frame/tile/token agreement -- shared constants, never transmitted

The split is real, not stylistic: `Satellite` imports nothing from `satlink.ground`, so the flight
bundle can be shipped without the decoder present on disk. Importing `satlink` itself pulls in
neither -- both are resolved on first attribute access (see satlink/__init__.py) so that a
satellite process never loads a 337.8M-parameter decoder just by touching the package.

Quick start -- satellite:

    import satlink
    sat = satlink.Satellite("satlink/weights/sat_encoder.pt", 512, 512)
    frame = sat.encode_file("scene.png")        # PNG or BMP, any size
    sat.save(frame, "codes.json")               # or: wire = sat.to_wire(frame)

Quick start -- ground:

    import satlink
    frame = satlink.load_codes("codes.json")
    satlink.Ground("satlink/weights/ground_decoder.pt").reconstruct(frame).save("recon.png")

See satlink/API.md for the full reference.
"""
import json
import os
import time

import numpy as np
import torch

from .device import max_memory_allocated, reset_peak_memory_stats, synchronize
from .frame import VQFrame
from .imageio import StripReader
from .protocol import LinkProfile, pack_frame, unpack_frame

__all__ = ["Satellite", "Ground", "GroundVLM", "VQFrame", "LinkProfile", "TOKENIZERS",
           "save_codes", "load_codes", "frame_from_wire", "read_image", "encode_one",
           "profile_for"]


# ---------------------------------------------------------------------------- tokenizer backends
# Two tokenizers are supported and they are NOT interchangeable at a given profile -- each pins its
# own tokens-per-crop and codebook width, so the profile is derived from the backend rather than
# passed in and hoped for.
#
#   opd    One-D-Piece      48 tokens x 12 bits =  72 B/crop, ~910x, ~21 dB
#   rstok  RS-Tokenizer     256 tokens x 18 bits = 576 B/crop, ~114x, ~41 dB
#
# 8x the downlink buys about 20 dB. Which one is right depends on whether the link or the imagery
# is the scarce resource, so both ship.
TOKENIZERS = {
    "opd":   {"n_tok": 48,  "bits": 12, "tile": 256},
    "rstok": {"n_tok": 256, "bits": 18, "tile": 256},
}


def profile_for(tokenizer, width, height, **kw):
    """-> a LinkProfile with the token geometry that `tokenizer` actually requires."""
    if tokenizer not in TOKENIZERS:
        raise ValueError(f"unknown tokenizer {tokenizer!r}; expected one of {sorted(TOKENIZERS)}")
    spec = dict(TOKENIZERS[tokenizer])
    spec.update(kw)
    return LinkProfile(width, height, **spec)


def _sniff(weights):
    """-> "opd" | "rstok", read off the checkpoint's own config.

    Sniffing rather than making the caller declare it: the two checkpoints are structurally
    distinct, and a mismatched declaration would fail deep inside a state-dict load with an
    unhelpful message.

    mmap=True because this reads a small dict out of a 1.35 GB file. Without it torch.load pulls
    every tensor storage into host RAM to answer a question about metadata, which is a 1.3 GB
    spike for a lookup -- and it happens in a constructor that otherwise allocates nothing."""
    import torch
    try:
        ck = torch.load(weights, map_location="cpu", weights_only=False, mmap=True)
    except (RuntimeError, TypeError, ValueError):
        # mmap needs the zipfile serialisation format and a supported torch; fall back rather
        # than refuse to load a checkpoint that is otherwise perfectly good.
        ck = torch.load(weights, map_location="cpu", weights_only=False)
    cfg = ck.get("config", {})
    if "token_grid" in cfg or "quantizer" in cfg:
        return "rstok"
    if "model" in cfg:
        return "opd"
    raise ValueError(f"cannot tell which tokenizer {weights} holds")


# ---------------------------------------------------------------------------- codes on disk
def save_codes(frame, path, **extra):
    """Write a VQFrame as JSON, atomically.

    JSON rather than a binary blob for anything that touches disk: it is the interchange format
    between the two halves during testing and integration, and a result file that a human cannot
    read is a result file nobody checks. The downlink itself is still packed bits -- see
    Satellite.to_wire -- this is the on-disk form, not the on-air form."""
    p = frame.profile
    obj = {"width": p.width, "height": p.height, "tile": p.tile, "n_tok": p.n_tok,
           "bits": p.bits, "codes": frame.tolist()}
    obj.update(extra)
    tmp = str(path) + ".tmp"
    with open(tmp, "w") as f:
        json.dump(obj, f)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)          # atomic: a reader sees the old file or the complete new one
    return path


def load_codes(path):
    """-> VQFrame, from a file written by save_codes."""
    import torch
    j = json.load(open(path))
    prof = LinkProfile(j["width"], j["height"], tile=j.get("tile", 256),
                       n_tok=j["n_tok"], bits=j.get("bits", 12))
    return VQFrame(torch.tensor(j["codes"], dtype=torch.long), prof)


def frame_from_wire(wire, width=None, height=None, n_tok=48, tile=256, bits=12, profile=None):
    """-> VQFrame, from packed downlink bytes.

    Nothing in the stream says how big the frame is or how it was tokenized: frame size, tile size,
    token count and codebook width are deployment constants both ends already agree on, so
    spending downlink bytes on a header that always says the same thing is pure overhead. The
    caller therefore supplies the geometry.

    Pass `profile=` when you already have one (from the encoder, or from load_codes) -- it carries
    all four constants together. The loose width/height form defaults to One-D-Piece's geometry,
    which is wrong for the RS-Tokenizer, and getting only n_tok right while leaving bits at 12 is
    exactly the mistake that produces a wrong-length frame."""
    import torch
    prof = profile or LinkProfile(width, height, tile=tile, n_tok=n_tok, bits=bits)
    return VQFrame(torch.tensor(unpack_frame(prof, wire), dtype=torch.long), prof)


def read_image(path, max_pixels=None):
    """-> (uint8 array in the file's native channels, width, height). PNG and BMP, any size.

    Convenience for callers that want the pixels rather than a reader object; large frames should
    use satlink.imageio.StripReader directly and stay streaming."""
    from .imageio import StripReader
    with StripReader(path, max_pixels) as r:
        return r.strip(0, r.height), r.width, r.height


# ------------------------------------------------------- streaming encode (satellite side)
def iter_crops(sat, reader):
    """Yield 256x256 grayscale crops, in profile raster order, one tile row at a time.

    Reads the source a strip at a time so peak host memory is one strip (32768x256 = 8 MB) rather
    than the whole 1 Gpixel plane. Crops are yielded individually and batched by the caller, so the
    batch size is a free parameter instead of being pinned to however many crops a row contains."""
    # Imported here rather than at module scope so that a ground-only process touching
    # satlink.Ground does not pull the satellite encoder into memory.
    from .sat.encoder import _LUMA
    p = sat.profile
    t = p.tile
    luma = torch.tensor(_LUMA, dtype=torch.float32)
    for y in p.starts(1):
        raw = torch.from_numpy(np.ascontiguousarray(reader.strip(y, t)))
        g = raw.float().div_(255.0)
        if g.ndim == 3:                                   # [h, W, 3] -> the encoder's own luma
            g = (g * luma).sum(-1)
        if g.shape[0] < t:                                # short frame: replicate, never zero-pad
            g = torch.nn.functional.pad(g[None, None], (0, 0, 0, t - g.shape[0]),
                                        mode="replicate")[0, 0]
        for x in p.starts(0):
            c = g[:, x:x + t]
            if c.shape[1] < t:
                c = torch.nn.functional.pad(c[None, None], (0, t - c.shape[1], 0, 0),
                                            mode="replicate")[0, 0]
            yield c
        del g, raw


def encode_one(sat, reader, batch):
    """-> (VQFrame, seconds, peak_bytes). `batch` crops are pushed through the encoder per call."""
    p = sat.profile
    reset_peak_memory_stats(sat.device)
    synchronize(sat.device)
    t0 = time.time()
    outs, buf = [], []
    for c in iter_crops(sat, reader):
        buf.append(c)
        if len(buf) == batch:
            outs.append(_flush(sat, buf))
            buf = []
    if buf:
        outs.append(_flush(sat, buf))
    synchronize(sat.device)
    secs = time.time() - t0
    codes = torch.cat(outs)
    if codes.shape[0] != p.n_tiles:
        raise RuntimeError(f"encoded {codes.shape[0]} crops, profile expects {p.n_tiles}")
    return VQFrame(codes, p), secs, max_memory_allocated(sat.device)


def _flush(sat, buf):
    """One encoder call over exactly len(buf) crops -> [B, n_tok] on the host."""
    b = torch.stack(buf)[:, None].expand(-1, 3, -1, -1).contiguous()
    return sat.indices_batch(b, chunk=len(buf)).cpu()


# ---------------------------------------------------------------------------- satellite
class Satellite:
    """Flight side. Pixels in, codes out.

        sat = satlink.Satellite("sat_encoder.pt", 32768, 30720, n_tok=48, dtype="float16")
        frame = sat.encode_file("scene.bmp", batch=64)
        wire  = sat.to_wire(frame)                       # 1.11 MB for a 32K x 30K frame

    `batch` is the number of 256x256 crops pushed through the encoder per call. It trades device
    memory for throughput and is yours to tune. Note that it is not numerically inert: different
    batch sizes select different GEMM kernels, whose rounding flips roughly 0.2-0.5% of codes that
    sat near a quantiser boundary. Mean reconstruction impact is under 0.4/255, but it does mean
    two runs at different batch sizes are not required to be bit-identical.
    """

    def __init__(self, weights, width, height, tokenizer="auto", n_tok=None, tile=None,
                 bits=None, device="auto", dtype="float16"):
        import torch
        self.tokenizer = _sniff(weights) if tokenizer == "auto" else tokenizer
        over = {k: v for k, v in (("n_tok", n_tok), ("tile", tile), ("bits", bits))
                if v is not None}
        self.profile = profile_for(self.tokenizer, width, height, **over)
        dt = getattr(torch, dtype) if isinstance(dtype, str) else dtype
        if self.tokenizer == "rstok":
            from .sat.rstok import RSTokEncoder as E
        else:
            from .sat.encoder import SatEncoder as E
        self.encoder = E(weights, self.profile, device=device, dtype=dt)
        self.device, self.dtype = self.encoder.device, dt

    @property
    def n_crops(self):
        return self.profile.n_tiles

    @property
    def frame_bytes(self):
        """Downlink size for one frame."""
        return self.profile.frame_bytes

    @property
    def compression_ratio(self):
        """vs the 8-bit grayscale source, the conservative denominator."""
        return self.profile.width * self.profile.height / self.profile.frame_bytes

    def encode(self, image, batch=64):
        """-> VQFrame. `image` is a PIL image, an HxW/HxWx3 array, or a CHW torch tensor."""
        return self.encoder.encode(image, chunk=batch)

    def encode_file(self, path, batch=64):
        """-> VQFrame, streaming a PNG or BMP from disk one tile row at a time.

        Peak host memory is one strip rather than the whole plane, so a 1-Gpixel frame does not
        need a 4 GB float buffer before the first crop is encoded."""
        with StripReader(path) as r:
            if (r.width, r.height) != (self.profile.width, self.profile.height):
                raise ValueError(f"{os.path.basename(path)} is {r.width}x{r.height}, this "
                                 f"Satellite is configured for "
                                 f"{self.profile.width}x{self.profile.height}")
            frame, _, _ = encode_one(self.encoder, r, batch)
        return frame

    def to_wire(self, frame):
        """-> packed downlink bytes, no header."""
        return pack_frame(frame.profile, frame.tolist())

    def save(self, frame, path, **extra):
        return save_codes(frame, path, **extra)


# ---------------------------------------------------------------------------- ground
class Ground:
    """Ground side, pixel reconstruction only. Codes in, image out.

    bfloat16 by default: measured identical to float32 to 0.01 dB on both backends, at half the
    memory. float16 is NOT the default because the RS-Tokenizer decoder overflows in it -- the
    whole frame comes back NaN, which casts to a black image rather than to an error.
    """

    def __init__(self, weights, tokenizer="auto", device="auto", dtype="bfloat16"):
        import torch
        from .device import resolve
        self.tokenizer = _sniff(weights) if tokenizer == "auto" else tokenizer
        dt = getattr(torch, dtype) if isinstance(dtype, str) else dtype
        self.device = resolve(device)
        if self.tokenizer == "rstok":
            from .ground.rstok_ground import RSTokDecoder as D
        else:
            from .ground.decoder import GroundDecoder as D
        self.decoder = D(weights, dtype=dt).to(self.device)

    def decode(self, frame):
        """-> float tensor [H, W] in [0, 1], overlaps averaged."""
        import torch
        with torch.no_grad():
            x = self.decoder(frame.to(self.device)).float()
        # GroundDecoder.stitch collapses the decoder's 3 channels to luma and then re-expands
        # them, so the [3, H, W] it returns holds three identical planes. Averaging rather than
        # taking x[0] keeps this correct either way: for identical channels it is exact, and if
        # the decoder ever produced genuine colour it would degrade to gray instead of silently
        # discarding two thirds of the output.
        while x.ndim > 3:
            x = x[0]
        x = x.mean(0) if x.ndim == 3 else x
        # A non-finite decode must be an error, not a picture. numpy casts NaN to 0 without
        # complaint, so an overflowed decoder silently produces a plausible-looking all-black PNG
        # that passes every check except a PSNR comparison the caller may not be running.
        if not torch.isfinite(x).all():
            raise RuntimeError(
                f"decoder produced non-finite output ({int((~torch.isfinite(x)).sum())} of "
                f"{x.numel()} values). The RS-Tokenizer decoder overflows in float16 -- "
                f"use dtype='float32' or 'bfloat16'.")
        return x

    def reconstruct(self, frame):
        """-> PIL grayscale image."""
        from PIL import Image
        x = self.decode(frame).clamp(0, 1).mul(255).round().to("cpu")
        return Image.fromarray(x.numpy().astype("uint8"), mode="L")


class GroundVLM:
    """Ground side, full stack. Codes in, answer out -- one module, no bolt-on adapter.

        vlm = satlink.GroundVLM("model/RSThinker", decoder="satlink/weights/ground_decoder.pt")
        print(vlm.ask(frame, "where is the ship"))

    Two paths, pick one:
      decoder=...        pixel path -- codes are decoded to pixels inside the model, then embedded
      latent_ground=...  latent path -- a partial decode is projected straight into the vision
                         tower's patch-embedding space; no pixels are ever produced

    On the pixel path the model is loaded LAZILY, on the first ask(), because its decoder is bound
    to a frame geometry that construction does not know yet. Pass `profile=` to load eagerly.
    """

    def __init__(self, model_dir, decoder=None, latent_ground=None, lora=None, device="auto",
                 profile=None):
        """`profile` controls WHEN the 10.6B model is loaded.

        The pixel path binds its decoder to a frame geometry, so without a profile there is nothing
        to build yet and construction returns immediately having touched no GPU at all -- the model
        appears on the device on the first ask(). That surprises people watching nvidia-smi, so
        pass `profile=LinkProfile(w, h)` (or satlink.profile_for(...)) to load eagerly instead, and
        pay the ~15 s and 21 GB where you expect them.

        The latent path has no such binding and always loads here.
        """
        from .device import resolve
        # A merged checkpoint carries its own path marker, so neither argument is needed -- and
        # supplying one would be ambiguous about which decoder actually gets used.
        sl = os.path.join(model_dir, "satlink.json")
        self.merged = json.load(open(sl)) if os.path.isfile(sl) else None
        if self.merged:
            if decoder or latent_ground or lora:
                raise ValueError(f"{model_dir} is a merged {self.merged['path']} checkpoint; "
                                 f"drop decoder=/latent_ground=/lora=")
        elif bool(decoder) == bool(latent_ground):
            raise ValueError("give exactly one of decoder= or latent_ground=")
        # Only the pixel path supports the RS-Tokenizer: the latent path projects a PARTIAL decode
        # into the vision tower, through a bridge trained against One-D-Piece's particular
        # intermediate features. Pointing it at another tokenizer would run, and would be wrong.
        if self.merged:
            dc = self.merged.get("decoder_config") or self.merged.get("bridge_config") or {}
            self.tokenizer = "rstok" if ("token_grid" in dc or "quantizer" in dc) else "opd"
        else:
            self.tokenizer = _sniff(decoder) if decoder else "opd"
        self.device = resolve(device)
        self.model_dir = model_dir
        self.latent = bool(latent_ground) or (self.merged or {}).get("path") == "latent"
        if self.merged and not self.latent:
            # Merged pixel path: the decoder is inside the shards, so there is no profile to bind
            # and the model can be built right here.
            from .ground.modeling_vqvlm import VQGlm4vForConditionalGeneration as M
            self.model = M.from_pretrained(model_dir, profile=profile, device=self.device)
            self._tail = "<|assistant|>\n"
        elif self.latent:
            from .ground.modeling_latentvlm import LatentGlm4vForConditionalGeneration as M
            self.model = M.from_pretrained(model_dir, latent_ground, lora=lora,
                                           device=self.device)
            self._tail = "<|assistant|>"
        else:
            from .ground.modeling_vqvlm import VQGlm4vForConditionalGeneration as M
            self._M, self._decoder_w = M, decoder
            self.model = None                 # profile-bound; built here only if one was given
            if profile is not None:
                self._ensure(profile)
            self._tail = "<|assistant|>\n"
        self._tok = None

    def _tokenizer(self):
        if self._tok is None:
            from tokenizers import Tokenizer
            self._tok = Tokenizer.from_file(os.path.join(self.model_dir, "tokenizer.json"))
        return self._tok

    def _ensure(self, profile):
        """The pixel path binds the decoder to a profile, so the model is built once we see one."""
        if self.model is None:
            self.model = self._M.from_pretrained(self.model_dir, self._decoder_w, profile,
                                                 device=self.device)
        return self.model

    def ask(self, frame, question="describe the image", max_new=256):
        """-> the model's answer as a string."""
        import torch
        from .ground.infer import greedy
        from .ground.modeling_vqvlm import mrope_positions
        model = self._ensure(frame.profile)
        p = frame.profile
        # Both model classes expose these; passing the frame's own profile rather than the
        # deployment default is what lets one loaded model answer about frames of several sizes,
        # which is exactly what a benchmark or an acceptance sweep does.
        n_img, (gh, gw) = model.n_image_tokens(p), model.image_grid(p)
        tok, IMG = self._tokenizer(), 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|>{question}{self._tail}", add_special_tokens=False).ids,
            device=self.device)
        # The two model classes spell the spatial-merge factor differently (pixel path:
        # merge_size, latent path: merge), so read whichever is present rather than assuming.
        merge = getattr(model, "merge_size", None) or model.merge
        pos = mrope_positions(ids, IMG, gh, gw, merge).to(self.device)
        with model.reference_attention():
            out = greedy(model, ids, pos, frame.to(self.device), max_new=max_new)
        return tok.decode(out, skip_special_tokens=True)