import pandas as pd
import pytest

from flownarrator.formats.alpaca_format import to_alpaca_format

# ---------- HAPPY PATH TESTS -------------------------------------------------


def test_to_alpaca_format_basic_two_rows():
    df = pd.DataFrame(
        {
            "sentence": ["Hello", "Question"],
            "response": ["Answer1", "Answer2"],
        }
    )

    result = to_alpaca_format(df)

    assert isinstance(result, list)
    assert len(result) == 2

    assert result[0] == {
        "instruction": "Hello",
        "input": "",
        "output": "Answer1",
    }
    assert result[1] == {
        "instruction": "Question",
        "input": "",
        "output": "Answer2",
    }


def test_to_alpaca_format_with_custom_column_names():
    df = pd.DataFrame(
        {
            "prompt": ["P1"],
            "answer": ["A1"],
            "extra": ["ignored"],  # ensure extra columns are ignored
        }
    )

    result = to_alpaca_format(df, instruction_col="prompt", output_col="answer")

    assert result == [
        {
            "instruction": "P1",
            "input": "",
            "output": "A1",
        }
    ]


def test_to_alpaca_format_preserves_row_order():
    df = pd.DataFrame(
        {
            "sentence": [f"sent_{i}" for i in range(5)],
            "response": [f"resp_{i}" for i in range(5)],
        }
    )

    result = to_alpaca_format(df)

    assert [r["instruction"] for r in result] == [f"sent_{i}" for i in range(5)]
    assert [r["output"] for r in result] == [f"resp_{i}" for i in range(5)]


def test_to_alpaca_format_casts_non_string_values_to_string():
    df = pd.DataFrame(
        {
            "sentence": [123, None, 3.14],
            "response": [456, True, None],
        }
    )

    result = to_alpaca_format(df)

    # Compare using str() of the dataframe values to avoid pandas-type differences
    assert result[0]["instruction"] == str(df.loc[0, "sentence"])
    assert result[0]["output"] == str(df.loc[0, "response"])

    assert result[1]["instruction"] == str(df.loc[1, "sentence"])
    assert result[1]["output"] == str(df.loc[1, "response"])

    assert result[2]["instruction"] == str(df.loc[2, "sentence"])
    assert result[2]["output"] == str(df.loc[2, "response"])


# ---------- EDGE CASES -------------------------------------------------------


def test_to_alpaca_format_empty_dataframe():
    df = pd.DataFrame(columns=["sentence", "response"])

    result = to_alpaca_format(df)

    assert isinstance(result, list)
    assert result == []  # no rows -> empty list


def test_to_alpaca_format_single_row_dataframe():
    df = pd.DataFrame(
        {
            "sentence": ["Only one"],
            "response": ["Single response"],
        }
    )

    result = to_alpaca_format(df)

    assert len(result) == 1
    assert result[0] == {
        "instruction": "Only one",
        "input": "",
        "output": "Single response",
    }


def test_to_alpaca_format_with_nan_values():
    df = pd.DataFrame(
        {
            "sentence": ["Text", pd.NA],
            "response": [pd.NA, "Response"],
        }
    )

    result = to_alpaca_format(df)

    # pandas NA converts to 'NA' or '<NA>' depending on type; use str() for comparison
    assert result[0]["instruction"] == str(df.loc[0, "sentence"])
    assert result[0]["output"] == str(df.loc[0, "response"])
    assert result[1]["instruction"] == str(df.loc[1, "sentence"])
    assert result[1]["output"] == str(df.loc[1, "response"])


def test_to_alpaca_format_output_record_shape():
    """Each record must contain exactly instruction, input, output keys."""
    df = pd.DataFrame(
        {
            "sentence": ["Hello"],
            "response": ["World"],
            "another_col": ["ignored"],
        }
    )

    result = to_alpaca_format(df)
    record = result[0]

    assert set(record.keys()) == {"instruction", "input", "output"}
    assert record["input"] == ""  # input must always be empty string


def test_to_alpaca_format_does_not_modify_original_dataframe():
    df = pd.DataFrame(
        {
            "sentence": ["Hello"],
            "response": ["World"],
        }
    )
    df_copy = df.copy(deep=True)

    _ = to_alpaca_format(df)

    # Ensure DataFrame is unchanged
    pd.testing.assert_frame_equal(df, df_copy)


def test_to_alpaca_format_handles_very_long_strings():
    long_text = "x" * 10_000
    df = pd.DataFrame(
        {
            "sentence": [long_text],
            "response": [long_text],
        }
    )

    result = to_alpaca_format(df)

    assert len(result[0]["instruction"]) == len(long_text)
    assert len(result[0]["output"]) == len(long_text)


# ---------- ERROR CASES / VALIDATION -----------------------------------------


def test_to_alpaca_format_raises_type_error_for_non_dataframe():
    with pytest.raises(TypeError) as excinfo:
        to_alpaca_format("not a dataframe")  # type: ignore[arg-type]

    assert "'df' must be a pandas.DataFrame" in str(excinfo.value)


@pytest.mark.parametrize(
    "columns",
    [
        # missing response column
        ["sentence"],
        # missing sentence column
        ["response"],
        # missing both
        [],
        # completely different columns
        ["foo", "bar"],
    ],
)
def test_to_alpaca_format_raises_value_error_for_missing_required_columns(columns):
    df = pd.DataFrame({col: [] for col in columns})

    with pytest.raises(ValueError) as excinfo:
        to_alpaca_format(df)

    msg = str(excinfo.value)
    assert "DataFrame is missing required columns" in msg
    # Ensure message mentions at least one of the required columns
    assert ("sentence" in msg) or ("response" in msg)


def test_to_alpaca_format_raises_value_error_for_missing_custom_columns():
    df = pd.DataFrame({"sentence": ["ok"], "response": ["ok"]})

    with pytest.raises(ValueError) as excinfo:
        to_alpaca_format(df, instruction_col="prompt", output_col="answer")

    msg = str(excinfo.value)
    assert "DataFrame is missing required columns" in msg
    assert "prompt" in msg
    assert "answer" in msg
