# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "polars>=0.20.0",
#   "sqlalchemy>=2.0.0",
#   "pyodbc>=5.0.0",
#   "clickhouse-connect>=0.7.0",
#   "python-dotenv>=1.0.0",
#   "pyarrow>=18.0.0",
# ]
# ///

"""
ETL Pipeline: Coverage KPI — SQL Server → Polars → ClickHouse
--------------------------------------------------------------
Run with:   uv run etl_coverage_kpi.py

Why SQLAlchemy instead of raw pyodbc?
  - pl.read_database() speaks SQLAlchemy Engine natively
  - Connection pooling built-in (reuses connections, no reconnect overhead)
  - Parameterised queries use :param syntax — cleaner and SQL-injection safe
  - One engine object shared across all steps — no passing conn handles around
  - Works with ANY database backend, not just SQL Server


  1. Build one SQLAlchemy engine (the "pipe" to SQL Server)
  2. Ask: which store visits happened yesterday? → get MIDs
  3. Pull Coverage rows for those MIDs in chunks → Polars DataFrames
  4. Clean + type-cast everything in Polars
  5. Delete stale ClickHouse rows, insert fresh batch
  6. Verify row count matches
"""

import os
import pyarrow
import sys
import logging
from datetime import date, timedelta

import polars as pl
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine, URL
import clickhouse_connect
from dotenv import load_dotenv

load_dotenv()

# ── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("etl_coverage")

# ── Config ────────────────────────────────────────────────────────────────────
# SQL Server credentials — read from .env, never hardcoded

MSSQL_SERVER = os.getenv("MSSQL_SERVER", "your_server")
MSSQL_DB     = os.getenv("MSSQL_DB",     "CPMIndiaBusinessInsight")
MSSQL_USER   = os.getenv("MSSQL_USER",   "sa")
MSSQL_PASS   = os.getenv("MSSQL_PASS",   "")

# pyodbc connection string — built separately so special characters in
# passwords (@ / + # etc.) never corrupt the URL.
# The safe pattern: build the full pyodbc DSN as a plain string,
# then hand it to URL.create() via the "odbc_connect" query key.
# URL.create() percent-encodes it internally — no manual urllib needed.
_ODBC_DSN = (
    f"DRIVER={{ODBC Driver 18 for SQL Server}};"
    f"SERVER={MSSQL_SERVER};"
    f"DATABASE={MSSQL_DB};"
    f"UID={MSSQL_USER};"
    f"PWD={MSSQL_PASS};"
    "TrustServerCertificate=yes;"
    "timeout=30;"
)
# Final SQLAlchemy URL — credentials live inside odbc_connect, not in the URL
# This is the officially recommended pattern in SQLAlchemy docs for mssql+pyodbc
MSSQL_SA_URL = URL.create(
    drivername="mssql+pyodbc",
    query={"odbc_connect": _ODBC_DSN},  # URL.create handles encoding internally
)

CH_HOST  = os.getenv("CH_HOST",  "localhost")
CH_PORT  = int(os.getenv("CH_PORT", "8123"))
CH_USER  = os.getenv("CH_USER",  "default")
CH_PASS  = os.getenv("CH_PASS",  "")
CH_DB    = os.getenv("CH_DB",    "kelloggs")
CH_TABLE = "coverage_kpi"

BATCH_SIZE = 5000   # rows per ClickHouse insert
MID_CHUNK  = 1000   # SQL Server IN-clause chunk size
PROJECT_ID = 40148


# ── Step 1: Build SQLAlchemy Engine ───────────────────────────────────────────
def build_engine() -> Engine:
    """
    The Engine is not a connection — it's the factory that makes connections.
    Think of it as the water main under the street, not the tap in your kitchen.

    How pyodbc + SQLAlchemy work together here:
      - SQLAlchemy's mssql+pyodbc dialect is the translator
        (it knows how to speak SQL Server's dialect)
      - pyodbc is the actual wire — it talks ODBC to the SQL Server driver
      - URL.create(drivername="mssql+pyodbc", query={"odbc_connect": DSN})
        is the safe way to pass the pyodbc DSN through SQLAlchemy without
        URL-encoding bugs when passwords have special characters

    fast_executemany=True:
      pyodbc has a mode where it sends many rows to SQL Server in one
      network round-trip instead of one-row-at-a-time. We're only reading
      here, but it's good practice to enable it on the engine level.

    pool_size=3, max_overflow=2:
      Keep 3 live connections in the pool, allow 2 extras under burst load.
      After each `with engine.connect()` block exits, the connection goes
      back to the pool — not closed — so the next query reuses it instantly.
    """
    log.info("Building SQLAlchemy + pyodbc engine for SQL Server...")
    engine = create_engine(
        MSSQL_SA_URL,
        pool_size=3,
        max_overflow=2,
        pool_pre_ping=True,        # verify connection is alive before use
        echo=False,                # True = print every SQL to console (debug)
        connect_args={
            "fast_executemany": True,   # pyodbc bulk-send optimisation
        },
    )
    # Smoke test — borrow one connection, run SELECT 1, return it to pool
    with engine.connect() as conn:
        conn.execute(text("SELECT 1"))
    log.info("SQL Server engine ready.")
    return engine


# ── Step 2: Collect MIDs ──────────────────────────────────────────────────────
def collect_mids(engine: Engine, target_date: date) -> list[int]:
    """
    SQLAlchemy parameterised query using :param syntax.
    This is SQL-injection safe — SQLAlchemy never concatenates
    the date value into the string directly.

    We use engine.connect() as a context manager —
    the connection is returned to the pool automatically when done.
    """
    sql = text("""
        SELECT MID FROM OneApp_KelloggsMT.dbo.T_StoreCoverage
        WHERE CONVERT(date, CreateDate) = :target_date
        UNION
        SELECT MID FROM OneApp_KelloggsMT.dbo.T_StoreCoverage
        WHERE CONVERT(date, UpdateDate) = :target_date
    """)
    log.info(f"Collecting MIDs for: {target_date}")
    with engine.connect() as conn:
        result = conn.execute(sql, {"target_date": target_date})
        mids = [row[0] for row in result.fetchall()]
    log.info(f"Found {len(mids):,} MIDs")
    return mids


# ── Step 3: Fetch Coverage data via pl.read_database() ───────────────────────
def fetch_coverage_data(engine: Engine, mids: list[int]) -> pl.DataFrame:
    """
    This is the key upgrade from pyodbc to SQLAlchemy+Polars.

    pl.read_database(query, connection) accepts a SQLAlchemy Engine directly.
    Polars handles the cursor, column name detection, and type inference
    automatically — no more manual cursor.fetchall() loops.

    We still chunk the MID list (SQL Server IN clause limit ~2100 params),
    but each chunk is now one clean pl.read_database() call.

    Feynman analogy: instead of manually scooping water bucket-by-bucket
    (pyodbc cursor loop), we turn on a tap and let it fill the container
    (pl.read_database manages the stream).
    """
    if not mids:
        log.warning("No MIDs — nothing to fetch.")
        return pl.DataFrame()

    # SQL mirrors the Coverage INSERT from the stored procedure
    # :mid_list is replaced per chunk below
    coverage_sql = """
        SELECT
            {project_id}                        AS project_id,
            JP.MID,
            sm.StoreId                          AS store_id,
            JP.EmpId                            AS employee_id,
            CONVERT(date, JP.VisitDate)         AS visit_date,
            JP.InTime                           AS in_time,
            JP.OutTime                          AS out_time,
            CASE
                WHEN JP.OutTime IS NULL OR JP.InTime IS NULL THEN NULL
                WHEN JP.OutTime < JP.InTime THEN 0
                ELSE DATEDIFF(SECOND, JP.InTime, JP.OutTime) / 60
            END                                 AS duration_minutes,
            CASE
                WHEN (
                    SELECT TOP 1 EmpId
                    FROM OneApp_KelloggsMT.dbo.T_StoreCoverage SC
                    WHERE SC.EmpId = JP.EmpId
                      AND SC.StoreId = JP.StoreId
                      AND SC.VisitDate = JP.VisitDate
                      AND SC.ReasonId IN (0,1,3,9,10,19,20)
                ) > 0
                THEN 'Y' ELSE 'N'
            END                                 AS is_covered,
            CASE JP.Deviation
                WHEN 0 THEN 'Planned'
                WHEN 1 THEN 'Adhoc'
                WHEN 2 THEN 'Beat Plan'
                WHEN 3 THEN 'Non Merchandised'
                WHEN 4 THEN 'Add New Store'
                WHEN 5 THEN 'Non Program'
                ELSE ''
            END                                 AS coverage_type,
            sm.StoreTypeId                      AS storetype_id,
            Em.SupervisorId                     AS supervisor_id,
            ISNULL(JP.ReasonId, 0)              AS reason_id,
            sm.CameraAllow                      AS camera_allow,
            CAST(
                CASE
                    WHEN sm.Latitude  IS NULL OR sm.Latitude  = 0 THEN 0
                    WHEN sm.Longitude IS NULL OR sm.Longitude = 0 THEN 0
                    WHEN JP.Latitude  IS NULL OR JP.Latitude  = 0 THEN 0
                    WHEN JP.Longitude IS NULL OR JP.Longitude = 0 THEN 0
                    ELSE SQRT(
                        POWER(69.1 * (JP.Latitude  - sm.Latitude),  2) +
                        POWER(69.1 * (sm.Longitude - JP.Longitude)
                              * COS(JP.Latitude / 57.3), 2)
                    ) * 1000
                END AS FLOAT
            )                                   AS distance_meters,
            GETDATE()                           AS update_date,
            'ETL-SQLAlchemy'                    AS update_by
        FROM OneApp_KelloggsMT.dbo.T_StoreCoverage JP WITH (NOLOCK)
        INNER JOIN OneApp_KelloggsMT.dbo.vw_StoreDetail sm
            ON JP.StoreId = sm.StoreId
        INNER JOIN OneApp_KelloggsMT.dbo.vw_Employee_Detail Em
            ON JP.EmpId = Em.EmpId
        WHERE JP.MID IN ({mid_placeholders})
          AND Em.UserName NOT LIKE 'test%'
    """

    frames: list[pl.DataFrame] = []

    for i in range(0, len(mids), MID_CHUNK):
        chunk = mids[i : i + MID_CHUNK]
        # Build the IN list as a literal — SQLAlchemy text() for IN-list with
        # variable length is cleanest done this way (all values are integers,
        # so no injection risk)
        mid_list = ",".join(str(m) for m in chunk)
        sql = coverage_sql.format(
            project_id=PROJECT_ID,
            mid_placeholders=mid_list,
        )

        log.info(
            f"Fetching chunk {i // MID_CHUNK + 1} "
            f"(MIDs {i+1}–{min(i+MID_CHUNK, len(mids))} of {len(mids)})..."
        )

        # pl.read_database() takes the raw SQL string + the SQLAlchemy engine.
        # It opens a connection from the pool, streams the result into a
        # Polars DataFrame, then releases the connection back automatically.
        df_chunk = pl.read_database(query=sql, connection=engine)
        if len(df_chunk) > 0:
            frames.append(df_chunk)

    if not frames:
        log.warning("No rows returned from SQL Server.")
        return pl.DataFrame()

    df = pl.concat(frames)
    log.info(f"Fetched {len(df):,} total rows from SQL Server")
    return df


# ── Step 4: Transform with Polars ─────────────────────────────────────────────
def transform(df: pl.DataFrame) -> pl.DataFrame:
    """
    SQLAlchemy + Polars gives us richer type inference than raw pyodbc —
    but we still normalise everything explicitly so ClickHouse never
    gets a surprise type.

    The chain of .with_columns() calls is like an assembly line:
    each station does one job and passes the belt to the next.
    """
    log.info("Transforming...")

    df = (
        df
        # ── IDs ──────────────────────────────────────────────────────────────
        .with_columns([
            pl.col("project_id").cast(pl.Int32),
            pl.col("MID").cast(pl.Int64),
            pl.col("store_id").cast(pl.Int64),
            pl.col("employee_id").cast(pl.Int64),
            pl.col("supervisor_id").cast(pl.Int64).fill_null(0),
            pl.col("storetype_id").cast(pl.Int32).fill_null(0),
            pl.col("reason_id").cast(pl.Int32).fill_null(0),
        ])

        # ── Dates ─────────────────────────────────────────────────────────────
        # SQLAlchemy may return visit_date as datetime — cast to Date
        .with_columns(
            pl.col("visit_date").cast(pl.Date)
        )

        # ── Numerics ──────────────────────────────────────────────────────────
        # DATEDIFF/CAST in SQL Server comes through as Decimal via ODBC
        # → cast to Float64 first, then to Int32 (avoids Decimal overflow)
        .with_columns(
            pl.col("duration_minutes")
              .cast(pl.Float64)
              .fill_null(0.0)
              .cast(pl.Int32)
        )
        .with_columns(
            pl.col("distance_meters")
              .cast(pl.Float64)
              .fill_null(0.0)
              .round(2)
        )

        # ── Flags ─────────────────────────────────────────────────────────────
        # is_covered: 'Y'/'N' → UInt8 (1/0) — ClickHouse has no Bool type
        .with_columns(
            pl.when(pl.col("is_covered") == "Y")
              .then(pl.lit(1, dtype=pl.UInt8))
              .otherwise(pl.lit(0, dtype=pl.UInt8))
              .alias("is_covered")
        )
        # camera_allow: SQL Server BIT → UInt8
        .with_columns(
            pl.col("camera_allow").cast(pl.UInt8).fill_null(0)
        )

        # ── Strings ───────────────────────────────────────────────────────────
        # ClickHouse String column rejects NULL — replace with empty string
        .with_columns([
            pl.col("coverage_type").fill_null(""),
            pl.col("update_by").fill_null("ETL-SQLAlchemy"),
        ])

        # ── Timestamps ────────────────────────────────────────────────────────
        .with_columns(
            pl.col("update_date").cast(pl.Datetime)
        )

        # ── Dedup ─────────────────────────────────────────────────────────────
        # MID is unique per project — drop any duplicates from chunked joins
        .unique(subset=["project_id", "MID"])
    )

    log.info(f"After transform: {len(df):,} rows | schema: {df.schema}")
    return df


# ── Step 5: Load to ClickHouse ────────────────────────────────────────────────
def load_to_clickhouse(df: pl.DataFrame) -> tuple[int, date]:
    """
    clickhouse-connect v0.7+ supports insert_arrow() which accepts
    a Polars DataFrame via its Arrow IPC backing — no pandas bridge needed.

    Delete-then-insert mirrors the stored procedure's pattern:
    wipe the day's data, reload fresh. Idempotent = safe to re-run.
    """
    target_date: date = df["visit_date"].min()  # always yesterday

    log.info(f"Connecting to ClickHouse {CH_HOST}:{CH_PORT} db={CH_DB}...")
    client = clickhouse_connect.get_client(
        host=CH_HOST, port=CH_PORT,
        username=CH_USER, password=CH_PASS,
        database=CH_DB,
    )

    # Delete stale rows for this date + project
    log.info(f"Deleting existing rows: project_id={PROJECT_ID} visit_date={target_date}")
    client.command(
        f"ALTER TABLE {CH_TABLE} DELETE "
        f"WHERE project_id = {PROJECT_ID} "
        f"AND visit_date = '{target_date}'"
    )

    # Insert in batches using Arrow — Polars → Arrow is zero-copy
    total = 0
    for start in range(0, len(df), BATCH_SIZE):
        batch = df.slice(start, BATCH_SIZE)
        # to_arrow() converts the Polars DataFrame to Apache Arrow Table
        # insert_arrow() sends it directly — no pandas, no row-by-row overhead
        client.insert_arrow(CH_TABLE, batch.to_arrow())
        total += len(batch)
        log.info(f"  Inserted {total:,} / {len(df):,} rows")

    log.info(f"ClickHouse load complete: {total:,} rows")
    return total, target_date


# ── Step 6: Verify ────────────────────────────────────────────────────────────
def verify(expected: int, target_date: date) -> bool:
    """
    Count rows in ClickHouse for this date and compare.
    If they don't match — the ETL failed silently somewhere.
    Exit code 2 so your scheduler/cron knows it's a data problem, not a crash.
    """
    client = clickhouse_connect.get_client(
        host=CH_HOST, port=CH_PORT,
        username=CH_USER, password=CH_PASS,
        database=CH_DB,
    )
    result = client.query(
        f"SELECT count() FROM {CH_TABLE} "
        f"WHERE project_id = {PROJECT_ID} "
        f"AND visit_date = '{target_date}'"
    )
    actual = result.result_rows[0][0]
    log.info(f"Verify — expected={expected:,}  clickhouse={actual:,}")
    if actual != expected:
        log.error(f"MISMATCH: sent {expected:,}, ClickHouse has {actual:,}")
        return False
    log.info("Verify passed.")
    return True


# ── Main ──────────────────────────────────────────────────────────────────────
def main() -> None:
    #target_date = date.today() - timedelta(days=1)
    for i in range(1, 10):
        target_date = date(2023, 3, 9) + timedelta(days=i)
        print(target_date)
        log.info(f"=== Coverage KPI ETL | date={target_date} ===")

        # Build the SQLAlchemy engine once — shared across all steps
        engine = build_engine()

        # Collect which MIDs need processing
        mids = collect_mids(engine, target_date)
        if not mids:
            log.info("No MIDs for yesterday. Nothing to do.")
            sys.exit(0)

        # Fetch raw data using pl.read_database()
        df_raw = fetch_coverage_data(engine, mids)
        engine.dispose()   # return all pooled connections to OS cleanly

        if df_raw.is_empty():
            log.warning("Empty result set. Exiting.")
            sys.exit(0)

        # Transform
        df_clean = transform(df_raw)

        # Load
        rows_inserted, inserted_date = load_to_clickhouse(df_clean)

        # Verify
        if not verify(rows_inserted, inserted_date):
            sys.exit(2)

        log.info(f"=== Done. {rows_inserted:,} rows loaded for {inserted_date} ===")


if __name__ == "__main__":
    main()