"""
Functions for testing the parsing functionality of the ftio package.
"""

import os

import pandas as pd

from ftio.freq.helper import get_mode
from ftio.parse.extract import extract_fields, get_time_behavior_and_args
from ftio.parse.scales import Scales


def test_scales():
    """
    Test the Scales class initialization.
    """
    file = os.path.join(
        os.path.dirname(__file__),
        "../examples/tmio/ior/collective/1536_new.json",
    )
    args = [
        "ftio",
        file,
        "-e",
        "no",
    ]
    data = Scales(args)

    assert isinstance(data, Scales)
    assert data.args is not None
    assert data.n >= 1
    assert len(data.s) >= 1


def test_assign_data():  # type: ignore
    """
    Test the assign_data method of the Scales class.
    """
    file = os.path.join(
        os.path.dirname(__file__),
        "../examples/tmio/ior/collective/1536_new.json",
    )
    args = [
        "ftio",
        file,
        "-e",
        "no",
    ]
    data = Scales(args)
    # assign the different fields in data (read/write sync/async and io time)
    data.assign_data()
    # extract mode of interest
    mode = get_mode(data, data.args.mode)

    assert hasattr(data, "df_wst")
    assert hasattr(data, "df_wat")
    assert hasattr(data, "df_rst")
    assert hasattr(data, "df_rat")
    assert hasattr(data, "df_time")

    assert mode is not None
    assert isinstance(mode, tuple)


def test_get_io_mode():
    """
    Test the get_io_mode method of the Scales class.
    """
    file = os.path.join(
        os.path.dirname(__file__),
        "../examples/tmio/ior/collective/1536_new.json",
    )
    args = [
        "ftio",
        file,
        "-e",
        "no",
    ]
    data = Scales(args)
    # assign the different fields in data (read/write sync/async and io time)
    args = data.args
    df = data.get_io_mode(args.mode)

    assert isinstance(df, tuple)
    assert len(df) == 5
    assert all(isinstance(d, pd.DataFrame) for d in df)


def test_get_time_behavior_and_args():
    """
    Test the get_time_behavior_and_args function.
    """
    file = os.path.join(
        os.path.dirname(__file__),
        "../examples/tmio/ior/collective/1536_new.json",
    )
    args = [
        "ftio",
        file,
        "-e",
        "no",
    ]
    data, args = get_time_behavior_and_args(args)

    assert isinstance(data, list)
    assert len(data) > 0
    assert "bandwidth" in data[0]
    assert "time" in data[0]
    assert args is not None


def test_extract_fields():
    """
    Test the extract_fields function.
    """
    file = os.path.join(
        os.path.dirname(__file__),
        "../examples/tmio/ior/collective/1536_new.json",
    )
    args = [
        "ftio",
        file,
        "-e",
        "no",
    ]
    data, args = get_time_behavior_and_args(args)
    b_sampled, time_b, ranks, total_bytes = extract_fields(data)

    assert len(b_sampled) > 0
    assert len(time_b) > 0
    assert isinstance(ranks, int)
    assert isinstance(total_bytes, int)
