"""Postgres loaders for forecast_engine.py.

Each loader returns a DataFrame that is structurally identical (same columns,
same dtypes) to what `pd.read_csv()` would have returned for the corresponding
CSV. The engine consumes them unchanged.

Strategy: DB first, CSV fallback. This preserves bit-exact behaviour on
systems where both are present (typical dev / QA), and lets the engine work
on DB-only deployments (production) and CSV-only deployments (Monika's
laptop without DB access).

The fallback path uses `pd.read_csv(path)` with the SAME path the engine
would have used — no changes to the engine's resolved-path logic.

Columns NOT yet in the unified schema (and how we handle them):
  - is_wholesale_spike  → computed in SQL via rolling-median window function
  - is_any_promo        → read from erp_promo_weeks (ground-truth ERP flag,
                          slightly cleaner than the statistical CSV version)
  - retail_discount_pct → defaulted to 0 in DB mode. The engine uses
                          `r.get('retail_discount_pct', 0)`, so this gracefully
                          degrades — the load_promo_data() discount-cleaning
                          branch is skipped on those rows.
  - avg_ppp_retail/webshop → defaulted to 0 in DB mode. Engine falls back to
                          `sku_prices.avg_sell_price` via get_latest_prices,
                          which we also load from DB.
"""
from __future__ import annotations

import os
import time
from datetime import datetime
from pathlib import Path
from typing import Optional

import pandas as pd

# ---------------------------------------------------------------------------
# DB connection — lazy so module import doesn't require DB
# ---------------------------------------------------------------------------

_engine = None
_engine_init_attempted = False


def _try_get_engine():
    """Return SQLAlchemy engine, or None if DB unavailable. Result cached."""
    global _engine, _engine_init_attempted
    if _engine is not None or _engine_init_attempted:
        return _engine
    _engine_init_attempted = True
    try:
        # Late import — let CLI users without sqlalchemy still use CSV path
        from sqlalchemy import create_engine, text   # noqa: F401
        from backend.config import settings
        _engine = create_engine(
            settings.DATABASE_URL,
            pool_size=2,
            max_overflow=4,
            pool_pre_ping=True,
            future=True,
        )
        # Smoke-test the connection
        with _engine.connect() as conn:
            from sqlalchemy import text as _t
            conn.execute(_t("SELECT 1"))
        return _engine
    except Exception as exc:
        print(f'  [DB] Connection unavailable, will use CSV fallback: {exc}')
        _engine = None
        return None


def db_available() -> bool:
    """True if the Postgres engine can be reached. Cached after first call."""
    return _try_get_engine() is not None


# ---------------------------------------------------------------------------
# Shared helper
# ---------------------------------------------------------------------------

def _csv_fallback(path: Optional[str], reason: str) -> Optional[pd.DataFrame]:
    if path and os.path.exists(path):
        print(f'  [DB] {reason}; falling back to {os.path.basename(path)}')
        try:
            return pd.read_csv(path)
        except Exception as exc:
            print(f'  [DB] CSV fallback for {path} also failed: {exc}')
    return None


def _df_from_query(sql: str, params: Optional[dict] = None) -> Optional[pd.DataFrame]:
    eng = _try_get_engine()
    if eng is None:
        return None
    try:
        from sqlalchemy import text
        with eng.connect() as conn:
            return pd.read_sql(text(sql), conn, params=params or {})
    except Exception as exc:
        print(f'  [DB] Query failed: {exc}')
        return None


# ===========================================================================
# Loader functions — one per CSV the engine reads
# ===========================================================================

def load_sales_clean(path: Optional[str] = None) -> pd.DataFrame:
    """Returns DataFrame with columns matching sales_clean.csv that the engine reads:
        sku, year, week, qty_total, qty_retail, qty_webshop, qty_wholesale,
        is_wholesale_spike, is_any_promo, retail_discount_pct,
        avg_ppp_retail, avg_ppp_webshop

    DB path:
      qty_*           ← v_sales_weekly
      is_any_promo    ← erp_promo_weeks (left-joined on (product, year, week))
      is_wholesale_spike ← computed in SQL via 5-week rolling median
                            (threshold matches WS_SPIKE_MULT in constants.py)
      retail_discount_pct → 0 (not yet in DB)
      avg_ppp_*           → 0 (engine falls back to erp_prices via get_latest_prices)

    CSV fallback preserves all columns bit-exactly.
    """
    # Try CSV first when it exists locally — bit-exact behaviour
    if path and os.path.exists(path):
        return pd.read_csv(path)

    # Postgres doesn't support PERCENTILE_CONT() OVER () (ordered-set aggregate
    # can't be windowed). We compute the rolling-median spike flag in pandas
    # after a simple SQL pull. This is vectorised via groupby.transform.
    sql = """
        SELECT
            p.sku,
            vsw.year::int AS year,
            vsw.week::int AS week,
            COALESCE(vsw.qty_retail, 0)::float    AS qty_retail,
            COALESCE(vsw.qty_webshop, 0)::float   AS qty_webshop,
            COALESCE(vsw.qty_wholesale, 0)::float AS qty_wholesale,
            COALESCE(vsw.qty_total, 0)::float     AS qty_total,
            CASE WHEN COALESCE(epw.is_erp_promo, false) THEN 1 ELSE 0 END AS is_any_promo
        FROM v_sales_weekly vsw
        JOIN dim_products p ON p.id = vsw.product_id
        LEFT JOIN erp_promo_weeks epw
            ON epw.product_id = vsw.product_id
            AND epw.year = vsw.year AND epw.week = vsw.week
        ORDER BY p.sku, vsw.year, vsw.week
    """
    df = _df_from_query(sql)
    if df is None:
        fb = _csv_fallback(path, "sales_clean: DB read failed")
        if fb is not None:
            return fb
        raise FileNotFoundError("sales_clean: neither DB nor CSV available")

    # is_wholesale_spike — per-SKU rolling 5-week median of non-zero wholesale weeks.
    # Matches update_sales.py's WS_SPIKE_MULT=2.5 semantics. Done in pandas
    # because Postgres can't window PERCENTILE_CONT.
    WS_MULT = 2.5
    wh = df['qty_wholesale']
    wh_pos = wh.where(wh > 0)
    # Rolling 5-prior median (exclude current week) within each SKU
    rolling_med = (
        wh_pos.groupby(df['sku'])
              .transform(lambda s: s.shift(1).rolling(window=5, min_periods=2).median())
    )
    df['is_wholesale_spike'] = (
        (wh > 0) & rolling_med.notna() & (wh > WS_MULT * rolling_med)
    ).astype(int)

    # Columns not yet in DB — engine reads via .get(col, 0)
    df['retail_discount_pct'] = 0.0
    df['avg_ppp_retail']      = 0.0
    df['avg_ppp_webshop']     = 0.0

    # Match the column order the engine expects (matters for some legacy code)
    df = df[[
        'sku', 'year', 'week',
        'qty_total', 'qty_retail', 'qty_webshop', 'qty_wholesale',
        'is_wholesale_spike', 'is_any_promo',
        'retail_discount_pct', 'avg_ppp_retail', 'avg_ppp_webshop',
    ]]
    return df


def load_erp_promo_calendar(path: Optional[str] = None) -> pd.DataFrame:
    """Columns: sku, year, week, is_erp_promo, promo_types
    Source: erp_promo_weeks JOIN dim_products.
    """
    if path and os.path.exists(path):
        return pd.read_csv(path)

    sql = """
        SELECT
            p.sku,
            epw.year::int AS year,
            epw.week::int AS week,
            CASE WHEN COALESCE(epw.is_erp_promo, false) THEN 1 ELSE 0 END AS is_erp_promo,
            COALESCE(epw.promo_types, '') AS promo_types
        FROM erp_promo_weeks epw
        JOIN dim_products p ON p.id = epw.product_id
        ORDER BY p.sku, epw.year, epw.week
    """
    df = _df_from_query(sql)
    if df is None:
        fb = _csv_fallback(path, "erp_promo_calendar: DB read failed")
        if fb is not None:
            return fb
        return pd.DataFrame(columns=['sku', 'year', 'week', 'is_erp_promo', 'promo_types'])
    return df


def load_sku_plan_list(path: Optional[str] = None) -> pd.DataFrame:
    """Columns: sku, name, cat, oznaka, vpc
    Source: sku_planning JOIN dim_products + dim_categories + erp_costs.
    The engine reads: sku_list, name, cat, oznaka (tier), vpc (cost).
    """
    if path and os.path.exists(path):
        return pd.read_csv(path)

    sql = """
        SELECT
            p.sku,
            COALESCE(p.name, '')   AS name,
            COALESCE(dc.name, '')  AS cat,
            COALESCE(sp.tier, '')  AS oznaka,
            COALESCE(ec.cost_price, 0)::float AS vpc,
            COALESCE(sp.total_xyz, '') AS total_xyz,
            COALESCE(sp.ws_xyz, '')    AS ws_xyz,
            COALESCE(sp.total_cv, 0)::float AS total_cv,
            COALESCE(sp.ws_cv, 0)::float    AS ws_cv,
            COALESCE(sp.ws_share_26w, 0)::float AS ws_share_26w
        FROM sku_planning sp
        JOIN dim_products p   ON p.id = sp.product_id
        LEFT JOIN dim_categories dc ON dc.id = p.category_id
        LEFT JOIN erp_costs ec ON ec.product_id = sp.product_id
        ORDER BY p.sku
    """
    df = _df_from_query(sql)
    if df is None:
        fb = _csv_fallback(path, "sku_plan_list: DB read failed")
        if fb is not None:
            return fb
        return pd.DataFrame(columns=['sku', 'name', 'cat', 'oznaka', 'vpc'])
    return df


def load_sku_uplift(path: Optional[str] = None) -> pd.DataFrame:
    """Columns: sku, promo_uplift
    No matching DB table. Returns CSV fallback or empty df (engine handles
    empty via cat_uplift fallback)."""
    if path and os.path.exists(path):
        return pd.read_csv(path)
    return pd.DataFrame(columns=['sku', 'promo_uplift'])


def load_cat_uplift(path: Optional[str] = None) -> pd.DataFrame:
    """Columns: cat, cat_promo_uplift
    No matching DB table. CSV fallback or empty df."""
    if path and os.path.exists(path):
        return pd.read_csv(path)
    return pd.DataFrame(columns=['cat', 'cat_promo_uplift'])


def load_sku_subcat_map(path: Optional[str] = None) -> pd.DataFrame:
    """Columns: sku, sub_cat
    Source: dim_products.subcategory.
    """
    if path and os.path.exists(path):
        return pd.read_csv(path)

    sql = """
        SELECT
            p.sku,
            COALESCE(p.subcategory, '') AS sub_cat
        FROM dim_products p
        WHERE p.active = true
        ORDER BY p.sku
    """
    df = _df_from_query(sql)
    if df is None:
        return pd.DataFrame(columns=['sku', 'sub_cat'])
    return df


def load_sku_prices(path: Optional[str] = None) -> pd.DataFrame:
    """Columns: sku, avg_sell_price, normal_retail_ppp, normal_webshop_ppp,
                qty_retail, qty_webshop, qty_wholesale, weeks_active
    Source: erp_prices JOIN dim_products.
    """
    if path and os.path.exists(path):
        return pd.read_csv(path)

    sql = """
        SELECT
            p.sku,
            COALESCE(ep.avg_sell_price, 0)::float     AS avg_sell_price,
            COALESCE(ep.normal_retail_ppp, 0)::float  AS normal_retail_ppp,
            COALESCE(ep.normal_webshop_ppp, 0)::float AS normal_webshop_ppp,
            0::int AS qty_retail,
            0::int AS qty_webshop,
            0::int AS qty_wholesale,
            COALESCE(ep.weeks_active, 0)::int AS weeks_active
        FROM dim_products p
        LEFT JOIN erp_prices ep ON ep.product_id = p.id
        WHERE p.active = true
        ORDER BY p.sku
    """
    df = _df_from_query(sql)
    if df is None:
        return pd.DataFrame(columns=[
            'sku', 'avg_sell_price', 'normal_retail_ppp', 'normal_webshop_ppp',
            'qty_retail', 'qty_webshop', 'qty_wholesale', 'weeks_active',
        ])
    return df


def load_sku_category_map(path: Optional[str] = None) -> pd.DataFrame:
    """Columns: sku, cat
    Source: dim_products JOIN dim_categories.
    """
    if path and os.path.exists(path):
        return pd.read_csv(path)

    sql = """
        SELECT
            p.sku,
            COALESCE(dc.name, 'OTHER') AS cat
        FROM dim_products p
        LEFT JOIN dim_categories dc ON dc.id = p.category_id
        WHERE p.active = true
        ORDER BY p.sku
    """
    df = _df_from_query(sql)
    if df is None:
        return pd.DataFrame(columns=['sku', 'cat'])
    return df


def load_ontop_input(channel, path: Optional[str] = None) -> pd.DataFrame:
    """Read on-top demand for a channel set, pivot to wide format with
    columns: sku, CW{n}, CW{n+1}, …

    `channel` accepts either a string ('wholesale') or a list of strings
    (['retail', 'food retail']) — the KAM/CM wizard writes 'food retail'
    for the supermarket-style MP channel while older inputs use 'retail',
    so MP must match both to avoid silently dropping rows.

    **DB is the authoritative source.** The CSV path is only used as a last-
    resort fallback when DB is unavailable.
    """
    if isinstance(channel, str):
        channels = [channel]
    else:
        channels = list(channel)

    eng = _try_get_engine()
    if path and os.path.exists(path) and eng is None:
        return pd.read_csv(path)

    sql = """
        SELECT
            p.sku,
            (oti.year_week % 100)::int AS cw,
            SUM(oti.quantity)::float AS qty
        FROM on_top_inputs oti
        JOIN dim_products p ON p.id = oti.product_id
        WHERE oti.channel = ANY(:channels)
        GROUP BY p.sku, oti.year_week
        ORDER BY p.sku, cw
    """
    long_df = _df_from_query(sql, params={"channels": channels})
    if long_df is None or long_df.empty:
        return pd.DataFrame(columns=['sku'])

    # Pivot to wide format: one row per SKU, columns CW{cw}
    wide = long_df.pivot_table(
        index='sku', columns='cw', values='qty',
        aggfunc='sum', fill_value=0,
    ).reset_index()
    wide.columns = (
        ['sku'] + [f'CW{int(c)}' for c in wide.columns[1:]]
    )
    return wide


def load_vp_input(path: Optional[str] = None) -> pd.DataFrame:
    """vp_input.csv equivalent — wholesale on-top demand, sku × CW grid."""
    return load_ontop_input('wholesale', path)


def load_mp_input(path: Optional[str] = None) -> pd.DataFrame:
    """mp_input.csv equivalent — retail on-top demand, sku × CW grid.

    Accepts both 'retail' and 'food retail' channels (the wizard writes
    'food retail' for food/supplement MP, 'retail' for other retail).
    """
    return load_ontop_input(['retail', 'food retail'], path)


def load_factor_history(path: Optional[str] = None) -> pd.DataFrame:
    """Columns: sku, run_date, run_year, run_week, target_year, target_week, factor"""
    if path and os.path.exists(path):
        return pd.read_csv(path)

    sql = """
        SELECT
            p.sku,
            COALESCE(fh.run_date::text, '') AS run_date,
            fh.run_year::int    AS run_year,
            fh.run_week::int    AS run_week,
            fh.target_year::int AS target_year,
            fh.target_week::int AS target_week,
            COALESCE(fh.factor, 1.0)::float AS factor
        FROM factor_history fh
        JOIN dim_products p ON p.id = fh.product_id
        ORDER BY fh.run_year, fh.run_week, p.sku
    """
    df = _df_from_query(sql)
    if df is None:
        return pd.DataFrame(columns=[
            'sku', 'run_date', 'run_year', 'run_week',
            'target_year', 'target_week', 'factor',
        ])
    return df


def load_latest_planner_factors() -> dict:
    """Return {sku: {target_cw: factor}} for the most recent run.

    Source priority:
      1. factor_history.csv (canonical planner-edited file, CSV-first)
      2. forecasts table in DB (fallback when CSV missing)

    factor_history.csv columns: run_date, run_year, run_week,
    target_year, target_week, sku, factor
    """
    # --- CSV-first: try data/factor_history.csv then factor_history.csv ---
    for candidate in ("data/factor_history.csv", "factor_history.csv"):
        if os.path.exists(candidate):
            try:
                df = pd.read_csv(candidate)
                if df.empty:
                    break
                # Latest (run_year, run_week)
                latest_idx = df.groupby(["run_year", "run_week"]).ngroups
                if latest_idx == 0:
                    break
                max_run = df[["run_year", "run_week"]].max()
                mask = (df["run_year"] == max_run["run_year"]) & \
                       (df["run_week"] == max_run["run_week"])
                latest = df[mask]
                out: dict = {}
                for _, r in latest.iterrows():
                    sku = str(r["sku"])
                    cw = int(r["target_week"])
                    factor = float(r["factor"])
                    out.setdefault(sku, {})[cw] = factor
                return out
            except Exception:
                break  # fall through to DB

    # --- DB fallback ---
    sql = """
        WITH latest AS (
            SELECT id FROM forecast_runs ORDER BY started_at DESC NULLS LAST, id DESC LIMIT 1
        )
        SELECT p.sku, f.week, COALESCE(f.planner_factor, 1.0)::float AS factor
        FROM forecasts f
        JOIN dim_products p ON p.id = f.product_id
        WHERE f.run_id IN (SELECT id FROM latest)
    """
    df = _df_from_query(sql)
    if df is None or df.empty:
        return {}
    out = {}
    for _, r in df.iterrows():
        out.setdefault(str(r["sku"]), {})[int(r["week"])] = float(r["factor"])
    return out


def load_latest_ontops() -> tuple[dict, dict]:
    """Return (old_vp, old_mp) dicts in the same shape as read_existing_corrections.

    The engine's signature expects:
      old_vp[sku] = {f'on-top demand_{cw}': qty, f'regular increase_{cw}': qty}
      old_mp[sku] = same shape

    We map: forecasts.on_top_wholesale → 'on-top demand_{cw}', regular_increase = 0.
    """
    sql = """
        WITH latest AS (
            SELECT id FROM forecast_runs ORDER BY started_at DESC NULLS LAST, id DESC LIMIT 1
        )
        SELECT
            p.sku, f.week,
            COALESCE(f.on_top_wholesale, 0)::float AS vp,
            COALESCE(f.on_top_retail, 0)::float    AS mp
        FROM forecasts f
        JOIN dim_products p ON p.id = f.product_id
        WHERE f.run_id IN (SELECT id FROM latest)
    """
    df = _df_from_query(sql)
    old_vp: dict = {}
    old_mp: dict = {}
    if df is None or df.empty:
        return old_vp, old_mp
    for _, r in df.iterrows():
        sku = str(r['sku'])
        cw  = int(r['week'])
        vp  = float(r['vp'])
        mp  = float(r['mp'])
        if vp > 0:
            old_vp.setdefault(sku, {})[f'on-top demand_{cw}'] = vp
            old_vp[sku][f'regular increase_{cw}'] = 0.0
        if mp > 0:
            old_mp.setdefault(sku, {})[f'on-top demand_{cw}'] = mp
            old_mp[sku][f'regular increase_{cw}'] = 0.0
    return old_vp, old_mp


# ===========================================================================
# DB WRITES — replace forecast_log.csv + factor_history.csv with direct inserts
# ===========================================================================

def write_forecast_run(
    *,
    run_type: str,
    run_by_id: Optional[int],
    year_week: int,
    n_skus: int,
    started_at: Optional[str] = None,
) -> Optional[int]:
    """Create sop_cycle (if missing) + forecast_runs row. Returns run_id, or
    None if DB unavailable."""
    eng = _try_get_engine()
    if eng is None:
        return None
    from sqlalchemy import text
    with eng.begin() as conn:
        # Ensure cycle exists
        row = conn.execute(
            text("SELECT id FROM sop_cycles WHERE year_week = :yw"),
            {"yw": year_week},
        ).first()
        if row:
            cycle_id = int(row[0])
        else:
            ins = conn.execute(
                text("""
                    INSERT INTO sop_cycles (year_week, status)
                    VALUES (:yw, 'active')
                    RETURNING id
                """),
                {"yw": year_week},
            ).first()
            cycle_id = int(ins[0]) if ins else 0

        ins = conn.execute(
            text("""
                INSERT INTO forecast_runs (cycle_id, run_by_id, run_type, n_skus, params)
                VALUES (:cid, :rid, :rtype, :n, CAST(:params AS JSONB))
                RETURNING id
            """),
            {
                "cid":    cycle_id,
                "rid":    run_by_id,
                "rtype":  run_type,
                "n":      n_skus,
                "params": '{"engine": "forecast_engine.py", "version": "3.7-db"}',
            },
        ).first()
        return int(ins[0]) if ins else None


def write_forecasts(*, run_id: int, rows: list[dict]) -> int:
    """Bulk insert into forecasts. `rows` items must have:
       sku, year, week, baseline, total, on_top_vp, on_top_mp,
       promo_uplift, planner_factor, model_used, channel_mode

    Optional per-channel split keys (added May 2026):
       forecast_retail    — engine's retail-channel forecast incl. mp on-top
       forecast_wholesale — engine's wholesale-channel forecast incl. vp on-top

    When present, these populate the per-channel columns so downstream
    (monthly plan loader, Margin Bridge) doesn't have to fall back to
    ws_share_26w. Both should sum to `total`.

    SKU→product_id is resolved in bulk.
    """
    eng = _try_get_engine()
    if eng is None or not rows:
        return 0
    from sqlalchemy import text

    skus = sorted({r['sku'] for r in rows})
    with eng.begin() as conn:
        pid_rows = conn.execute(
            text("SELECT sku, id FROM dim_products WHERE sku = ANY(:skus)"),
            {"skus": skus},
        ).fetchall()
        sku_to_pid = {r[0]: int(r[1]) for r in pid_rows}

        insert_rows = []
        for r in rows:
            pid = sku_to_pid.get(r['sku'])
            if pid is None:
                continue
            fc_r = r.get('forecast_retail')
            fc_w = r.get('forecast_wholesale')
            insert_rows.append({
                "run_id":            run_id,
                "product_id":        pid,
                "year":              int(r['year']),
                "week":              int(r['week']),
                "baseline":          float(r.get('baseline') or 0),
                "total":             float(r.get('total') or 0),
                "on_top_wholesale":  float(r.get('on_top_vp') or 0),
                "on_top_retail":     float(r.get('on_top_mp') or 0),
                "forecast_retail":    float(fc_r) if fc_r is not None else None,
                "forecast_wholesale": float(fc_w) if fc_w is not None else None,
                "promo_uplift":      float(r.get('promo_uplift') or 1.0),
                "planner_factor":    float(r.get('planner_factor') or 1.0),
                "model_used":        (r.get('model_used') or '')[:64] or None,
                "channel_mode":      (r.get('channel_mode') or '')[:64] or None,
            })

        if not insert_rows:
            return 0

        sql = text("""
            INSERT INTO forecasts
                (run_id, product_id, year, week, baseline, total,
                 on_top_wholesale, on_top_retail,
                 forecast_retail, forecast_wholesale,
                 promo_uplift, planner_factor, model_used, channel_mode)
            VALUES
                (:run_id, :product_id, :year, :week, :baseline, :total,
                 :on_top_wholesale, :on_top_retail,
                 :forecast_retail, :forecast_wholesale,
                 :promo_uplift, :planner_factor, :model_used, :channel_mode)
        """)
        n = 0
        batch = 1000
        for i in range(0, len(insert_rows), batch):
            conn.execute(sql, insert_rows[i:i + batch])
            n += min(batch, len(insert_rows) - i)
        return n


def generate_demand_plan_xlsx(run_id: Optional[int] = None) -> bytes:
    """Build a Polleo_Demand_Plan-style workbook from the forecasts table
    and return the xlsx as bytes.

    For first iteration this produces a **single-sheet** workbook ("Demand
    Planning") with the planner-facing grid that Monika needs for round-trip
    editing:

      cols:  SKU | Name | Category | Tier | model | (RR-3..RR-0 weeks) |
             baseline | (CW{w}..CW{w+12} forecasts) | planner_factor |
             on_top_vp | on_top_mp | total

    Editable cells per Monika's workflow:
      - planner_factor  → multiplied through to `total` on upload
      - on_top_vp       → adds to `total`
      - on_top_mp       → adds to `total`

    The full 7-sheet Monika xlsx (Demand Input VP/MP, Revenue Dashboard,
    Forecast Detail, Demand Output Total/On-Top, Price Reference) remains
    available via the engine's legacy `generate_xlsx=True` path. This
    function is for the **fast DB-only roundtrip** path.

    Args:
      run_id: forecasts.run_id to read. If None, picks the latest run.

    Returns:
      bytes — the xlsx content. Caller writes/streams it.
    """
    import io as _io
    from openpyxl import Workbook
    from openpyxl.styles import Font, PatternFill, Alignment

    eng = _try_get_engine()
    if eng is None:
        raise RuntimeError("Database not available — cannot generate plan xlsx")

    from sqlalchemy import text
    with eng.connect() as conn:
        if run_id is None:
            row = conn.execute(text("""
                SELECT id FROM forecast_runs
                ORDER BY started_at DESC NULLS LAST, id DESC LIMIT 1
            """)).first()
            if not row:
                raise RuntimeError("No forecast_runs found")
            run_id = int(row[0])

        # Pull forecast rows + metadata in one query
        q = """
            SELECT
                p.sku, p.name,
                COALESCE(dc.name, '') AS cat,
                COALESCE(sp.tier, '') AS tier,
                f.year, f.week,
                COALESCE(f.baseline, 0)::float       AS baseline,
                COALESCE(f.total, 0)::float          AS total,
                COALESCE(f.on_top_wholesale, 0)::float AS on_top_vp,
                COALESCE(f.on_top_retail, 0)::float    AS on_top_mp,
                COALESCE(f.planner_factor, 1.0)::float AS planner_factor,
                COALESCE(f.model_used, '') AS model_used
            FROM forecasts f
            JOIN dim_products p ON p.id = f.product_id
            LEFT JOIN dim_categories dc ON dc.id = p.category_id
            LEFT JOIN sku_planning sp ON sp.product_id = p.id
            WHERE f.run_id = :rid
            ORDER BY p.sku, f.year, f.week
        """
        df = pd.read_sql(text(q), conn, params={"rid": run_id})

    if df.empty:
        raise RuntimeError(f"No forecast rows for run_id={run_id}")

    # Pivot to one row per SKU with CW columns
    cws_sorted = sorted({(int(y), int(w)) for y, w in zip(df['year'], df['week'])},
                         key=lambda yw: yw[0] * 100 + yw[1])
    cw_labels = [f'CW{w}' for (_y, w) in cws_sorted]

    # Per-SKU aggregate row
    sku_meta = (df.groupby('sku')
                  .agg(name=('name', 'first'),
                        cat=('cat', 'first'),
                        tier=('tier', 'first'),
                        model=('model_used', 'first'))
                  .reset_index())

    # Build sheet
    wb = Workbook()
    ws = wb.active
    ws.title = 'Demand Planning'

    HDR_FILL = PatternFill('solid', fgColor='1F4E79')
    FC_FILL  = PatternFill('solid', fgColor='E8F4E8')   # forecast cells
    PF_FILL  = PatternFill('solid', fgColor='FFF0F0')   # planner_factor (editable)
    OT_FILL  = PatternFill('solid', fgColor='FFFAE6')   # on-top (editable)
    HDR_FONT = Font(color='FFFFFF', bold=True)
    CENTER   = Alignment(horizontal='center', vertical='center')

    # Header row
    headers = ['SKU', 'Name', 'Category', 'Tier', 'Model'] + cw_labels + [
        'Planner Factor', 'On-Top VP', 'On-Top MP', 'Total',
    ]
    for c, h in enumerate(headers, 1):
        cell = ws.cell(1, c, h)
        cell.fill = HDR_FILL
        cell.font = HDR_FONT
        cell.alignment = CENTER

    # Build pivot
    fc_pivot = df.pivot_table(
        index='sku', columns=['year', 'week'], values='total',
        aggfunc='first',
    )
    pf_pivot = df.pivot_table(
        index='sku', columns=['year', 'week'], values='planner_factor',
        aggfunc='first',
    )
    vp_pivot = df.pivot_table(
        index='sku', columns=['year', 'week'], values='on_top_vp',
        aggfunc='first',
    )
    mp_pivot = df.pivot_table(
        index='sku', columns=['year', 'week'], values='on_top_mp',
        aggfunc='first',
    )

    row = 2
    for _, m in sku_meta.iterrows():
        sku = m['sku']
        ws.cell(row, 1, sku)
        ws.cell(row, 2, m['name'])
        ws.cell(row, 3, m['cat'])
        ws.cell(row, 4, m['tier'])
        ws.cell(row, 5, m['model'])

        # CW forecast columns
        for j, yw in enumerate(cws_sorted):
            v = fc_pivot.loc[sku, yw] if sku in fc_pivot.index and yw in fc_pivot.columns else 0
            cell = ws.cell(row, 6 + j, float(v) if pd.notna(v) else 0)
            cell.fill = FC_FILL
            cell.number_format = '#,##0'

        # Aggregate scalars — use first non-1 planner_factor or first row's VP/MP
        # (these are typically constant per SKU, edited in bulk by planner)
        first_yw = cws_sorted[0] if cws_sorted else None
        if first_yw and sku in pf_pivot.index:
            pf_val = float(pf_pivot.loc[sku, first_yw]) if first_yw in pf_pivot.columns else 1.0
            vp_val = float(vp_pivot.loc[sku, first_yw]) if first_yw in vp_pivot.columns and sku in vp_pivot.index else 0.0
            mp_val = float(mp_pivot.loc[sku, first_yw]) if first_yw in mp_pivot.columns and sku in mp_pivot.index else 0.0
        else:
            pf_val, vp_val, mp_val = 1.0, 0.0, 0.0

        n_cw = len(cws_sorted)
        pf_cell = ws.cell(row, 6 + n_cw, round(pf_val, 4))
        pf_cell.fill = PF_FILL
        pf_cell.number_format = '0.0000'

        vp_cell = ws.cell(row, 7 + n_cw, vp_val)
        vp_cell.fill = OT_FILL
        vp_cell.number_format = '#,##0'

        mp_cell = ws.cell(row, 8 + n_cw, mp_val)
        mp_cell.fill = OT_FILL
        mp_cell.number_format = '#,##0'

        # Total (sum of CW columns)
        total_val = sum(float(fc_pivot.loc[sku, yw]) if sku in fc_pivot.index
                         and yw in fc_pivot.columns and pd.notna(fc_pivot.loc[sku, yw])
                         else 0 for yw in cws_sorted)
        tcell = ws.cell(row, 9 + n_cw, total_val)
        tcell.font = Font(bold=True)
        tcell.number_format = '#,##0'

        row += 1

    # Column widths
    ws.column_dimensions['A'].width = 16
    ws.column_dimensions['B'].width = 36
    ws.column_dimensions['C'].width = 16
    ws.column_dimensions['D'].width = 10
    ws.column_dimensions['E'].width = 12
    for j in range(len(cw_labels)):
        ws.column_dimensions[chr(ord('F') + j)].width = 10
    ws.freeze_panes = 'F2'
    ws.auto_filter.ref = ws.dimensions

    # Embed metadata in a hidden second sheet so the upload endpoint can
    # validate which run_id this download came from.
    meta = wb.create_sheet('_meta')
    meta.cell(1, 1, 'run_id'); meta.cell(1, 2, int(run_id))
    meta.cell(2, 1, 'generated_at'); meta.cell(2, 2, datetime.now().isoformat())
    meta.sheet_state = 'hidden'

    buf = _io.BytesIO()
    wb.save(buf)
    return buf.getvalue()


def parse_uploaded_plan(file_bytes: bytes) -> dict:
    """Parse an edited Polleo_Demand_Plan.xlsx that was downloaded via
    generate_demand_plan_xlsx(). Reads the 'Demand Planning' sheet and
    extracts the per-SKU planner_factor + on_top_vp + on_top_mp columns.

    Returns:
      {
        'run_id': int (from _meta sheet, or None),
        'updates': list of dicts {sku, planner_factor, on_top_vp, on_top_mp},
        'warnings': list of strings,
      }
    """
    import io as _io
    from openpyxl import load_workbook

    out = {'run_id': None, 'updates': [], 'warnings': []}
    wb = load_workbook(_io.BytesIO(file_bytes), data_only=True)

    # Read embedded run_id metadata
    if '_meta' in wb.sheetnames:
        ms = wb['_meta']
        try:
            for row in ms.iter_rows(min_row=1, max_row=5, values_only=True):
                if row and row[0] == 'run_id':
                    out['run_id'] = int(row[1])
                    break
        except Exception:
            pass

    # Find the Demand Planning sheet
    ws = None
    for name in ('Demand Planning', 'Sheet1'):
        if name in wb.sheetnames:
            ws = wb[name]
            break
    if ws is None:
        ws = wb[wb.sheetnames[0]]

    # Read header row to locate columns
    headers = [str(c.value or '').strip() for c in ws[1]]

    def col_idx(name):
        for i, h in enumerate(headers, 1):
            if h.lower() == name.lower():
                return i
        return None

    col_sku = col_idx('SKU')
    col_pf  = col_idx('Planner Factor')
    col_vp  = col_idx('On-Top VP')
    col_mp  = col_idx('On-Top MP')
    if col_sku is None or col_pf is None:
        out['warnings'].append(
            'Could not find required columns "SKU" and "Planner Factor" in uploaded sheet. '
            f'Headers found: {headers}'
        )
        return out

    for r in range(2, ws.max_row + 1):
        sku = ws.cell(r, col_sku).value
        if not sku:
            continue
        pf = ws.cell(r, col_pf).value if col_pf else 1.0
        vp = ws.cell(r, col_vp).value if col_vp else 0
        mp = ws.cell(r, col_mp).value if col_mp else 0
        try:
            pf = float(pf) if pf is not None else 1.0
            vp = float(vp) if vp is not None else 0.0
            mp = float(mp) if mp is not None else 0.0
        except (TypeError, ValueError):
            out['warnings'].append(f'SKU {sku}: non-numeric values; skipped')
            continue
        out['updates'].append({
            'sku': str(sku).strip(),
            'planner_factor': pf,
            'on_top_vp': vp,
            'on_top_mp': mp,
        })
    return out


def apply_plan_updates(run_id: int, updates: list[dict]) -> dict:
    """Apply updates from a parsed plan upload to the forecasts table.

    For each (sku, planner_factor, on_top_vp, on_top_mp) update:
      UPDATE forecasts
         SET planner_factor = pf,
             on_top_wholesale = vp,
             on_top_retail    = mp,
             total = baseline * pf + vp + mp
       WHERE run_id = :rid AND product_id = (resolved from sku)

    Returns: {n_skus_updated, n_rows_updated, factors_changed: [...]}
    """
    eng = _try_get_engine()
    if eng is None:
        raise RuntimeError("Database not available — cannot apply plan updates")
    from sqlalchemy import text

    skus = sorted({u['sku'] for u in updates})
    with eng.begin() as conn:
        # Resolve SKU→product_id
        pid_rows = conn.execute(
            text("SELECT sku, id FROM dim_products WHERE sku = ANY(:skus)"),
            {"skus": skus},
        ).fetchall()
        sku_to_pid = {r[0]: int(r[1]) for r in pid_rows}

        # Read existing planner_factor per SKU before update (for the diff)
        before = conn.execute(text("""
            SELECT p.sku, AVG(COALESCE(f.planner_factor, 1.0))::float AS pf
            FROM forecasts f JOIN dim_products p ON p.id = f.product_id
            WHERE f.run_id = :rid AND p.sku = ANY(:skus)
            GROUP BY p.sku
        """), {"rid": run_id, "skus": skus}).fetchall()
        before_map = {r[0]: float(r[1]) for r in before}

        n_skus_updated = 0
        n_rows_updated = 0
        factors_changed = []
        for u in updates:
            pid = sku_to_pid.get(u['sku'])
            if pid is None:
                continue
            res = conn.execute(text("""
                UPDATE forecasts
                   SET planner_factor   = :pf,
                       on_top_wholesale = :vp,
                       on_top_retail    = :mp,
                       total            = COALESCE(baseline, 0) * :pf + :vp + :mp
                 WHERE run_id = :rid AND product_id = :pid
            """), {
                "pf":  u['planner_factor'],
                "vp":  u['on_top_vp'],
                "mp":  u['on_top_mp'],
                "rid": run_id,
                "pid": pid,
            })
            if (res.rowcount or 0) > 0:
                n_skus_updated += 1
                n_rows_updated += int(res.rowcount)
                old_pf = before_map.get(u['sku'])
                if old_pf is not None and abs(old_pf - u['planner_factor']) > 0.0001:
                    factors_changed.append({
                        'sku':        u['sku'],
                        'old_factor': round(old_pf, 4),
                        'new_factor': round(u['planner_factor'], 4),
                    })

        # Audit log entry
        try:
            conn.execute(text("""
                INSERT INTO audit_log (action, table_name, record_id, changes)
                VALUES ('plan_upload', 'forecasts', :rid::text, CAST(:changes AS JSONB))
            """), {
                "rid":     str(run_id),
                "changes": pd.io.json.dumps({
                    "n_skus_updated":  n_skus_updated,
                    "n_rows_updated":  n_rows_updated,
                    "factors_changed": factors_changed[:50],
                }) if hasattr(pd.io, 'json') else
                    '{"n_skus_updated": ' + str(n_skus_updated) + '}',
            })
        except Exception:
            pass

    return {
        'n_skus_updated':  n_skus_updated,
        'n_rows_updated':  n_rows_updated,
        'factors_changed': factors_changed,
    }


def write_factor_history(
    *, run_year: int, run_week: int, rows: list[dict],
    run_date: Optional[str] = None,
) -> int:
    """Insert factor_history rows. Idempotent on (run_year, run_week)."""
    eng = _try_get_engine()
    if eng is None or not rows:
        return 0
    from sqlalchemy import text
    from datetime import datetime

    skus = sorted({r['sku'] for r in rows})
    with eng.begin() as conn:
        pid_rows = conn.execute(
            text("SELECT sku, id FROM dim_products WHERE sku = ANY(:skus)"),
            {"skus": skus},
        ).fetchall()
        sku_to_pid = {r[0]: int(r[1]) for r in pid_rows}

        # Idempotent: delete same-run rows first
        conn.execute(
            text("DELETE FROM factor_history WHERE run_year = :y AND run_week = :w"),
            {"y": run_year, "w": run_week},
        )

        if run_date:
            try:
                rd = datetime.strptime(run_date, "%Y-%m-%d %H:%M")
            except Exception:
                rd = datetime.now()
        else:
            rd = datetime.now()

        insert_rows = []
        for r in rows:
            pid = sku_to_pid.get(r['sku'])
            if pid is None:
                continue
            insert_rows.append({
                "product_id":  pid,
                "run_year":    int(r['run_year']),
                "run_week":    int(r['run_week']),
                "target_year": int(r['target_year']),
                "target_week": int(r['target_week']),
                "factor":      float(r.get('factor') or 1.0),
                "run_date":    rd,
            })
        if not insert_rows:
            return 0

        sql = text("""
            INSERT INTO factor_history
                (product_id, run_year, run_week, target_year, target_week, factor, run_date)
            VALUES
                (:product_id, :run_year, :run_week, :target_year, :target_week, :factor, :run_date)
        """)
        n = 0
        batch = 1000
        for i in range(0, len(insert_rows), batch):
            conn.execute(sql, insert_rows[i:i + batch])
            n += min(batch, len(insert_rows) - i)
        return n