"""RFM (Recency, Frequency, Monetary) feature computation."""

from datetime import date

import pandas as pd
from sqlalchemy import Engine, text

RFM_SEGMENT_MAP = {
    (4, 4): "Champions",
    (4, 3): "Loyal",
    (3, 4): "Loyal",
    (3, 3): "Loyal",
    (4, 2): "Potential Loyalist",
    (3, 2): "Potential Loyalist",
    (4, 1): "Promising",
    (3, 1): "Promising",
    (2, 4): "At Risk",
    (2, 3): "At Risk",
    (2, 2): "Need Attention",
    (2, 1): "About To Sleep",
    (1, 4): "Lost",
    (1, 3): "Lost",
    (1, 2): "Lost",
    (1, 1): "Lost",
}


def _assign_rfm_segment(r_score: int, f_score: int) -> str:
    return RFM_SEGMENT_MAP.get((r_score, f_score), "Unknown")


def compute_rfm(engine: Engine, reference_date: date | None = None) -> pd.DataFrame:
    """
    Compute RFM scores for each customer_unique_id.

    Returns a DataFrame with columns:
        customer_unique_id, customer_key, recency_days, frequency, monetary,
        r_score, f_score, m_score, rfm_score, rfm_segment
    """
    if reference_date is None:
        with engine.connect() as conn:
            result = conn.execute(
                text("SELECT MAX(order_date_id) FROM warehouse.fact_orders")
            )
            reference_date = result.scalar()

    sql = text("""
        SELECT
            dc.customer_unique_id,
            :ref_date - MAX(f.order_date_id)   AS recency_days,
            COUNT(DISTINCT f.order_id)          AS frequency,
            SUM(f.price + f.freight_value)      AS monetary
        FROM warehouse.fact_orders f
        JOIN warehouse.dim_customers dc ON f.customer_key = dc.customer_key
        WHERE f.order_status NOT IN ('canceled', 'unavailable')
        GROUP BY dc.customer_unique_id
    """)

    with engine.connect() as conn:
        df = pd.read_sql(sql, conn, params={"ref_date": reference_date})

    # Quartile scoring (1=worst, 4=best)
    df["r_score"] = pd.qcut(df["recency_days"], q=4, labels=[4, 3, 2, 1]).astype(int)
    df["f_score"] = pd.qcut(df["frequency"].rank(method="first"), q=4, labels=[1, 2, 3, 4]).astype(int)
    df["m_score"] = pd.qcut(df["monetary"].rank(method="first"), q=4, labels=[1, 2, 3, 4]).astype(int)

    df["rfm_score"] = df["r_score"].astype(str) + df["f_score"].astype(str) + df["m_score"].astype(str)
    df["rfm_segment"] = df.apply(
        lambda row: _assign_rfm_segment(row["r_score"], row["f_score"]), axis=1
    )

    return df