import os
import pathlib
import random
import unittest
from unittest.mock import MagicMock

import pandas as pd
import pytest

from llm_studio.app_utils.default_datasets import (
    prepare_default_dataset_causal_language_modeling,
)
from llm_studio.src.datasets.conversation_chain_handler import ConversationChainHandler
from llm_studio.src.utils.data_utils import load_train_valid_data


@pytest.fixture
def cfg_mock():
    cfg = MagicMock()
    cfg.dataset.train_dataframe = "/path/to/train/data"
    cfg.dataset.validation_dataframe = "/path/to/validation/data"

    cfg.dataset.system_column = "None"
    cfg.dataset.prompt_column = "prompt"
    cfg.dataset.answer_column = "answer"

    cfg.dataset.validation_size = 0.2
    return cfg


@pytest.fixture
def read_dataframe_drop_missing_labels_mock(monkeypatch):
    data = {
        "prompt": [f"Prompt{i}" for i in range(100)],
        "answer": [f"Answer{i}" for i in range(100)],
        "id": list(range(100)),
    }
    df = pd.DataFrame(data)
    mock = MagicMock(return_value=df)
    monkeypatch.setattr(
        "llm_studio.src.utils.data_utils.read_dataframe_drop_missing_labels", mock
    )
    return mock


numbers = list(range(100))
random.shuffle(
    numbers,
)
groups = [numbers[n::13] for n in range(13)]


@pytest.fixture
def conversation_chain_ids_mock(monkeypatch):
    def mocked_init(self, *args, **kwargs):
        self.conversation_chain_ids = groups

    with unittest.mock.patch.object(
        ConversationChainHandler, "__init__", new=mocked_init
    ):
        yield


def test_get_data_custom_validation_strategy(
    cfg_mock, read_dataframe_drop_missing_labels_mock
):
    cfg_mock.dataset.validation_strategy = "custom"
    train_df, val_df = load_train_valid_data(cfg_mock)
    assert len(train_df), len(val_df) == 100


def test_get_data_automatic_split(
    cfg_mock, read_dataframe_drop_missing_labels_mock, conversation_chain_ids_mock
):
    cfg_mock.dataset.validation_strategy = "automatic"
    train_df, val_df = load_train_valid_data(cfg_mock)
    train_ids = set(train_df["id"].tolist())
    val_ids = set(val_df["id"].tolist())

    assert len(train_ids.intersection(val_ids)) == 0
    assert len(train_ids) + len(val_ids) == 100

    shared_groups = [
        i for i in groups if not train_ids.isdisjoint(i) and not val_ids.isdisjoint(i)
    ]
    assert len(shared_groups) == 0


def test_oasst_data_automatic_split(tmp_path: pathlib.Path):
    prepare_default_dataset_causal_language_modeling(tmp_path)
    assert len(os.listdir(tmp_path)) > 0, tmp_path
    cfg_mock = MagicMock()
    for file in os.listdir(tmp_path):
        if file.endswith(".pq"):
            cfg_mock.dataset.train_dataframe = os.path.join(tmp_path, file)

            cfg_mock.dataset.system_column = "None"
            cfg_mock.dataset.prompt_column = ("instruction",)
            cfg_mock.dataset.answer_column = "output"
            cfg_mock.dataset.parent_id_column = "parent_id"

            cfg_mock.dataset.validation_strategy = "automatic"

            for validation_size in [0.05, 0.1, 0.2, 0.3, 0.4, 0.5]:
                cfg_mock.dataset.validation_size = validation_size

                train_df, val_df = load_train_valid_data(cfg_mock)
                assert set(train_df["parent_id"].dropna().values).isdisjoint(
                    set(val_df["id"].dropna().values)
                )
                assert set(val_df["parent_id"].dropna().values).isdisjoint(
                    set(train_df["id"].dropna().values)
                )
                assert (len(val_df) / (len(train_df) + len(val_df))) == pytest.approx(
                    validation_size, 0.05
                )