"""
Generate churn labels from historical data.
For each customer observation point, determine if they churned within 30 days.

Label = 1 (churned): Customer had status 3 or 7 within 30 days of observation
Label = 0 (retained): Customer remained active (status 2) for 30+ days after observation
"""

import sys
import argparse
import psycopg2
from clickhouse_client import ClickHouseClient
from config import CLICKHOUSE_CONFIG, CHURN_TRAINING_LABELS_TABLE, DB_CONFIG, DEFAULT_APP_ID


def get_clickhouse_client():
    return ClickHouseClient(CLICKHOUSE_CONFIG)


def generate_labels(app_id, lookback_months=6):
    """Generate training labels from historical subscription data."""
    
    query = """
    WITH customer_events AS (
        -- Get all subscription status changes
        SELECT 
            sms_cust_id,
            app_id,
            status,
            start_date,
            end_date,
            created_at,
            ROW_NUMBER() OVER (PARTITION BY sms_cust_id ORDER BY created_at) AS event_seq
        FROM tbl_customer_subscriptions
        WHERE app_id = %(app_id)s
          AND created_at >= NOW() - INTERVAL '%(months)s months'
    ),
    observation_points AS (
        -- Create monthly observation points for each customer
        SELECT DISTINCT
            ce.sms_cust_id,
            ce.app_id,
            date_trunc('month', gs.observation_month)::DATE AS observation_date
        FROM customer_events ce
        CROSS JOIN generate_series(
            (NOW() - INTERVAL '%(months)s months')::DATE,
            (NOW() - INTERVAL '30 days')::DATE,
            INTERVAL '1 month'
        ) AS gs(observation_month)
        WHERE ce.created_at <= gs.observation_month
    ),
    churn_labels AS (
        SELECT 
            op.sms_cust_id,
            op.app_id,
            op.observation_date,
            -- Did this customer have a churn event within 30 days of observation?
            CASE 
                WHEN EXISTS (
                    SELECT 1 
                    FROM tbl_customer_subscriptions cs
                    WHERE cs.sms_cust_id = op.sms_cust_id
                      AND cs.app_id = op.app_id
                      AND cs.status IN (3, 7)
                      AND cs.end_date BETWEEN op.observation_date 
                                          AND op.observation_date + INTERVAL '30 days'
                ) THEN TRUE
                ELSE FALSE
            END AS churned,
            -- Actual churn date if churned
            (
                SELECT MIN(cs.end_date)::DATE
                FROM tbl_customer_subscriptions cs
                WHERE cs.sms_cust_id = op.sms_cust_id
                  AND cs.app_id = op.app_id
                  AND cs.status IN (3, 7)
                  AND cs.end_date BETWEEN op.observation_date 
                                      AND op.observation_date + INTERVAL '30 days'
            ) AS churn_date
        FROM observation_points op
    )
    SELECT sms_cust_id, app_id, churned, observation_date, churn_date
    FROM churn_labels
    ORDER BY observation_date, sms_cust_id;
    """
    
    try:
        print(f"Connecting to DB and generating labels for app_id={app_id} (lookback: {lookback_months}m)...")
        clickhouse = get_clickhouse_client()
        with psycopg2.connect(**DB_CONFIG) as conn:
            count_before_rows = clickhouse.select(f"""
                SELECT count() AS total
                FROM {clickhouse.qualified_table(CHURN_TRAINING_LABELS_TABLE)}
                WHERE `app_id` = {int(app_id)}
            """)
            count_before = int(count_before_rows[0]['total']) if count_before_rows else 0

            with conn.cursor(name='churn_training_labels_cursor') as cur:
                cur.itersize = 2000
                cur.execute(query, {'app_id': app_id, 'months': lookback_months})

                while True:
                    rows = cur.fetchmany(2000)
                    if not rows:
                        break

                    payload = []
                    tuples = []
                    for sms_cust_id, row_app_id, churned, observation_date, churn_date in rows:
                        sms_cust_id = str(sms_cust_id or '').strip()
                        if sms_cust_id == '':
                            continue

                        observation_date_str = observation_date.strftime('%Y-%m-%d')
                        tuples.append(
                            f"({int(row_app_id)}, '{clickhouse.escape(sms_cust_id)}', toDate('{clickhouse.escape(observation_date_str)}'))"
                        )
                        payload.append({
                            'sms_cust_id': sms_cust_id,
                            'app_id': int(row_app_id),
                            'churned': 1 if churned else 0,
                            'observation_date': observation_date,
                            'churn_date': churn_date,
                        })

                    if not payload:
                        continue

                    clickhouse.command(f"""
                        ALTER TABLE {clickhouse.qualified_table(CHURN_TRAINING_LABELS_TABLE)}
                        DELETE WHERE tuple(`app_id`, `sms_cust_id`, `observation_date`) IN ({", ".join(tuples)})
                        SETTINGS mutations_sync = 2
                    """)
                    clickhouse.insert_json_rows(
                        CHURN_TRAINING_LABELS_TABLE,
                        payload,
                        columns=['sms_cust_id', 'app_id', 'churned', 'observation_date', 'churn_date'],
                    )

            count_after_rows = clickhouse.select(f"""
                SELECT count() AS total
                FROM {clickhouse.qualified_table(CHURN_TRAINING_LABELS_TABLE)}
                WHERE `app_id` = {int(app_id)}
            """)
            count_after = int(count_after_rows[0]['total']) if count_after_rows else 0

            print(f"Success! Labels generated. Added {count_after - count_before} new observation points.")
            print(f"Total labels for app_id={app_id}: {count_after}")
                
    except Exception as e:
        print(f"Database error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate churn labels for training.")
    parser.add_argument("--app-id", type=int, default=DEFAULT_APP_ID, help="App ID to process")
    parser.add_argument("--months", type=int, default=6, help="Months of history to look back")
    args = parser.parse_args()
    
    if args.app_id <= 0:
        print("Error: --app-id must be a positive integer.", file=sys.stderr)
        sys.exit(1)
        
    generate_labels(args.app_id, args.months)