"""Python mirrors of the native sampling-input structs, for the compiled Neuron verification (issue #29).

Dynamo cannot read the nanobind structs' fields (each access breaks the graph) nor construct them, so the driver's
decode inputs enter the compiled ``Sampler._verification`` as these dataclasses: same field names, plain tensors.
A greedy batch carries no seeds or sample positions (``None``): only the RNG draw reads them, and the expansion of
the uint64 seeds is int64 arithmetic neuronx-cc graphs reject.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch


@dataclass
class LogitsMirror:
    temperatures: torch.Tensor
    top_ps: torch.Tensor
    top_ks: torch.Tensor


@dataclass
class KernelMirror:
    logits: LogitsMirror
    seq_seeds: torch.Tensor | None
    sample_positions: torch.Tensor | None


@dataclass
class SpeculativeMirror:
    mtp_info_read_rows: torch.Tensor
    mtp_info_write_slots: torch.Tensor
    nextn_tokens: torch.Tensor
    stop_token_ids: torch.Tensor


@dataclass
class DriverMirror:
    kernel: KernelMirror
    speculative: SpeculativeMirror | None
    golden_output: torch.Tensor | None
    max_top_logprobs: int


@dataclass
class DecodeMirror:
    mtp_first_token_positions: torch.Tensor
    mtp_seqlens_q: torch.Tensor


def mirror_kernel(kernel, greedy: bool) -> KernelMirror:  # noqa: ANN001
    logits = kernel.logits
    return KernelMirror(
        LogitsMirror(logits.temperatures, logits.top_ps, logits.top_ks),
        None if greedy else kernel.seq_seeds,
        None if greedy else kernel.sample_positions,
    )


def mirror_driver(driver, greedy: bool) -> DriverMirror:  # noqa: ANN001
    spec = driver.speculative
    return DriverMirror(
        kernel=mirror_kernel(driver.kernel, greedy),
        speculative=None
        if spec is None
        else SpeculativeMirror(spec.mtp_info_read_rows, spec.mtp_info_write_slots, spec.nextn_tokens, spec.stop_token_ids),
        golden_output=driver.golden_output,
        max_top_logprobs=driver.max_top_logprobs,
    )


def mirror_decode(decode) -> DecodeMirror:  # noqa: ANN001
    return DecodeMirror(decode.mtp_first_token_positions, decode.mtp_seqlens_q)