"""Cohort retention analytics queries."""

import pandas as pd
from sqlalchemy import Engine, text


def cohort_retention(engine: Engine) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Compute monthly cohort retention rates.

    Returns
    -------
    retention_pivot : pd.DataFrame
        Pivoted table with cohort months as index, month offsets (0, 1, 2, …)
        as columns, and retention percentages as values.
    cohort_sizes : pd.DataFrame
        Cohort month and size (number of unique customers in each acquisition cohort).
    """
    # NOTE: In the Olist dataset, customer_id (and therefore customer_key) is
    # unique per ORDER, not per customer. customer_unique_id is the true
    # cross-order customer identifier. We must join through dim_customers to
    # get customer_unique_id before computing cohort membership.
    sql = text("""
        WITH first_purchase AS (
            SELECT
                dc.customer_unique_id,
                strftime('%Y-%m-01', MIN(f.order_date_id)) AS cohort_month
            FROM fact_orders f
            JOIN dim_customers dc ON f.customer_key = dc.customer_key
            WHERE f.order_status NOT IN ('canceled', 'unavailable')
            GROUP BY dc.customer_unique_id
        ),
        monthly_active AS (
            SELECT DISTINCT
                dc.customer_unique_id,
                strftime('%Y-%m-01', f.order_date_id) AS order_month
            FROM fact_orders f
            JOIN dim_customers dc ON f.customer_key = dc.customer_key
            WHERE f.order_status NOT IN ('canceled', 'unavailable')
        ),
        cohort_activity AS (
            SELECT
                fp.cohort_month,
                CAST(
                    (CAST(strftime('%Y', ma.order_month) AS INTEGER) - CAST(strftime('%Y', fp.cohort_month) AS INTEGER)) * 12 +
                    (CAST(strftime('%m', ma.order_month) AS INTEGER) - CAST(strftime('%m', fp.cohort_month) AS INTEGER))
                AS INTEGER) AS month_offset,
                COUNT(DISTINCT ma.customer_unique_id) AS active_customers
            FROM first_purchase fp
            JOIN monthly_active ma ON fp.customer_unique_id = ma.customer_unique_id
            WHERE ma.order_month >= fp.cohort_month
            GROUP BY fp.cohort_month, month_offset
        ),
        cohort_sizes AS (
            SELECT cohort_month, COUNT(DISTINCT customer_unique_id) AS cohort_size
            FROM first_purchase
            GROUP BY cohort_month
        )
        SELECT
            ca.cohort_month,
            ca.month_offset,
            ROUND(100.0 * ca.active_customers / cs.cohort_size, 1) AS retention_pct,
            cs.cohort_size
        FROM cohort_activity ca
        JOIN cohort_sizes cs ON ca.cohort_month = cs.cohort_month
        ORDER BY ca.cohort_month, ca.month_offset
    """)

    with engine.connect() as conn:
        df = pd.read_sql(sql, conn)

    df["cohort_month"] = pd.to_datetime(df["cohort_month"]).dt.strftime("%Y-%m")

    # Drop tiny cohorts — statistically meaningless (e.g. 2016-09 has 2 customers,
    # 2016-12 has 1, so any single return shows as 50–100 % retention).
    MIN_COHORT_SIZE = 100
    valid_cohorts = df.loc[df["cohort_size"] >= MIN_COHORT_SIZE, "cohort_month"].unique()
    df = df[df["cohort_month"].isin(valid_cohorts)]

    # Pivot: rows = cohort month, columns = month offset
    retention_pivot = df.pivot(
        index="cohort_month", columns="month_offset", values="retention_pct"
    )
    retention_pivot.columns = [f"Month {c}" for c in retention_pivot.columns]
    retention_pivot.index.name = "Cohort"

    # Drop Month 0 — always 100 % by definition (first purchase IS the acquisition
    # event), so it adds no information and compresses the colorscale for months 1+.
    retention_pivot = retention_pivot.drop(columns=["Month 0"], errors="ignore")

    cohort_sizes = (
        df[["cohort_month", "cohort_size"]]
        .drop_duplicates()
        .rename(columns={"cohort_month": "Cohort", "cohort_size": "Customers"})
        .reset_index(drop=True)
    )

    return retention_pivot, cohort_sizes