import pandas as pd
import pytest

from suntdataset.visualizer import SUNTVisualizer


nx = pytest.importorskip("networkx")
OD_URL = "https://github.com/LabIA-UFBA/SUNT/raw/main/data/od/od-2024-03-01.parquet"


def _sample_od_df():
    return pd.DataFrame(
        {
            "trip_id": ["t1", "t1", "t1", "t2", "t2"],
            "stop_id": ["A", "B", "C", "A", "B"],
            "pt_sequence": [1, 2, 3, 1, 2],
            "stop_time": [
                "2024-03-01 08:00:00",
                "2024-03-01 08:10:00",
                "2024-03-01 08:20:00",
                "2024-03-01 09:00:00",
                "2024-03-01 09:10:00",
            ],
            "n-boardings": [10, 20, 30, 5, 7],
        }
    )


@pytest.fixture(scope="module")
def real_od_df():
    return pd.read_parquet(OD_URL)


def test_build_od_graph_aggregates_edges():
    viz = SUNTVisualizer(_sample_od_df())
    graph = viz.build_od_graph()

    assert isinstance(graph, nx.DiGraph)
    assert graph.number_of_nodes() == 3
    assert graph.number_of_edges() == 2

    assert graph.has_edge("A", "B")
    assert graph["A"]["B"]["trips"] == 2.0
    assert graph["A"]["B"]["passengers"] == 15.0
    assert graph["A"]["B"]["weight"] == 15.0

    assert graph.has_edge("B", "C")
    assert graph["B"]["C"]["trips"] == 1.0
    assert graph["B"]["C"]["passengers"] == 20.0


def test_build_od_graph_works_with_time_order_only():
    df = _sample_od_df().drop(columns=["pt_sequence"])
    viz = SUNTVisualizer(df)
    graph = viz.build_od_graph(sequence_col=None, time_col="stop_time")

    assert graph.number_of_edges() == 2
    assert graph.has_edge("A", "B")
    assert graph.has_edge("B", "C")


def test_build_od_graph_missing_columns_raises():
    df = pd.DataFrame({"trip_id": ["t1"], "stop_id": ["A"]})
    viz = SUNTVisualizer(df)

    with pytest.raises(ValueError, match="Columns not found|Missing required columns|Missing"):
        viz.build_od_graph()


def test_export_od_graph_graphml_and_gexf(tmp_path):
    viz = SUNTVisualizer(_sample_od_df())

    graphml_path = viz.export_od_graph(tmp_path / "od.graphml", fmt="graphml")
    gexf_path = viz.export_od_graph(tmp_path / "od.gexf", fmt="gexf")

    assert graphml_path.exists()
    assert graphml_path.suffix == ".graphml"
    assert gexf_path.exists()
    assert gexf_path.suffix == ".gexf"


def test_export_od_graph_invalid_format_raises(tmp_path):
    viz = SUNTVisualizer(_sample_od_df())

    with pytest.raises(ValueError, match="Invalid Format|Invalid"):
        viz.export_od_graph(tmp_path / "od.invalid", fmt="json")


def test_build_od_graph_with_real_od_data(real_od_df):
    viz = SUNTVisualizer(real_od_df)
    graph = viz.build_od_graph(
        origin_col="stop_id",
        trip_col="trip_id",
        sequence_col="pt_sequence",
        time_col="stop_time",
        weight_col="n-boardings",
    )

    assert isinstance(graph, nx.DiGraph)
    assert graph.number_of_nodes() > 0
    assert graph.number_of_edges() > 0
    assert sum(data["passengers"] for _, _, data in graph.edges(data=True)) > 0


def test_plot_gis_lines_creates_polylines():
    folium = pytest.importorskip("folium")

    df = pd.DataFrame(
        {
            "shape_id": ["s1", "s1", "s1", "s2", "s2"],
            "shape_pt_lat": [-12.98, -12.97, -12.96, -13.01, -13.0],
            "shape_pt_lon": [-38.51, -38.5, -38.49, -38.52, -38.51],
            "shape_pt_sequence": [2, 1, 3, 2, 1],
        }
    )

    viz = SUNTVisualizer(df)
    fmap = viz.plot_gis_lines()

    assert isinstance(fmap, folium.Map)
    polyline_count = sum(
        1 for child in fmap._children.values() if isinstance(child, folium.vector_layers.PolyLine)
    )
    assert polyline_count == 2
