from sqllineage.utils.entities import ColumnQualifierTuple
from ...helpers import assert_column_lineage_equal


def test_select_column_from_table_join():
    sql = """INSERT INTO tab1
SELECT tab2.col1,
       tab3.col2
FROM tab2
         INNER JOIN tab3
                    ON tab2.id = tab3.id"""
    assert_column_lineage_equal(
        sql,
        [
            (
                ColumnQualifierTuple("col1", "tab2"),
                ColumnQualifierTuple("col1", "tab1"),
            ),
            (
                ColumnQualifierTuple("col2", "tab3"),
                ColumnQualifierTuple("col2", "tab1"),
            ),
        ],
    )
    sql = """INSERT INTO tab1
SELECT tab2.col1 AS col3,
       tab3.col2 AS col4
FROM tab2
         INNER JOIN tab3
                    ON tab2.id = tab3.id"""
    assert_column_lineage_equal(
        sql,
        [
            (
                ColumnQualifierTuple("col1", "tab2"),
                ColumnQualifierTuple("col3", "tab1"),
            ),
            (
                ColumnQualifierTuple("col2", "tab3"),
                ColumnQualifierTuple("col4", "tab1"),
            ),
        ],
    )
    sql = """INSERT INTO tab1
SELECT a.col1 AS col3,
       b.col2 AS col4
FROM tab2 a
         INNER JOIN tab3 b
                    ON a.id = b.id"""
    assert_column_lineage_equal(
        sql,
        [
            (
                ColumnQualifierTuple("col1", "tab2"),
                ColumnQualifierTuple("col3", "tab1"),
            ),
            (
                ColumnQualifierTuple("col2", "tab3"),
                ColumnQualifierTuple("col4", "tab1"),
            ),
        ],
    )


def test_select_column_without_table_qualifier_from_table_join():
    sql = """INSERT INTO tab3
SELECT f1
FROM ( SELECT f1 FROM tab1)
LEFT JOIN ( SELECT f1 FROM tab2) USING (f1)"""
    assert_column_lineage_equal(
        sql,
        [
            (
                ColumnQualifierTuple("f1", "tab1"),
                ColumnQualifierTuple("f1", "tab3"),
            ),
            (
                ColumnQualifierTuple("f1", "tab2"),
                ColumnQualifierTuple("f1", "tab3"),
            ),
        ],
    )


def test_select_column_from_same_table_multiple_time_using_different_alias():
    sql = """INSERT INTO tab1
SELECT a.col1 AS col2,
       b.col1 AS col3
FROM tab2 a
         JOIN tab2 b
              ON a.parent_id = b.id"""
    assert_column_lineage_equal(
        sql,
        [
            (
                ColumnQualifierTuple("col1", "tab2"),
                ColumnQualifierTuple("col2", "tab1"),
            ),
            (
                ColumnQualifierTuple("col1", "tab2"),
                ColumnQualifierTuple("col3", "tab1"),
            ),
        ],
    )


def test_column_reference_with_ansi89_join():
    sql = """INSERT INTO tab3
SELECT a.id,
       a.name AS name1,
       b.name AS name2
FROM (SELECT id, name
      FROM tab1) a,
     (SELECT id, name
      FROM tab2) b
WHERE a.id = b.id"""
    assert_column_lineage_equal(
        sql,
        [
            (ColumnQualifierTuple("id", "tab1"), ColumnQualifierTuple("id", "tab3")),
            (
                ColumnQualifierTuple("name", "tab1"),
                ColumnQualifierTuple("name1", "tab3"),
            ),
            (
                ColumnQualifierTuple("name", "tab2"),
                ColumnQualifierTuple("name2", "tab3"),
            ),
        ],
    )


def test_column_lineage_multiple_paths_for_same_column():
    sql = """INSERT INTO tab2
SELECT tab1.id,
       coalesce(join_table_1.col1, join_table_2.col1, join_table_3.col1) AS col1
FROM tab1
         LEFT JOIN (SELECT id, col1 FROM tab1 WHERE flag = 1) AS join_table_1
                   ON tab1.id = join_table_1.id
         LEFT JOIN (SELECT id, col1 FROM tab1 WHERE flag = 2) AS join_table_2
                   ON tab1.id = join_table_2.id
         LEFT JOIN (SELECT id, col1 FROM tab1 WHERE flag = 3) AS join_table_3
                   ON tab1.id = join_table_3.id"""
    assert_column_lineage_equal(
        sql,
        [
            (
                ColumnQualifierTuple("id", "tab1"),
                ColumnQualifierTuple("id", "tab2"),
            ),
            (
                ColumnQualifierTuple("col1", "tab1"),
                ColumnQualifierTuple("col1", "tab2"),
            ),
        ],
    )
