import great_expectations as gx
from great_expectations.core.batch import RuntimeBatchRequest

context = gx.get_context()

suite = context.add_or_update_expectation_suite("omni_reservations_suite")

def build_expectation_suite(validator):
    # Uniqueness
    validator.expect_column_values_to_be_unique("res_id")

    # Completeness — structurally required fields
    for col in ["guest_name", "check_in", "check_out", "property", "room_type"]:
        validator.expect_column_values_to_not_be_null(col)

    # Validity — rate range ($89–$1500 covers standard to Presidential)
    validator.expect_column_values_to_be_between("rate", min_value=89, max_value=1500)

    # Validity — email regex
    validator.expect_column_values_to_match_regex(
        "email", regex=r"^[\w.+-]+@[\w-]+\.[\w.]+$"
    )

    # Validity — loyalty tier set
    validator.expect_column_values_to_be_in_set(
        "loyalty",
        value_set=["Select Guest", "Gold", "Platinum"]
    )

    # Consistency — checkout after checkin (custom SQL expectation)
    validator.expect_column_pair_values_to_be_equal(
        column_A="check_out", column_B="check_in",
        or_equal=False,
        meta={"description": "check_out must be after check_in"}
    )

    return validator.get_expectation_suite()

import pandas as pd
from ydata_profiling import ProfileReport
import snowflake.connector

def profile_reservations(date_range: tuple[str, str]) -> ProfileReport:
    conn = snowflake.connector.connect(
        user="etl_service", account="omni-hotels.us-east-1",
        warehouse="ANALYTICS_WH", database="OPERA_PROD"
    )
    df = pd.read_sql(f"""
        SELECT res_id, guest_name, email, phone, check_in, check_out,
               room_type, rate, loyalty_tier, property_code
        FROM   reservations
        WHERE  created_date BETWEEN '{date_range[0]}' AND '{date_range[1]}'
    """, conn)

    profile = ProfileReport(
        df,
        title=f"Omni Reservations {date_range[0]} to {date_range[1]}",
        explorative=True,
        correlations={"auto": {"calculate": True}},
        missing_diagrams={"heatmap": True, "dendrogram": True}
    )

    # Key stats: missing rates, value distributions, correlations
    completeness = (1 - df.isnull().mean()) * 100
    print(completeness.sort_values())          # spot bottom fields
    print(df["rate"].describe())              # outlier check
    return profile

import numpy as np
from dataclasses import dataclass

@dataclass
class QualityAlert:
    date: str; score: float; z_score: float; prev_scores: list[float]

class QualityRegressionDetector:
    """Detect unusual drops in daily DQ score using rolling Z-Score.
    Fires if today's score is >2.5 sigma below the rolling mean."""
    def __init__(self, window: int = 14, threshold: float = 2.5):
        self.window = window; self.threshold = threshold
        self._history: list[tuple[str, float]] = []

    def update(self, date: str, score: float) -> QualityAlert | None:
        self._history.append((date, score))
        if len(self._history) <= self.window:
            return None

        recent = [s for _, s in self._history[-self.window - 1:-1]]
        mu = np.mean(recent); sigma = np.std(recent, ddof=1)
        if sigma == 0: return None

        z = (score - mu) / sigma
        if z < -self.threshold:
            return QualityAlert(date=date, score=score, z_score=z, prev_scores=recent)
        return None

# Usage: called nightly after GE checkpoint
detector = QualityRegressionDetector(window=14)
for date, score in daily_scores:  # from DQ metrics table
    alert = detector.update(date, score)
    if alert: send_pagerduty(alert)   # page duty manager

from enum import Enum
from dataclasses import dataclass, field

class Severity(str, Enum):
    CRITICAL = "critical"   # blocks load, pages duty manager
    HIGH     = "high"       # quarantines row, creates Jira P1
    MEDIUM   = "medium"     # loads with quality flag, Jira P2
    LOW      = "low"        # logs only, next-sprint backlog

DIMENSION_SLA = {
    "CONSISTENCY": Severity.CRITICAL,  # checkout < checkin
    "RANGE":       Severity.HIGH,       # rate = $0 or $5000+
    "UNIQUE":      Severity.HIGH,       # duplicate res_id
    "NOT NULL":    Severity.MEDIUM,     # missing required field
    "REGEX":       Severity.MEDIUM,     # invalid email/phone
    "SET":         Severity.LOW,        # unexpected loyalty tier
    "DATE PAST":   Severity.LOW,        # future date in historical extract
}

@dataclass
class RemediationAction:
    res_id: str; rule_type: str; severity: Severity; action: str

def route_failures(failures: list[dict]) -> list[RemediationAction]:
    actions = []
    for f in failures:
        sev = DIMENSION_SLA.get(f["rule_type"], Severity.LOW)
        action_map = {
            Severity.CRITICAL: "BLOCK_LOAD|PAGE_MANAGER",
            Severity.HIGH:     "QUARANTINE|CREATE_JIRA_P1",
            Severity.MEDIUM:   "LOAD_FLAGGED|CREATE_JIRA_P2",
            Severity.LOW:      "LOAD_FLAGGED|LOG_ONLY",
        }
        actions.append(RemediationAction(
            res_id=f["res_id"], rule_type=f["rule_type"],
            severity=sev, action=action_map[sev]
        ))
    return actions