import pytest
from pytest_mock import MockerFixture
from pathlib import Path
from typing import Literal
from datasets import DatasetDict  # type: ignore

from artifex.models.base_model import BaseModel
from artifex.core import ValidationError

@pytest.mark.unit
@pytest.mark.parametrize(
    "user_instructions, output_path, num_samples",
    [
        ([1, 2, 3], "output/path", 100), # wrong type for user_instructions, should be a list[str]
        (["Inst A", "Inst B"], 1, 200), # wrong type for output_path, should be a str
        (["Inst A", "Inst B"], "single/instruction/output", "invalid"), # wrong type for num_samples, should be an int
    ]
)
def test_bm_build_tokenized_train_ds_validation_failure(    
    base_model: BaseModel,
    user_instructions: list[str],
    output_path: str,
    num_samples: int,
):
    """
    Test that the `_build_tokenized_train_ds` method of the BaseModel class raises a ValidationError when
    provided with invalid arguments.

    Args:
        base_model (BaseModel): An instance of the BaseModel class.
        user_instructions (list[str]): A list of user instructions.
        output_path (str): The output path where the dataset will be saved.
        num_samples (int): The number of samples to generate.
    """
    
    with pytest.raises(ValidationError):
        base_model._build_tokenized_train_ds( # type: ignore
            user_instructions=user_instructions,
            output_path=output_path,
            num_samples=num_samples
        )
        
@pytest.mark.unit
@pytest.mark.parametrize(
    "csv_content, base_model", 
    [
        ({"llm_output": "The sky is blue.", "labels": 1}, "llm_output")
    ],
    indirect=["base_model"]
)
def test_bm_build_tokenized_train_ds_success(
    mocker: MockerFixture,
    temp_synthetic_csv_file: Path,
    base_model: BaseModel,
    csv_content: dict[str, str | Literal[0, 1]],
):
    """
    Test that the `_build_tokenized_train_ds` method of the BaseModel class works correctly when
    provided with valid arguments.

    Args:
        mocker: A pytest fixture for mocking.
        temp_synthetic_csv_file (Path): A temporary file path for the mocked synthetic training data.
        base_model (BaseModel): An instance of the BaseModel class.
        csv_content (dict[str, str | Literal[0, 1]]): A dictionary representing the content of the CSV file.
        token_key (str): The key in the CSV content that will be tokenized.
    """
    
    mock_generate_data = mocker.patch("synthex.jobs_api.JobsAPI.generate_data")
    mocker.patch("artifex.models.base_model.BaseModel._await_data_generation")
    
    user_instr = ["user_instr_1", "user_instr_2"]
    
    output_path = str(temp_synthetic_csv_file.parent)
        
    out = base_model._build_tokenized_train_ds( # type: ignore
        user_instructions=user_instr,
        output_path=output_path
    )
    
    # Get the call arguments
    _, kwargs = mock_generate_data.call_args

    # Assert that the synthetic data generation job has received the correct instructions and schema
    assert set(kwargs["requirements"]) == set(user_instr + base_model._system_data_gen_instr) # type: ignore
    assert kwargs["schema_definition"] == base_model._synthetic_data_schema # type: ignore
    
    # Assert that the output has the correct type
    assert isinstance(out, DatasetDict)
    
    # Assert that the tokenized dataset contains the same keys as the non-tokenized dataset (other keys will also
    # be present, the ones generated by the tokenizer)
    assert set(csv_content.keys()).issubset(out["train"].column_names) # type: ignore

    # Check that each split contains tokenized fields (e.g., 'input_ids' or similar)
    for _, split_dataset in out.items(): # type: ignore
        # Ensure the split is not empty
        assert len(split_dataset) > 0 # type: ignore

        # Check for presence of tokenized fields in the first example; common tokenized fields are 
        # 'input_ids' and 'attention_mask'
        assert any(
            key in split_dataset[0] for key in ["input_ids", "attention_mask", "token_type_ids"]
        )
