import polars as pl
import pytest
from polars.testing import assert_frame_equal

from maplib import BlankNode

pl.Config.set_fmt_str_lengths(300)


@pytest.fixture(scope="function")
def pizzas_mapping():
    from maplib import (
        Mapping,
        Prefix,
        Template,
        Argument,
        Parameter,
        Variable,
        RDFType,
        Triple,
        a,
    )
    import polars as pl

    pl.Config.set_fmt_str_lengths(150)

    pi = "https://github.com/DataTreehouse/maplib/pizza#"
    df = pl.DataFrame(
        {
            "p": [pi + "Hawaiian", pi + "Grandiosa"],
            "c": [pi + "CAN", pi + "NOR"],
            "ings": [[pi + "Pineapple", pi + "Ham"], [pi + "Pepper", pi + "Meat"]],
        }
    )
    # print(df)

    pi = Prefix("pi", pi)

    p_var = Variable("p")
    c_var = Variable("c")
    ings_var = Variable("ings")

    template = Template(
        iri=pi.suf("PizzaTemplate"),
        parameters=[
            Parameter(variable=p_var, rdf_type=RDFType.IRI()),
            Parameter(variable=c_var, rdf_type=RDFType.IRI()),
            Parameter(variable=ings_var, rdf_type=RDFType.Nested(RDFType.IRI())),
        ],
        instances=[
            Triple(p_var, a(), pi.suf("Pizza")),
            Triple(p_var, pi.suf("fromCountry"), c_var),
            Triple(p_var, pi.suf("hasBlank"), BlankNode("MyBlank")),
            Triple(
                p_var,
                pi.suf("hasIngredient"),
                Argument(term=ings_var, list_expand=True),
                list_expander="cross",
            ),
        ],
    )

    m = Mapping()
    m.expand(template, df)
    hpizzas = """
    PREFIX pi:<https://github.com/DataTreehouse/maplib/pizza#>
    CONSTRUCT { ?p a pi:HeterodoxPizza } 
    WHERE {
        ?p a pi:Pizza .
        ?p pi:hasIngredient pi:Pineapple .
    }"""
    m.insert(hpizzas)
    return m


def test_simple_query_no_error(pizzas_mapping):
    res = pizzas_mapping.query(
        """
    PREFIX pi:<https://github.com/DataTreehouse/maplib/pizza#>

    SELECT ?p WHERE {
    ?p a pi:HeterodoxPizza
    }
    """
    )

    expected_df = pl.DataFrame(
        {"p": ["<https://github.com/DataTreehouse/maplib/pizza#Hawaiian>"]}
    )
    assert_frame_equal(res, expected_df)
