from pathlib import Path

import pytest
import torch
from transformers import AutoTokenizer

from processing_salmonn import SalmonnProcessor


ROOT = Path(__file__).resolve().parents[1]


@pytest.fixture(scope="module")
def processor():
    tokenizer = AutoTokenizer.from_pretrained(ROOT, fix_mistral_regex=False)
    return SalmonnProcessor(tokenizer)


def test_basic_audio_request(processor):
    inputs = processor(
        audios=torch.zeros(1600),
        sampling_rate=16000,
        instruction="Describe the audio.",
    )

    assert inputs["input_ids"].shape[0] == 1
    assert inputs["audio_features"].shape == (1, 10, 128)
    assert inputs["audio_lengths"].tolist() == [10]
    assert inputs["audio_counts"].tolist() == [1]


def test_text_only_context_does_not_add_audio(processor):
    prompt = processor.build_prompt(
        "Transcribe the speech.",
        context=["howes", "wszelaki"],
    )

    assert prompt.count("<audio>") == 1
    assert "[howes, wszelaki]" in prompt


def test_audio_text_context_preserves_placeholder_order(processor):
    waveform = torch.zeros(1600)
    context = [
        {"text": "howes", "audio": (waveform, 16000)},
        {"text": "wszelaki", "audio": (waveform, 16000)},
    ]
    prompt = processor.build_prompt("Transcribe the speech.", context=context)
    inputs = processor(
        audios=(waveform, 16000),
        instruction="Transcribe the speech.",
        context=context,
    )

    assert prompt.count("<audio>") == 3
    assert "<audio>howes\n<audio>wszelaki" in prompt
    assert inputs["audio_features"].shape == (3, 10, 128)
    assert inputs["audio_lengths"].tolist() == [10, 10, 10]
    assert inputs["audio_counts"].tolist() == [3]


def test_mixed_context_modalities_are_rejected(processor):
    with pytest.raises(ValueError, match="all include audio or all be text-only"):
        processor.build_prompt(
            "Transcribe the speech.",
            context=[
                "howes",
                {"text": "wszelaki", "audio": (torch.zeros(1600), 16000)},
            ],
        )


def test_formatted_prompt_preserves_custom_audio_placement(processor):
    waveform = torch.zeros(1600)
    inputs = processor(
        audios=[(waveform, 16000), (waveform, 16000)],
        formatted_prompt="<audio>Compare this recording with <audio>this example.",
    )

    assert (inputs["input_ids"] == 151652).sum().item() == 2
    assert inputs["audio_counts"].tolist() == [2]


def test_formatted_prompt_validates_placeholder_count(processor):
    with pytest.raises(ValueError, match="1 audio placeholders but received 2"):
        processor(
            audios=[(torch.zeros(1600), 16000), (torch.zeros(1600), 16000)],
            formatted_prompt="<audio>Describe both recordings.",
        )


def test_instruction_and_formatted_prompt_are_mutually_exclusive(processor):
    with pytest.raises(ValueError, match="exactly one"):
        processor(
            audios=(torch.zeros(1600), 16000),
            instruction="Describe the audio.",
            formatted_prompt="<audio>Describe the audio.",
        )