"""Unit tests for the flat raw prompt top logprob response format
(`return_flat_raw_top_logprobs`).
"""

import asyncio
import json
import os
import time
import unittest

import numpy as np

from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel

maybe_stub_sgl_kernel()

from sglang.srt.managers.io_struct import GenerateReqInput
from sglang.srt.managers.tokenizer_manager import (
    ReqState,
    TokenizerManager,
    _build_flat_input_top_logprobs_fields,
)
from sglang.srt.observability.req_time_stats import APIServerReqTimeStats

register_cpu_ci(est_time=10, suite="base-a-test-cpu")

# Synthetic per-position top-k rows (k=2). The leading None mirrors the first
# prompt position, which has no top logprobs.
_VAL_ROWS = [None, [-0.1, -2.5], [-0.3, -1.5], [-0.05, -4.0]]
_IDX_ROWS = [None, [11, 22], [33, 44], [55, 66]]


class _TokenizerManagerStub:
    """Borrow the real logprob meta_info methods without a full manager."""

    add_logprob_to_meta_info = TokenizerManager.add_logprob_to_meta_info
    detokenize_logprob_tokens = TokenizerManager.detokenize_logprob_tokens
    detokenize_top_logprobs_tokens = TokenizerManager.detokenize_top_logprobs_tokens


def _make_state(**obj_kwargs) -> ReqState:
    obj = GenerateReqInput(text="hello", **obj_kwargs)
    obj.normalize_batch_and_arguments()
    return ReqState(
        out_list=[],
        finished=False,
        event=asyncio.Event(),
        obj=obj,
        time_stats=APIServerReqTimeStats(),
    )


def _add_logprob_meta_info(state: ReqState, top_logprobs_num: int = 2) -> dict:
    meta_info = {}
    _TokenizerManagerStub().add_logprob_to_meta_info(
        meta_info,
        state,
        top_logprobs_num=top_logprobs_num,
        token_ids_logprob=None,
        return_text_in_logprobs=False,
    )
    return meta_info


class TestFlatRawTopLogprobsValidation(CustomTestCase):
    def test_flag_defaults_off_and_valid(self):
        for kwargs in (
            {},
            {"return_flat_raw_top_logprobs": True},
        ):
            req = GenerateReqInput(text="hello", **kwargs)
            req.normalize_batch_and_arguments()

    def test_flat_rejects_multi_item_scoring(self):
        req = GenerateReqInput(
            text="a<sep>b",
            return_flat_raw_top_logprobs=True,
            multi_item_delimiter_indices=[1],
        )
        with self.assertRaisesRegex(ValueError, "multi-item"):
            req.normalize_batch_and_arguments()

    def test_flag_propagates_to_batch_items(self):
        req = GenerateReqInput(
            text=["a", "b"],
            return_flat_raw_top_logprobs=True,
        )
        req.normalize_batch_and_arguments()
        for i in range(2):
            self.assertTrue(req[i].return_flat_raw_top_logprobs)


class TestFlatAssembly(CustomTestCase):
    def test_flat_matches_nested_rows(self):
        fields = _build_flat_input_top_logprobs_fields(
            _VAL_ROWS, _IDX_ROWS, top_logprobs_num=2
        )
        self.assertEqual(fields["input_top_logprobs_shape"], [3, 2])
        self.assertEqual(fields["input_top_logprobs_null_prefix"], 1)
        self.assertEqual(
            fields["input_top_logprobs_val_flat"],
            [v for row in _VAL_ROWS[1:] for v in row],
        )
        self.assertEqual(
            fields["input_top_logprobs_idx_flat"],
            [i for row in _IDX_ROWS[1:] for i in row],
        )
        # Reconstruct the nested rows: covered position i, entry j lives at
        # flat[(i - null_prefix) * k + j].
        rows, k = fields["input_top_logprobs_shape"]
        null_prefix = fields["input_top_logprobs_null_prefix"]
        flat_val = fields["input_top_logprobs_val_flat"]
        self.assertEqual(len(flat_val), rows * k)
        for i in range(null_prefix, null_prefix + rows):
            start = (i - null_prefix) * k
            self.assertEqual(flat_val[start : start + k], _VAL_ROWS[i])

    def test_all_null_rows(self):
        fields = _build_flat_input_top_logprobs_fields(
            [None], [None], top_logprobs_num=2
        )
        self.assertEqual(fields["input_top_logprobs_shape"], [0, 2])
        self.assertEqual(fields["input_top_logprobs_null_prefix"], 1)
        self.assertEqual(fields["input_top_logprobs_val_flat"], [])
        self.assertEqual(fields["input_top_logprobs_idx_flat"], [])

    def test_rejects_null_row_after_prefix(self):
        with self.assertRaisesRegex(ValueError, "leading prefix"):
            _build_flat_input_top_logprobs_fields(
                [None, [-0.1, -2.5], None, [-0.3, -1.5]],
                [None, [11, 22], None, [33, 44]],
                top_logprobs_num=2,
            )

    def test_rejects_ragged_rows(self):
        with self.assertRaisesRegex(ValueError, "rectangular"):
            _build_flat_input_top_logprobs_fields(
                [None, [-0.1, -2.5], [-0.3]],
                [None, [11, 22], [33]],
                top_logprobs_num=2,
            )


class TestAddLogprobToMetaInfo(CustomTestCase):
    def _extend_input_top(self, state: ReqState, val_rows, idx_rows):
        state.input_top_logprobs_val.extend(val_rows)
        state.input_top_logprobs_idx.extend(idx_rows)

    def test_nested_path_unchanged_when_flags_unset(self):
        state = _make_state(return_logprob=True, top_logprobs_num=2)
        self._extend_input_top(state, _VAL_ROWS, _IDX_ROWS)
        meta_info = _add_logprob_meta_info(state)
        self.assertEqual(
            meta_info["input_top_logprobs"],
            [
                None if row is None else [(v, i, None) for v, i in zip(row, idx_row)]
                for row, idx_row in zip(_VAL_ROWS, _IDX_ROWS)
            ],
        )
        self.assertIn("output_top_logprobs", meta_info)
        for key in meta_info:
            self.assertNotIn("_flat", key)
        self.assertNotIn("input_top_logprobs_shape", meta_info)
        self.assertNotIn("input_top_logprobs_null_prefix", meta_info)

    def test_flat_path_replaces_nested_input(self):
        state = _make_state(
            return_logprob=True,
            top_logprobs_num=2,
            return_flat_raw_top_logprobs=True,
        )
        self._extend_input_top(state, _VAL_ROWS, _IDX_ROWS)
        meta_info = _add_logprob_meta_info(state)
        self.assertNotIn("input_top_logprobs", meta_info)
        self.assertIn("output_top_logprobs", meta_info)
        self.assertEqual(meta_info["input_top_logprobs_shape"], [3, 2])
        self.assertEqual(meta_info["input_top_logprobs_null_prefix"], 1)
        self.assertEqual(
            meta_info["input_top_logprobs_val_flat"],
            [v for row in _VAL_ROWS[1:] for v in row],
        )

    def test_unrepresentable_rows_fall_back_to_nested(self):
        # The shared batch-output loop must not raise on unrepresentable
        # rows; the request degrades to the nested format.
        state = _make_state(
            return_logprob=True,
            top_logprobs_num=2,
            return_flat_raw_top_logprobs=True,
        )
        val_rows = [_VAL_ROWS[1], None, _VAL_ROWS[2]]
        idx_rows = [_IDX_ROWS[1], None, _IDX_ROWS[2]]
        self._extend_input_top(state, val_rows, idx_rows)
        meta_info = _add_logprob_meta_info(state)
        self.assertIn("input_top_logprobs", meta_info)
        self.assertNotIn("input_top_logprobs_val_flat", meta_info)
        self.assertNotIn("input_top_logprobs_shape", meta_info)
        self.assertEqual(len(meta_info["input_top_logprobs"]), 3)

    def test_chunked_accumulation_matches_one_shot(self):
        # One shot.
        one_shot = _make_state(
            return_logprob=True,
            top_logprobs_num=2,
            return_flat_raw_top_logprobs=True,
        )
        self._extend_input_top(one_shot, _VAL_ROWS, _IDX_ROWS)
        expected = _add_logprob_meta_info(one_shot)

        # Rows arriving across two chunks, with meta_info assembled after each
        # (as happens for streaming requests).
        chunked = _make_state(
            return_logprob=True,
            top_logprobs_num=2,
            return_flat_raw_top_logprobs=True,
        )
        self._extend_input_top(chunked, _VAL_ROWS[:2], _IDX_ROWS[:2])
        _add_logprob_meta_info(chunked)
        self._extend_input_top(chunked, _VAL_ROWS[2:], _IDX_ROWS[2:])
        got = _add_logprob_meta_info(chunked)

        flat_keys = [
            "input_top_logprobs_val_flat",
            "input_top_logprobs_idx_flat",
            "input_top_logprobs_shape",
            "input_top_logprobs_null_prefix",
        ]
        for key in flat_keys:
            self.assertEqual(got[key], expected[key])

        # No new rows -> the encoded payload is reused, not rebuilt.
        again = _add_logprob_meta_info(chunked)
        self.assertIs(
            again["input_top_logprobs_val_flat"],
            got["input_top_logprobs_val_flat"],
        )


@unittest.skipUnless(
    os.environ.get("SGLANG_BENCH_FLAT_RAW_TOP_LOGPROBS"),
    "Serialization microbenchmark; set SGLANG_BENCH_FLAT_RAW_TOP_LOGPROBS=1 to run.",
)
class BenchFlatRawTopLogprobsSerialization(CustomTestCase):
    """Round-trip cost of the formats: server assembly + json.dumps, then
    client json.loads + reconstruction into [rows, k] arrays."""

    def test_bench(self):
        num_positions, k = 32768, 2
        rng = np.random.default_rng(0)
        vals = rng.standard_normal((num_positions, k)).astype(np.float32)
        idxs = rng.integers(0, 150000, size=(num_positions, k), dtype=np.int32)
        val_rows = [None] + vals[1:].tolist()
        idx_rows = [None] + idxs[1:].tolist()

        def best_of(fn, iters=5):
            result = fn()
            elapsed = min(
                (lambda s=time.perf_counter(): (fn(), time.perf_counter() - s)[1])()
                for _ in range(iters)
            )
            return elapsed * 1e3, result

        def bench(name, build, decode):
            encode_ms, payload = best_of(lambda: json.dumps(build()))
            decode_ms, arrays = best_of(lambda: decode(payload))
            self.assertEqual(arrays[0].shape, (num_positions - 1, k))
            print(
                f"{name}: encode {encode_ms:.1f} ms, decode {decode_ms:.1f} ms, "
                f"{len(payload)} bytes"
            )

        def decode_nested(payload):
            rows = [r for r in json.loads(payload) if r is not None]
            return (
                np.array([[e[0] for e in r] for r in rows], dtype=np.float32),
                np.array([[e[1] for e in r] for r in rows], dtype=np.int32),
            )

        def decode_flat(payload):
            d = json.loads(payload)
            shape = d["input_top_logprobs_shape"]
            return (
                np.asarray(d["input_top_logprobs_val_flat"], np.float32).reshape(shape),
                np.asarray(d["input_top_logprobs_idx_flat"], np.int32).reshape(shape),
            )

        bench(
            "nested triples",
            lambda: [
                (None if row is None else [(v, i, None) for v, i in zip(row, idx_row)])
                for row, idx_row in zip(val_rows, idx_rows)
            ],
            decode_nested,
        )
        bench(
            "flat lists",
            lambda: _build_flat_input_top_logprobs_fields(
                val_rows, idx_rows, top_logprobs_num=k
            ),
            decode_flat,
        )


if __name__ == "__main__":
    unittest.main(verbosity=2)
