# Lakeflow Spark Declarative Pipelines source: business rules as expectations.
# Data Meaning guide: https://www.datameaning.com/guides/databricks-lakeflow-expectations-from-business-rules/
#
# Add this file as the source of a serverless pipeline whose target schema holds the two input tables.
# Unqualified table names resolve to the pipeline's target schema.
from pyspark import pipelines as dp

# ---- EDIT THIS BLOCK: your table and column names, and your business rules --------------------------
SOURCE_TABLE = "orders_raw"        # the table the rules apply to
LOOKUP_TABLE = "customers_ref"     # the table a key must exist in
SOURCE_KEY = "o_custkey"           # the column in SOURCE_TABLE that must exist in LOOKUP_TABLE
LOOKUP_KEY = "c_custkey"           # the matching column in LOOKUP_TABLE

# Each rule is: name -> SQL condition that is TRUE for a good row. The name is what the event log shows.
# A rule that stops the update. Use it only for a problem that makes every later number untrustworthy.
RULES_FAIL = {
    "order_key_present": "o_orderkey IS NOT NULL",
}
# Rules that remove the row from the clean table. The same rules fill the quarantine table.
RULES_DROP = {
    "price_positive": "o_totalprice > 0",
    "status_valid": "o_orderstatus IN ('F', 'O', 'P')",
    "order_date_present": "o_orderdate IS NOT NULL",
    "customer_exists": "customer_found",   # set by the join below: expectations cannot hold a subquery
}
# Rules that only count: the row is kept, the failures are recorded.
RULES_WARN = {
    "priority_known": "o_orderpriority IN ('1-URGENT', '2-HIGH', '3-MEDIUM', '4-NOT SPECIFIED', '5-LOW')",
}
# ------------------------------------------------------------------------------------------------------


def _passes(condition):
    # A rule that evaluates to NULL is a failure, so the quarantine treats it as one.
    return f"coalesce(({condition}), false)"


QUARANTINE_FILTER = "NOT (" + " AND ".join(_passes(c) for c in RULES_DROP.values()) + ")"
FAILED_RULES = "concat_ws(', ', " + ", ".join(
    f"CASE WHEN NOT {_passes(c)} THEN '{n}' END" for n, c in RULES_DROP.items()) + ")"


@dp.materialized_view(comment="Source rows plus a customer_found flag; stops the update if a rule in RULES_FAIL breaks")
@dp.expect_all_or_fail(RULES_FAIL)
def orders_checked():
    return spark.sql(f"""
        SELECT s.*, l.{LOOKUP_KEY} IS NOT NULL AS customer_found
        FROM {SOURCE_TABLE} s
        LEFT JOIN (SELECT DISTINCT {LOOKUP_KEY} FROM {LOOKUP_TABLE}) l ON s.{SOURCE_KEY} = l.{LOOKUP_KEY}
    """)


@dp.materialized_view(comment="Rows that pass every drop rule")
@dp.expect_all_or_drop(RULES_DROP)
@dp.expect_all(RULES_WARN)
def orders_clean():
    return spark.read.table("orders_checked")


@dp.materialized_view(comment="Rows that broke at least one drop rule, with the rules they broke")
def orders_quarantine():
    return (spark.read.table("orders_checked")
            .filter(QUARANTINE_FILTER)
            .selectExpr("*", f"{FAILED_RULES} AS failed_rules"))