"""
OMS Regression Data Service
Extracts and engineers features for ML models directly from OMS MySQL.
"""

from typing import List, Optional, Dict
import pandas as pd

from app.utils.helpers import get_logger
from app.models.mysql_model import MySQLModel


class OMSRegressionDataService:
    """Pulls and engineers ML features from OMS operational tables."""

    COLLECTION_FEATURES = [
        'total_collections', 'avg_settlement_lag', 'bounce_count',
        'bounce_rate', 'cheque_ratio', 'balance_credit_ratio',
    ]

    SALES_TARGET_FEATURES = [
        'target_qty', 'season_num', 'prev_attainment_pct',
        'product_head_id', 'district_id',
    ]

    def __init__(self):
        self.logger = get_logger("oms_regression_data_service")
        self._db = MySQLModel()

    # ------------------------------------------------------------------
    # Collection Risk
    # ------------------------------------------------------------------

    def get_collection_risk_training_data(self, company_id: int) -> pd.DataFrame:
        """
        Returns per-client feature matrix with binary label is_high_risk.
        Excludes clients with fewer than 3 credited collections (insufficient history).
        """
        sql = """
            SELECT
                col.client_id,
                COUNT(col.id)                                                           AS total_collections,
                COALESCE(AVG(DATEDIFF(col.credited_date, col.collection_date)), 0)      AS avg_settlement_lag,
                SUM(CASE WHEN col.bounce_date IS NOT NULL THEN 1 ELSE 0 END)            AS bounce_count,
                SUM(CASE WHEN col.bounce_date IS NOT NULL THEN 1 ELSE 0 END)
                    / COUNT(col.id)                                                     AS bounce_rate,
                SUM(CASE WHEN col.collection_mode = 2 THEN 1 ELSE 0 END)
                    / COUNT(col.id)                                                     AS cheque_ratio,
                COALESCE(
                    MAX(cb.closing_bal) / NULLIF(MAX(cl.total_credit_limit), 0),
                    0
                )                                                                       AS balance_credit_ratio,
                CASE
                    WHEN AVG(DATEDIFF(col.credited_date, col.collection_date)) > 30
                         OR SUM(CASE WHEN col.bounce_date IS NOT NULL THEN 1 ELSE 0 END) > 0
                    THEN 1 ELSE 0
                END                                                                     AS is_high_risk
            FROM collection col
            JOIN clients cl ON cl.id = col.client_id AND cl.deleted_at IS NULL
            LEFT JOIN client_balance cb
                ON cb.client_id = col.client_id AND cb.company_id = col.company_id
            WHERE col.company_id = %s
              AND col.deleted_at IS NULL
              AND col.credited_date IS NOT NULL
            GROUP BY col.client_id
            HAVING total_collections >= 3
        """
        rows = self._db.execute_query(sql, (company_id,))
        if not rows:
            return pd.DataFrame()
        df = pd.DataFrame(rows)
        numeric_cols = self.COLLECTION_FEATURES + ['is_high_risk', 'client_id']
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0)
        return df

    def get_collection_risk_scoring_features(
        self,
        company_id: int,
        client_ids: Optional[List[int]] = None
    ) -> pd.DataFrame:
        """
        Returns current feature values for scoring (no label column).
        Uses last 12 months of data for recency.
        client_ids=None means score all active clients.
        """
        client_filter = ""
        params = [company_id]
        if client_ids:
            placeholders = ",".join(["%s"] * len(client_ids))
            client_filter = f"AND col.client_id IN ({placeholders})"
            params.extend(client_ids)

        sql = f"""
            SELECT
                col.client_id,
                COUNT(col.id)                                                           AS total_collections,
                COALESCE(AVG(DATEDIFF(col.credited_date, col.collection_date)), 0)      AS avg_settlement_lag,
                SUM(CASE WHEN col.bounce_date IS NOT NULL THEN 1 ELSE 0 END)            AS bounce_count,
                SUM(CASE WHEN col.bounce_date IS NOT NULL THEN 1 ELSE 0 END)
                    / NULLIF(COUNT(col.id), 0)                                          AS bounce_rate,
                SUM(CASE WHEN col.collection_mode = 2 THEN 1 ELSE 0 END)
                    / NULLIF(COUNT(col.id), 0)                                          AS cheque_ratio,
                COALESCE(
                    MAX(cb.closing_bal) / NULLIF(MAX(cl.total_credit_limit), 0),
                    0
                )                                                                       AS balance_credit_ratio
            FROM collection col
            JOIN clients cl ON cl.id = col.client_id AND cl.deleted_at IS NULL
            LEFT JOIN client_balance cb
                ON cb.client_id = col.client_id AND cb.company_id = col.company_id
            WHERE col.company_id = %s
              AND col.deleted_at IS NULL
              AND col.credited_date IS NOT NULL
              AND col.collection_date >= DATE_SUB(NOW(), INTERVAL 12 MONTH)
              {client_filter}
            GROUP BY col.client_id
            HAVING total_collections >= 1
        """
        rows = self._db.execute_query(sql, tuple(params))
        if not rows:
            return pd.DataFrame()
        df = pd.DataFrame(rows)
        for col in self.COLLECTION_FEATURES + ['client_id']:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0)
        return df

    # ------------------------------------------------------------------
    # Sales Target Prediction
    # ------------------------------------------------------------------

    def get_sales_target_training_data(self, company_id: int) -> pd.DataFrame:
        """
        Returns historical attainment % per product_head × district × season.
        Adds lag-1 prev_attainment_pct feature via pandas shift.
        Requires at least 2 seasons per combination.
        """
        sql = """
            SELECT
                phdt.product_head_id,
                phdt.district_id,
                phdt.account_id                                                          AS account_year_id,
                phdt.target_qty,
                COALESCE(SUM(tsp.qty), 0)                                               AS actual_qty,
                COALESCE(SUM(tsp.qty), 0) / NULLIF(phdt.target_qty, 0) * 100           AS attainment_pct,
                ay.from_period
            FROM product_head_district_target phdt
            JOIN mas_accountyear ay ON ay.id = phdt.account_id
            LEFT JOIN transaction_sales ts
                ON ts.account_year_id = phdt.account_id
               AND ts.company_id = %s
               AND ts.status = 5
               AND ts.deleted_by > 0
            LEFT JOIN transaction_sales_products tsp ON tsp.tran_sales_id = ts.id
            LEFT JOIN products p
                ON p.id = tsp.product_id
               AND p.product_head_id = phdt.product_head_id
            WHERE phdt.target_qty > 0
            GROUP BY phdt.product_head_id, phdt.district_id, phdt.account_id,
                     phdt.target_qty, ay.from_period
            ORDER BY phdt.product_head_id, phdt.district_id, ay.from_period
        """
        rows = self._db.execute_query(sql, (company_id,))
        if not rows:
            return pd.DataFrame()

        df = pd.DataFrame(rows)
        for col in ['product_head_id', 'district_id', 'account_year_id',
                    'target_qty', 'actual_qty', 'attainment_pct']:
            df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0)

        # Season number and lag feature per product_head × district
        df = df.sort_values(['product_head_id', 'district_id', 'from_period'])
        df['season_num'] = df.groupby(['product_head_id', 'district_id']).cumcount() + 1
        df['prev_attainment_pct'] = df.groupby(
            ['product_head_id', 'district_id']
        )['attainment_pct'].shift(1).fillna(0)

        # Keep only combinations with at least 2 seasons (need lag-1)
        counts = df.groupby(['product_head_id', 'district_id'])['season_num'].count()
        valid = counts[counts >= 2].index
        df = df.set_index(['product_head_id', 'district_id'])
        df = df.loc[df.index.isin(valid)].reset_index()

        # Drop rows without lag (season_num == 1 and prev == 0 means first season, no history)
        df = df[df['season_num'] > 1]

        return df

    def get_sales_target_prediction_features(
        self,
        company_id: int,
        account_year_id: int,
        product_head_ids: Optional[List[int]] = None,
        district_ids: Optional[List[int]] = None
    ) -> pd.DataFrame:
        """
        Returns feature rows for the given account year to run predictions on.
        Includes lag feature from the most recent completed season.
        """
        ph_filter = ""
        d_filter = ""
        params = [account_year_id]

        if product_head_ids:
            placeholders = ",".join(["%s"] * len(product_head_ids))
            ph_filter = f"AND phdt.product_head_id IN ({placeholders})"
            params.extend(product_head_ids)
        if district_ids:
            placeholders = ",".join(["%s"] * len(district_ids))
            d_filter = f"AND phdt.district_id IN ({placeholders})"
            params.extend(district_ids)

        # Current season targets
        sql = f"""
            SELECT
                phdt.product_head_id,
                phdt.district_id,
                phdt.account_id AS account_year_id,
                phdt.target_qty,
                1 AS season_num
            FROM product_head_district_target phdt
            WHERE phdt.account_id = %s
              AND phdt.target_qty > 0
              {ph_filter}
              {d_filter}
        """
        rows = self._db.execute_query(sql, tuple(params))
        if not rows:
            return pd.DataFrame()

        df = pd.DataFrame(rows)
        for col in ['product_head_id', 'district_id', 'account_year_id', 'target_qty']:
            df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0)

        # Fetch prev season attainment
        prev_sql = f"""
            SELECT
                phdt.product_head_id,
                phdt.district_id,
                COALESCE(SUM(tsp.qty), 0) / NULLIF(phdt.target_qty, 0) * 100 AS prev_attainment_pct
            FROM product_head_district_target phdt
            JOIN mas_accountyear ay ON ay.id = phdt.account_id
            LEFT JOIN transaction_sales ts
                ON ts.account_year_id = phdt.account_id
               AND ts.company_id = %s
               AND ts.status = 5
               AND ts.deleted_by > 0
            LEFT JOIN transaction_sales_products tsp ON tsp.tran_sales_id = ts.id
            LEFT JOIN products p
                ON p.id = tsp.product_id
               AND p.product_head_id = phdt.product_head_id
            WHERE phdt.account_id = (
                SELECT id FROM mas_accountyear
                WHERE from_period < (SELECT from_period FROM mas_accountyear WHERE id = %s)
                ORDER BY from_period DESC LIMIT 1
            )
            GROUP BY phdt.product_head_id, phdt.district_id, phdt.target_qty
        """
        prev_rows = self._db.execute_query(prev_sql, (company_id, account_year_id))
        if prev_rows:
            prev_df = pd.DataFrame(prev_rows)
            for col in ['product_head_id', 'district_id', 'prev_attainment_pct']:
                prev_df[col] = pd.to_numeric(prev_df[col], errors='coerce').fillna(0)
            df = df.merge(
                prev_df[['product_head_id', 'district_id', 'prev_attainment_pct']],
                on=['product_head_id', 'district_id'],
                how='left'
            )
        else:
            df['prev_attainment_pct'] = 0.0

        df['prev_attainment_pct'] = df['prev_attainment_pct'].fillna(0)
        return df