"""DB loaders for the v4 engine (spec §2) — pull a SKU's granular weekly series
(retail/web per region, wholesale per buyer) from erp_transactions, aligned to a
common week grid."""
from __future__ import annotations

import numpy as np
import pandas as pd
from sqlalchemy import text

from backend.models.database import engine


def load_sku(pid: int, max_yw: int = 999999) -> dict | None:
    """Return aligned weekly series for one SKU up to max_yw (inclusive):
        weeks : list of yyyyww
        iso   : ISO week number per grid point (for seasonality)
        retail/web : {region: (qty[T], price[T])}
        wholesale  : {buyer_id: qty[T]}
    """
    with engine.connect() as conn:
        df = pd.read_sql(text("""
            SELECT EXTRACT(isoyear FROM transaction_date)::int*100
                     + EXTRACT(week FROM transaction_date)::int AS yw,
                   EXTRACT(week FROM transaction_date)::int AS wk,
                   cm.channel AS ch,
                   CASE WHEN cm.channel IN ('retail','webshop')
                        THEN COALESCE(ds.country,'UNK') ELSE 'ws' END AS region,
                   et.partner_id AS bpid, et.quantity AS q, et.total_value AS tv
            FROM erp_transactions et
            JOIN lookup_channel_map cm ON cm.id = et.channel_map_id
            LEFT JOIN dim_stores ds ON ds.id = et.store_id
            WHERE et.product_id = :pid
              AND (EXTRACT(isoyear FROM transaction_date)::int*100
                     + EXTRACT(week FROM transaction_date)::int) <= :mx
        """), conn, params={"pid": pid, "mx": max_yw})
        # ERP promo calendar (retail/web ground truth), keyed by product_id.
        prom = pd.read_sql(text("""
            SELECT year*100 + week AS yw, is_erp_promo
            FROM erp_promo_weeks
            WHERE product_id = :pid AND (year*100 + week) <= :mx
        """), conn, params={"pid": pid, "mx": max_yw})
    if df.empty:
        return None
    covered_yws = set(int(w) for w in prom["yw"]) if not prom.empty else set()
    promo_yws = set(int(w) for w in prom.loc[prom["is_erp_promo"] == True, "yw"]) if not prom.empty else set()

    weeks = sorted(df.yw.unique())
    idx = {w: i for i, w in enumerate(weeks)}
    iso = [int(df[df.yw == w].wk.iloc[0]) for w in weeks]

    def grid(sub, col="q"):
        a = np.zeros(len(weeks))
        for w, v in sub.groupby("yw")[col].sum().items():
            a[idx[w]] = v
        return a

    retail, web = {}, {}
    for ch, store in (("retail", retail), ("webshop", web)):
        sub = df[df.ch == ch]
        for reg in sub.region.unique():
            s = sub[sub.region == reg]
            qa, tva = grid(s), grid(s, "tv")
            price = np.divide(tva, qa, out=np.zeros_like(qa), where=qa > 0)
            store[reg] = (qa, price)

    wholesale = {}
    ws = df[df.ch == "wholesale"]
    for b in ws.bpid.dropna().unique():
        wholesale[int(b)] = grid(ws[ws.bpid == b])

    return {"weeks": weeks, "iso": iso, "retail": retail, "web": web,
            "wholesale": wholesale,
            "promo_yws": promo_yws, "covered_yws": covered_yws}


def next_iso_week(last_yw: int) -> int:
    """ISO week number following the last data week (1..52)."""
    wk = last_yw % 100
    return (wk % 52) + 1


def horizon_weeks(last_yw: int, h: int) -> list[tuple[int, int]]:
    """The h (year, week) pairs following last_yw (ISO weeks 1..52)."""
    y, w = last_yw // 100, last_yw % 100
    out = []
    for _ in range(h):
        w += 1
        if w > 52:
            w, y = 1, y + 1
        out.append((y, w))
    return out


def load_store_retail(pid: int, max_yw: int = 999999):
    """Per-store retail weekly series for a SKU (real stores only; warehouse
    excluded). Returns (weeks, {store_id: {'name','region','qty'[T]}})."""
    with engine.connect() as conn:
        df = pd.read_sql(text("""
            SELECT EXTRACT(isoyear FROM transaction_date)::int*100
                     + EXTRACT(week FROM transaction_date)::int AS yw,
                   et.store_id AS sid, ds.name AS sname,
                   COALESCE(ds.country,'UNK') AS region, SUM(et.quantity) AS q
            FROM erp_transactions et
            JOIN lookup_channel_map cm ON cm.id = et.channel_map_id
            JOIN dim_stores ds ON ds.id = et.store_id
            WHERE et.product_id = :pid AND cm.channel = 'retail'
              AND ds.is_warehouse IS NOT TRUE
              AND NOT (ds.name ILIKE ANY (ARRAY[
                    '%warehouse%','%transfer%','%proizvodnja%','%testni%',
                    '%FBA%','%oštećena%','%aggregate%']))
              AND (EXTRACT(isoyear FROM transaction_date)::int*100
                     + EXTRACT(week FROM transaction_date)::int) <= :mx
            GROUP BY 1, 2, 3, 4
        """), conn, params={"pid": pid, "mx": max_yw})
    if df.empty:
        return [], {}
    weeks = sorted(df.yw.unique())
    idx = {w: i for i, w in enumerate(weeks)}
    stores: dict[int, dict] = {}
    for sid, g in df.groupby("sid"):
        a = np.zeros(len(weeks))
        for w, v in g.groupby("yw")["q"].sum().items():
            a[idx[w]] = v
        stores[int(sid)] = {"name": g["sname"].iloc[0],
                            "region": g["region"].iloc[0], "qty": a}
    return weeks, stores