import datetime

import pandas as pd
import pytest

from metricflow.dataflow.sql_table import SqlTable
from metricflow.protocols.sql_client import SqlClient
from metricflow.sql_clients.sql_utils import make_df
from metricflow.test.fixtures.setup_fixtures import MetricFlowTestSessionState
from metricflow.time.time_constants import ISO8601_PYTHON_FORMAT
from metricflow.time.time_granularity import TimeGranularity

DEFAULT_DS = "ds"


def create_table(sql_client: SqlClient, sql_table: SqlTable, df: pd.DataFrame) -> None:  # noqa: D
    # TODO: Replace with table_exists() check.
    sql_client.drop_table(sql_table)

    sql_client.create_table_from_dataframe(
        sql_table=sql_table,
        df=df,
    )


@pytest.fixture(scope="session")
def create_simple_model_tables(mf_test_session_state: MetricFlowTestSessionState, sql_client: SqlClient) -> bool:
    """Creates tables with example source data for testing."""
    schema = mf_test_session_state.mf_source_schema

    cote_divoire = "cote d'ivoire"

    # NOTE: fct_bookings should have both a host and guest id, but we do not yet
    # have "entity roles" so that a user dimension can be attached to either
    # guest or host (e.g. bookings by user:guest/country)
    data = [
        # While datetime columns can be used here, somewhere there is a conversion to a string when testing
        # with sqlite, so when you read the datetime column, it comes out as a string.
        # Columns for reference
        # ["guest_id", "host_id", "listing_id", "booking_value", "is_instant", "ds", "ds_partitioned",
        #  "booking_paid_at", "referrer_id"]
        ("u0004114", "u0003141", "l3141592", 742.42, False, "2020-01-01", "2020-01-01", "2020-01-02", None),
        ("u0004114", "u0003141", "l3141592", 441.14, True, "2020-01-01", "2020-01-01", "2020-01-02", "u0003141"),
        ("u0003141", "u0003154", "l2718281", 499.99, True, "2020-01-01", "2020-01-01", "2020-01-02", None),
        ("u0003154", "u0003141", "l2718281", 719.89, False, "2020-01-01", "2020-01-01", "2020-01-02", None),
        ("u0004114", "u0003141", "l3141592", 319.85, True, "2020-01-01", "2020-01-01", "2020-01-02", "u0005432"),
        ("u0004114", "u0003141", "l2718281", 442.42, True, "2020-01-02", "2020-01-02", "2020-01-03", None),
        ("u0004114", "u0003141", "l2718281", 241.14, True, "2020-01-02", "2020-01-02", "2020-01-03", None),
        ("u0004114", "u0003141", "l3141592", 799.99, False, "2020-01-02", "2020-01-02", "2020-01-03", None),
        ("u0005432", "u0003452", "l2718281", 319.85, True, "2020-01-02", "2020-01-02", "2020-01-03", "u0003452"),
        ("u0004114", "u1004114", "l9658588-incomplete", 0.0, False, "2020-01-02", "2020-01-02", "2020-01-03", None),
        ("u0004114", "u1004114", "l8912456-incomplete", 0.0, False, "2020-01-02", "2020-01-02", "2020-01-03", None),
        ("u0003452", "u0005432", "l3141592", 519.89, False, "2020-01-02", "2020-01-02", "2020-01-03", "u0003141"),
        ("u0003452", "u0004114", "l5948301", 332.23, False, "2020-01-02", "2020-01-02", "2020-01-03", "u0003141"),
        ("u1003452", "u1004114", "no_such_listing", 0.0, False, "2020-01-02", "2020-01-02", "2020-01-03", None),
        ("u0003452", "u0004114", "l5948301", 0.0, False, "2020-01-03", "2020-01-03", "2020-01-04", None),
        ("u0004114", "u0003141", "l3141592", 797.42, False, "2019-12-18", "2019-12-18", "2019-12-19", None),
        ("u0004114", "u0003141", "l3141592", 808.14, True, "2019-12-18", "2019-12-18", "2019-12-19", "u0003141"),
        ("u0003141", "u0003154", "l2718281", 531.99, True, "2019-12-18", "2019-12-18", "2019-12-19", None),
        ("u0003154", "u0003141", "l2718281", 285.89, False, "2019-12-18", "2019-12-18", "2019-12-19", None),
        ("u0004114", "u0003141", "l3141592", 714.85, True, "2019-12-18", "2019-12-18", "2019-12-19", "u0005432"),
        ("u0004114", "u0003141", "l2718281", 368.42, True, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0004114", "u0003141", "l2718281", 521.14, True, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0004114", "u0003141", "l3141592", 581.99, False, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0005432", "u0003452", "l2718281", 753.85, True, "2019-12-19", "2019-12-19", "2019-12-20", "u0003452"),
        ("u0004114", "u1004114", "l9658588-incomplete", 0.0, False, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0004114", "u1004114", "l8912456-incomplete", 0.0, False, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0003452", "u0005432", "l3141592", 754.89, False, "2019-12-19", "2019-12-19", "2019-12-20", "u0003141"),
        ("u0003452", "u0004114", "l5948301", 152.23, False, "2019-12-19", "2019-12-19", "2019-12-20", "u0003141"),
        ("u1003452", "u1004114", "no_such_listing", 0.0, False, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0003452", "u0004114", "l5948301", 0.0, False, "2019-12-20", "2019-12-20", "2019-12-21", None),
        ("u0004114", "u0003141", "l3141592", 352.42, False, "2019-12-18", "2019-12-18", "2019-12-19", None),
        ("u0004114", "u0003141", "l3141592", 444.14, True, "2019-12-18", "2019-12-18", "2019-12-19", "u0003141"),
        ("u0003141", "u0003154", "l2718281", 415.99, True, "2019-12-18", "2019-12-18", "2019-12-19", None),
        ("u0003154", "u0003141", "l2718281", 578.89, False, "2019-12-18", "2019-12-18", "2019-12-19", None),
        ("u0004114", "u0003141", "l3141592", 768.85, True, "2019-12-18", "2019-12-18", "2019-12-19", "u0005432"),
        ("u0004114", "u0003141", "l2718281", 532.42, True, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0004114", "u0003141", "l2718281", 258.14, True, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0004114", "u0003141", "l3141592", 753.99, False, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0005432", "u0003452", "l2718281", 852.85, True, "2019-12-19", "2019-12-19", "2019-12-20", "u0003452"),
        ("u0004114", "u1004114", "l9658588-incomplete", 0.0, False, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0004114", "u1004114", "l8912456-incomplete", 0.0, False, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0003452", "u0005432", "l3141592", 728.89, False, "2019-12-19", "2019-12-19", "2019-12-20", "u0003141"),
        ("u0003452", "u0004114", "l5948301", 951.23, False, "2019-12-19", "2019-12-19", "2019-12-20", "u0003141"),
        ("u1003452", "u1004114", "no_such_listing", 0.0, False, "2019-12-19", "2019-12-19", "2019-12-20", None),
        ("u0003452", "u0004114", "l5948301", 0.0, False, "2019-12-20", "2019-12-20", "2019-12-21", None),
        ("u0003452", "u0004114", "l5948301", 951.23, False, "2019-12-01", "2019-12-01", "2019-12-01", "u0003141"),
    ]
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="fct_bookings"),
        df=make_df(
            sql_client=sql_client,
            columns=[
                "guest_id",
                "host_id",
                "listing_id",
                "booking_value",
                "is_instant",
                DEFAULT_DS,
                "ds_partitioned",
                "booking_paid_at",
                "referrer_id",
            ],
            time_columns={DEFAULT_DS, "ds_partitioned", "booking_paid_at"},
            data=data,
        ),
    )

    # Like fct_bookings, but using dt instead of ds to verify that a primary time dimension named something other than
    # ds can be used.
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="fct_bookings_dt"),
        df=make_df(
            sql_client=sql_client,
            columns=[
                "guest_id",
                "host_id",
                "listing_id",
                "booking_value",
                "is_instant",
                "dt",
                "dt_partitioned",
                "booking_paid_at",
                "referrer_id",
            ],
            time_columns={"dt", "dt_partitioned"},
            data=data,
        ),
    )

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="fct_views"),
        df=make_df(
            sql_client=sql_client,
            columns=["user_id", "listing_id", DEFAULT_DS, "ds_partitioned"],
            time_columns={DEFAULT_DS, "ds_partitioned"},
            data=[
                ("u0004114", "l3141592", "2020-01-01", "2020-01-01"),
                ("u1612112", "l3141592", "2020-01-01", "2020-01-01"),
                ("u0004114", "l2718281", "2020-01-02", "2020-01-02"),
                ("u0004114", "l3141592", "2020-01-02", "2020-01-02"),
                ("u1612112", "l2718281", "2020-01-02", "2020-01-02"),
                ("u0004114", "", "2020-01-02", "2020-01-02"),
                ("u0004114", "l7891283-incomplete", "2020-01-02", "2020-01-02"),
            ],
        ),
    )

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="dim_listings_latest"),
        df=make_df(
            sql_client=sql_client,
            columns=["listing_id", "country", "capacity", "is_lux", "user_id", "created_at"],
            time_columns={"created_at"},
            data=[
                ("l3141592", "us", 3, True, "u0004114", "2020-01-01"),
                ("l5948301", "us", 5, True, "u0004114", "2020-01-02"),
                ("l2718281", cote_divoire, 4, False, "u0005432", "2020-01-02"),
                ("l9658588-incomplete", "us", None, None, "u1004114", "2020-01-02"),
                ("l8912456-incomplete", None, None, None, "u1004114", "2020-01-02"),
                ("l7891283-incomplete", "ca", None, False, "u1004114", "2020-01-02"),
            ],
        ),
    )

    # The dim_listings table is an SCD type II table in the style of a dbt snapshot with two columns setting the
    # the range of time when this row was "active", with a NULL value for "active_to" meaning it's still current
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="dim_listings"),
        df=make_df(
            sql_client=sql_client,
            columns=["listing_id", "country", "capacity", "is_lux", "user_id", "active_from", "active_to"],
            time_columns={"active_from", "active_to"},
            data=[
                ("l3141592", "us", 2, True, "u0004114", "2020-01-01", "2020-01-02"),
                ("l2718281", cote_divoire, 4, True, "u0005432", "2020-01-01", "2020-01-02"),
                ("l3141592", "us", 3, True, "u0004114", "2020-01-02", None),
                # This cote_divoire property changed hands to a person from Maryland who considers it not lux
                ("l2718281", cote_divoire, 4, False, "u0003154", "2020-01-02", None),
                ("l5948301", "us", 5, True, "u0004114", "2020-01-02", None),
                ("l9658588-incomplete", None, None, None, "u1004114", "2020-01-01", "2020-01-02"),
                ("l9658588-incomplete", "us", None, None, "u1004114", "2020-01-02", None),
                ("l8912456-incomplete", None, None, None, "u1004114", "2020-01-02", None),
                ("l7891283-incomplete", "ca", None, False, "u1004114", "2020-01-02", None),
            ],
        ),
    )
    create_table(
        sql_client=sql_client,  # successful identity verifications
        sql_table=SqlTable(schema_name=schema, table_name="fct_id_verifications"),
        df=make_df(
            sql_client=sql_client,
            columns=["verification_id", "user_id", "verification_type", DEFAULT_DS, "ds_partitioned"],
            time_columns={DEFAULT_DS, "ds_partitioned"},
            data=[
                ("v0000001", "u0004114", "drivers_license", "2020-01-01", "2020-01-01"),
                ("v0000002", "u1612112", "passport", "2020-01-01", "2020-01-01"),
                ("v0000003", "u0004114", "social_security", "2020-01-02", "2020-01-02"),
                ("v0000004", "u1612112", "drivers_license", "2020-01-02", "2020-01-02"),
            ],
        ),
    )

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="dim_users"),
        df=make_df(
            sql_client=sql_client,
            columns=["user_id", "home_state", "created_at", DEFAULT_DS, "ds_partitioned"],
            time_columns={"created_at", DEFAULT_DS, "ds_partitioned"},
            data=[
                ("u0004114", "CA", "2019-03-03", "2020-01-01", "2020-01-01"),
                ("u1612112", "CA", "2019-04-03", "2020-01-01", "2020-01-01"),
                ("u0005432", "TX", "2017-03-03", "2020-01-01", "2020-01-01"),
                ("u0003452", "HI", "2013-03-03", "2020-01-01", "2020-01-01"),
                ("u0003154", "MD", "2014-03-03", "2020-01-01", "2020-01-01"),
                ("u0003141", "NY", "2015-03-03", "2020-01-01", "2020-01-01"),
                ("u0004114", "CA", "2019-03-03", "2020-01-02", "2020-01-02"),
                ("u1612112", "WA", "2019-04-03", "2020-01-02", "2020-01-02"),
                ("u0005432", "TX", "2017-03-03", "2020-01-02", "2020-01-02"),
                ("u0003452", "HI", "2013-03-03", "2020-01-02", "2020-01-02"),
                ("u0003154", "MD", "2014-03-03", "2020-01-02", "2020-01-02"),
                ("u0003141", "NY", "2015-03-03", "2020-01-02", "2020-01-02"),
            ],
        ),
    )

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="dim_users_latest"),
        df=make_df(
            sql_client=sql_client,
            columns=["user_id", "home_state_latest", DEFAULT_DS],
            time_columns={DEFAULT_DS},
            data=[
                ("u0004114", "CA", "2020-01-02"),
                ("u1612112", "WA", "2020-01-02"),
                ("u0005432", "TX", "2020-01-02"),
                ("u0003452", "HI", "2020-01-02"),
                ("u0003154", "MD", "2020-01-02"),
                ("u0003141", "NY", "2020-01-02"),
            ],
        ),
    )

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="fct_revenue"),
        df=make_df(
            sql_client=sql_client,
            columns=["revenue", "created_at", "user_id"],
            time_columns={"created_at"},
            data=[
                (1000, "2020-01-01", "u0004114"),
                (4000, "2020-02-01", "u0004114"),
                (3000, "2020-03-01", "u0005432"),
                (4000, "2020-04-02", "u0004114"),
                (1000, "2021-01-01", "u0004114"),
                (1000, "2021-01-02", "u0004114"),
                (1000, "2021-01-03", "u0004114"),
                (1000, "2021-01-04", "u0004114"),
            ],
        ),
    )

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="dim_lux_listing_id_mapping"),
        df=make_df(
            sql_client=sql_client,
            columns=["listing_id", "lux_listing_id"],
            data=[
                ("l3141592", "ll_001"),
                ("l5948301", "LUX_TEST_ID"),
                ("l2718281", "ll_002"),
            ],
        ),
    )

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="dim_lux_listings"),
        df=make_df(
            sql_client=sql_client,
            columns=["lux_listing_id", "is_confirmed_lux", "valid_from", "valid_to"],
            time_columns={"valid_from", "valid_to"},
            data=[
                ("ll_001", True, "2020-01-01", None),
                ("LUX_TEST_ID", True, "2020-01-02", None),
                ("ll_002", True, "2020-01-01", "2020-01-02"),
                ("ll_002", False, "2020-01-02", None),
            ],
        ),
    )

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="dim_companies"),
        df=make_df(
            sql_client=sql_client,
            columns=["company_id", "user_id", "company_name"],
            data=[
                ("cpid_0", "u0003154", "MD Vacation Rentals LLC"),
            ],
        ),
    )

    accounts_data = [
        ("u0004114", 2135, "2020-01-01", "checking"),
        ("u0004114", 1234, "2020-01-01", "savings"),
        ("u0005432", 5234, "2020-01-03", "checking"),
        ("u0005432", 234, "2020-01-03", "savings"),
        ("u0003141", 24634, "2020-01-01", "checking"),
        ("u0003154", 23452, "2020-01-01", "checking"),
        ("u0004114", 1213, "2020-01-06", "savings"),
        ("u1612112", 5123, "2020-01-02", "checking"),
        ("u0004114", 523, "2020-01-07", "checking"),
        ("u0004114", 7434, "2020-01-10", "checkings"),
        ("u0005432", 8456, "2020-01-02", "savings"),
        ("u0003141", 12939, "2020-01-12", "checking"),
        ("u0003452", 6939, "2020-01-02", "checking"),
        ("u0003452", 35, "2020-01-12", "checking"),
        ("u0005414", 5582, "2020-01-04", "savings"),
    ]

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="fct_accounts"),
        df=make_df(
            sql_client=sql_client,
            columns=["user_id", "account_balance", DEFAULT_DS, "account_type"],
            time_columns={DEFAULT_DS},
            data=accounts_data,
        ),
    )

    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="dim_primary_accounts"),
        df=make_df(
            sql_client=sql_client,
            columns=["user_id", "account_type", "set_as_primary", "removed_as_primary"],
            time_columns={"set_as_primary", "removed_as_primary"},
            data=[
                ("u0004114", "savings", "2020-01-01", "2020-01-02"),
                ("u1612112", "checking", "2020-01-01", "2020-01-01"),
                ("u0005432", "savings", "2020-01-02", "2020-01-03"),
                ("u0003452", "checking", "2020-01-01", None),
                ("u0003154", "checking", "2020-01-01", None),
                ("u0003141", "checking", "2020-01-01", None),
                ("u0004114", "savings", "2020-01-02", None),
                ("u1612112", "checking", "2020-01-02", "2020-01-02"),
                ("u0005432", "checking", "2020-01-03", None),
            ],
        ),
    )

    return True


@pytest.fixture(scope="session")
def create_scd_model_tables(create_simple_model_tables: bool) -> bool:
    """This fixture ensures the scd tables are created when necessary

    The SCD model relies on the simple model tables, so we simply depend on that fixture and return it here.
    If, at any point, we decide to diverge from the simple model we can update this fixture independently.
    """
    return create_simple_model_tables


@pytest.fixture(scope="session")
def create_bridge_table(mf_test_session_state: MetricFlowTestSessionState, sql_client: SqlClient) -> bool:
    """Creates tables with example source data for multi-hop join testing."""
    schema = mf_test_session_state.mf_source_schema

    # NOTE: fct_bookings should have both a host and guest id, but we do not yet
    # have "entity roles" so that a user dimension can be attached to either
    # guest or host (e.g. bookings by user:guest/country)
    customer_accounts = [
        # ["customer_id", "customer_name", "customer_atomic_weight", "ds_partitioned"]
        ("0", "thorium", "90", "2020-01-01"),
        ("1", "osmium", "76", "2020-01-02"),
        ("2", "tellurium", "52", "2020-01-01"),
        ("3", "einsteinium", "99", "2020-01-04"),
    ]
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="customer_table"),
        df=make_df(
            sql_client=sql_client,
            columns=["customer_id", "customer_name", "customer_atomic_weight", "ds_partitioned"],
            time_columns={"ds_partitioned"},
            data=customer_accounts,
        ),
    )

    customer_other_data = [
        # ["customer_id", "country", "customer_third_hop_id"]
        ("0", "turkmenistan", "another_id0"),
        ("1", "paraguay", "another_id1"),
        ("2", "myanmar", "another_id2"),
        ("3", "djibouti", "another_id3"),
    ]
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="customer_other_data"),
        df=make_df(
            sql_client=sql_client,
            columns=["customer_id", "country", "customer_third_hop_id"],
            data=customer_other_data,
        ),
    )

    third_hop_table = [
        # ["customer_third_hop_id", "value"]
        ("another_id0", "citadel"),
        ("another_id1", "virtu"),
        ("another_id2", "two sigma"),
        ("another_id3", "jump"),
    ]
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="third_hop_table"),
        df=make_df(
            sql_client=sql_client,
            columns=["customer_third_hop_id", "value"],
            data=third_hop_table,
        ),
    )

    bridge_data = [
        # ["account_id", "customer_id", "extra_dim", "ds_partitioned"]
        ("a0", "0", "lux", "2020-01-01"),
        ("a0", "1", "lux", "2020-01-02"),
        ("a1", "2", "not_lux", "2020-01-01"),
        ("a2", "3", "super_lux", "2020-01-04"),
    ]
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="bridge_table"),
        df=make_df(
            sql_client=sql_client,
            columns=["account_id", "customer_id", "extra_dim", "ds_partitioned"],
            time_columns={"ds_partitioned"},
            data=bridge_data,
        ),
    )

    account_month_txns = [
        # While datetime columns can be used here, somewhere there is a conversion to a string when testing
        # with sqlite, so when you read the datetime column, it comes out as a string.
        # Columns for reference
        # ["account_id", DEFAULT_DS, "account_month", "txn_count", "txns_value", "ds_partitioned"]
        ("a0", "2020-01-01", "JANUARY", 3, 300, "2020-01-01"),
        ("a0", "2020-02-01", "FEBRUARY", 4, 400, "2020-02-01"),
        ("a0", "2020-03-01", "MARCH", 5, 1000, "2020-03-01"),
        ("a0", "2020-04-01", "APRIL", 4, 600, "2020-04-01"),
        ("a1", "2020-01-01", "JANUARY", 2, 200, "2020-01-01"),
        ("a1", "2020-02-01", "FEBRUARY", 5, 250, "2020-02-01"),
        ("a1", "2020-03-01", "MARCH", 6, 360, "2020-03-01"),
        ("a1", "2020-04-01", "APRIL", 10, 1000, "2020-04-01"),
        ("a2", "2020-01-01", "JANUARY", 5, 500, "2020-01-01"),
    ]
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="account_month_txns"),
        df=make_df(
            sql_client=sql_client,
            columns=["account_id", DEFAULT_DS, "account_month", "txn_count", "txns_value", "ds_partitioned"],
            time_columns={DEFAULT_DS, "ds_partitioned"},
            data=account_month_txns,
        ),
    )
    return True


@pytest.fixture(scope="session")
def create_message_source_tables(mf_test_session_state: MetricFlowTestSessionState, sql_client: SqlClient) -> bool:
    """Creates tables with example source data for testing."""
    schema = mf_test_session_state.mf_source_schema

    data = [
        ("t0", "u0", "m0", "2020-01-01"),
        ("t0", "u0", "m1", "2020-01-01"),
        ("t0", "u1", "m2", "2020-01-01"),
        ("t0", "u1", "m3", "2020-01-01"),
        ("t0", "u1", "m4", "2020-01-01"),
        ("t1", "u0", "m5", "2020-01-02"),
        ("t1", "u0", "m7", "2020-01-02"),
        ("t1", "u0", "m8", "2020-01-02"),
        ("t1", "u1", "m9", "2020-01-02"),
        ("t1", "u1", "m10", "2020-01-02"),
        ("t1", "u1", "m11", "2020-01-02"),
        ("t1", "u1", "m12", "2020-01-03"),
        ("t1", "u2", "m13", "2020-01-03"),
        ("t1", "u3", "m14", "2020-01-03"),
    ]
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="fct_messages"),
        df=make_df(
            sql_client=sql_client,
            columns=["team_id", "user_id", "message_id", DEFAULT_DS],
            time_columns={DEFAULT_DS},
            data=data,
        ),
    )
    user_data = [
        ("t0", "u0", "burkina faso", "2020-01-01", "u0_2"),
        ("t0", "u1", "vanuatu", "2020-01-01", "u1_2"),
        ("t1", "u0", "merica", "2020-01-01", "u0_2"),
        ("t1", "u1", "merica", "2020-01-01", "u1_2"),
        ("t1", "u2", "burkina faso", "2020-01-01", "u2_2"),
        ("t1", "u3", "vanuata", "2020-01-01", "u3_2"),
    ]
    # Like fct_bookings, but using dt instead of ds to verify that a primary time dimension named something other than
    # ds can be used.
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="fct_users"),
        df=make_df(
            sql_client=sql_client,
            columns=["team_id", "id", "country", "created_at", "ident_2"],
            time_columns={"created_at"},
            data=user_data,
        ),
    )

    users_more_data = [
        ("u0", "thorium", "u0_2"),
        ("u1", "tellurium", "u1_2"),
        ("u0", "osmium", "u0_2"),
        ("u1", "tecnetium", "u1_2"),
        ("u2", "einsteinium", "u2_2"),
        ("u3", "rubidium", "u3_2"),
    ]
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="fct_users_more"),
        df=pd.DataFrame(
            columns=["id", "user_element", "ident_2"],
            data=users_more_data,
        ),
    )
    return True


@pytest.fixture(scope="session")
def create_extended_date_model_tables(mf_test_session_state: MetricFlowTestSessionState, sql_client: SqlClient) -> bool:
    """Creates fct_bookings_extended / dim_listings_extended with a year's worth of dates for granularity testing."""
    schema = mf_test_session_state.mf_source_schema

    # Create fct_bookings_extended that has rows for every date in 2020.
    # Also create dim_listings_extended to match.
    start_date = datetime.datetime(2020, 1, 1)
    end_date = datetime.datetime(2020, 12, 31)

    current_date = start_date
    fct_bookings_extended_data = []
    dim_listings_extended_data = []
    booking_id = 0
    listing_id = 0
    while current_date <= end_date:
        ds_as_str = current_date.strftime(ISO8601_PYTHON_FORMAT)
        fct_bookings_extended_data.append((booking_id, 1, listing_id, booking_id % 2 == 0, ds_as_str))
        dim_listings_extended_data.append((listing_id, ds_as_str))

        listing_id += 1
        # Have 10 different listing IDs so that we get different distinct values when we set the granularity
        # to daily / weekly / monthly
        listing_id = listing_id % 10
        current_date = current_date + datetime.timedelta(days=1)

        booking_id += 1

    sql_table = SqlTable(schema_name=schema, table_name="fct_bookings_extended")
    sql_client.drop_table(sql_table)

    sql_client.create_table_from_dataframe(
        sql_table=sql_table,
        df=make_df(
            sql_client=sql_client,
            columns=["booking_id", "booking", "listing_id", "is_instant", DEFAULT_DS],
            time_columns={DEFAULT_DS},
            data=fct_bookings_extended_data,
        ),
    )

    sql_table = SqlTable(schema_name=schema, table_name="dim_listings_extended")
    sql_client.drop_table(sql_table)

    sql_client.create_table_from_dataframe(
        sql_table=sql_table,
        df=make_df(
            sql_client=sql_client,
            columns=["listing_id", "listing_creation_ds"],
            time_columns={"listing_creation_ds"},
            data=dim_listings_extended_data,
        ),
    )

    # Create fct_bookings_extended_monthly with data at a monthly granularity.
    # Example rows:
    # ["bookings_monthly", "is_instant", "ds"]
    # 10                   True          2020-01-01
    # 5                    False         2020-01-01
    # 10                   True          2020-02-01
    # 5                    False         2020-02-01
    current_date = pd.Timestamp(start_date)
    fct_bookings_extended_monthly_data = []
    while current_date <= end_date:
        fct_bookings_extended_monthly_data.append((10, True, current_date.strftime(ISO8601_PYTHON_FORMAT)))
        fct_bookings_extended_monthly_data.append((5, False, current_date.strftime(ISO8601_PYTHON_FORMAT)))

        current_date += TimeGranularity.MONTH.period_begin_offset

    sql_table = SqlTable(schema_name=schema, table_name="fct_bookings_extended_monthly")
    sql_client.drop_table(sql_table=sql_table)
    sql_client.create_table_from_dataframe(
        sql_table=sql_table,
        df=make_df(
            sql_client=sql_client,
            columns=["bookings_monthly", "is_instant", DEFAULT_DS],
            time_columns={DEFAULT_DS},
            data=fct_bookings_extended_monthly_data,
        ),
    )

    return True


@pytest.fixture(scope="session")
def create_data_warehouse_validation_model_tables(
    mf_test_session_state: MetricFlowTestSessionState, sql_client: SqlClient
) -> bool:
    """Creates tables with example source data for testing."""
    schema = mf_test_session_state.mf_source_schema

    data = [
        (True, "2022-06-01"),
    ]
    create_table(
        sql_client=sql_client,
        sql_table=SqlTable(schema_name=schema, table_name="fct_animals"),
        df=make_df(
            sql_client=sql_client,
            columns=["is_dog", DEFAULT_DS],
            time_columns={DEFAULT_DS},
            data=data,
        ),
    )

    return True
