"""Tests for multimodal splice (T2.3).

Verifies the splice [属性][时间戳][TS tok][事件][问题][回答] produces:
- inputs_embeds [1, seq, hidden] with seq_len <= max_len
- attention_mask [1, seq] all ones (unpadded single sample)
- labels [1, seq] with -100 everywhere EXCEPT the answer segment

Uses the real Qwen2.5-0.5B tokenizer (local) + a stub embedding layer so the
splice *logic* is validated without loading the full LLM weights.
"""
import pytest
import torch
from transformers import AutoTokenizer

from tsmm.model.multimodal import MultimodalSplicer

TOKENIZER_PATH = "/home/zhangzp/models/Qwen2.5-0.5B-Instruct"
HIDDEN = 896


@pytest.fixture(scope="module")
def tokenizer():
    return AutoTokenizer.from_pretrained(TOKENIZER_PATH)


@pytest.fixture(scope="module")
def stub_embed(tokenizer):
    # mimic LLM get_input_embeddings(): Embedding(vocab, hidden).
    # Use len(tokenizer) which includes added special tokens (eos=151645 etc.);
    # tokenizer.vocab_size (151643) excludes them and is too small.
    return torch.nn.Embedding(len(tokenizer), HIDDEN)


@pytest.fixture
def splicer(tokenizer, stub_embed):
    s = MultimodalSplicer(tokenizer, hidden_size=HIDDEN, max_len=1024)
    s.bind(stub_embed)
    return s


def make_ts_embeds(n_patches=127, hidden=HIDDEN):
    return torch.randn(n_patches, hidden)


class TestSplice:
    def test_shapes(self, splicer):
        out = splicer.splice(
            attributes="变量：CPU利用率 单位：% 采样率：1Hz",
            timestamps="时间范围：2026-01-01 00:00 ~ 00:08",
            ts_embeds=make_ts_embeds(127),
            events="",
            question="描述这段时序的整体趋势",
            answer="整体呈上升趋势。",
        )
        assert out.inputs_embeds.dim() == 3
        assert out.inputs_embeds.shape[0] == 1
        assert out.inputs_embeds.shape[2] == HIDDEN
        seq = out.inputs_embeds.shape[1]
        assert out.attention_mask.shape == (1, seq)
        assert out.labels.shape == (1, seq)
        # within context budget
        assert seq <= 1024

    def test_attention_mask_all_ones_unpadded(self, splicer):
        out = splicer.splice(
            attributes="CPU", timestamps="t0..t1", ts_embeds=make_ts_embeds(10),
            events="", question="Q?", answer="A.",
        )
        assert torch.all(out.attention_mask == 1)

    def test_labels_only_answer_non_masked(self, splicer, tokenizer):
        out = splicer.splice(
            attributes="CPU", timestamps="t0", ts_embeds=make_ts_embeds(10),
            events="", question="Q?", answer="上升趋势",
        )
        labels = out.labels[0]
        # there must be at least one non -100 position (the answer)
        assert (labels != -100).sum() > 0
        # all non-answer positions must be -100
        n_non_masked = int((labels != -100).sum())
        # the answer token count should match tokenizing answer (+ eos)
        ans_ids = tokenizer("上升趋势", add_special_tokens=False)["input_ids"]
        # +1 for appended eos
        assert n_non_masked == len(ans_ids) + 1

    def test_no_answer_means_all_masked(self, splicer):
        # inference mode: no answer → all labels -100
        out = splicer.splice(
            attributes="CPU", timestamps="t0", ts_embeds=make_ts_embeds(10),
            events="", question="Q?", answer=None,
        )
        assert torch.all(out.labels == -100)

    def test_ts_tokens_in_sequence(self, splicer):
        # the TS token block length should appear in the total seq length
        n_patches = 50
        out_no_ts = splicer.splice(
            attributes="CPU", timestamps="t0", ts_embeds=torch.zeros(0, HIDDEN),
            events="", question="Q?", answer="A.",
        )
        out_with_ts = splicer.splice(
            attributes="CPU", timestamps="t0", ts_embeds=make_ts_embeds(n_patches),
            events="", question="Q?", answer="A.",
        )
        delta = out_with_ts.inputs_embeds.shape[1] - out_no_ts.inputs_embeds.shape[1]
        assert delta == n_patches

    def test_seq_len_within_budget(self, splicer):
        out = splicer.splice(
            attributes="CPU 内存 磁盘 网络 温度 请求量",
            timestamps="2026-01-01 00:00:00 至 2026-01-01 00:08:32 UTC",
            ts_embeds=make_ts_embeds(127),
            events="t=300 处发生部署事件：服务重启",
            question="请判断这段时序是否存在异常，并以 JSON 给出异常区间。",
            answer='{"segments": [[300, 350]], "reason": "部署后延迟突增"}',
        )
        assert out.inputs_embeds.shape[1] <= 1024

    def test_no_nan(self, splicer):
        out = splicer.splice(
            attributes="CPU", timestamps="t0", ts_embeds=make_ts_embeds(20),
            events="", question="Q?", answer="A.",
        )
        assert not torch.isnan(out.inputs_embeds).any()

    def test_truncation_when_overlong(self, splicer):
        # force an overlong sequence; splicer should truncate to max_len
        splicer.max_len = 40
        out = splicer.splice(
            attributes="CPU", timestamps="t0",
            ts_embeds=make_ts_embeds(127),
            events="", question="Q?", answer="A.",
        )
        assert out.inputs_embeds.shape[1] <= 40
        assert out.attention_mask.shape == (1, out.inputs_embeds.shape[1])
        assert out.labels.shape == (1, out.inputs_embeds.shape[1])