"""Populate all dimension tables from existing CSV / JSON sources.

Idempotent: TRUNCATEs every target table first (with CASCADE so any
downstream tables that have been populated in earlier work get cleaned
out too). Loads in strict FK-dependency order:

    dim_suppliers
    dim_categories
    users
    product_families         (intentionally skipped — see note below)
    dim_products
    dim_stores
    dim_partners
    sku_planning

Why product_families is skipped: resolution lives in
PromoTool/parent_map.py and is XLSX-driven (Polleo Help svi artikli.xlsx
keyed on Naziv) with a heuristic fallback for SKUs absent from the
spreadsheet. Wiring that into a Postgres migrator is non-trivial and not
needed yet — dim_products.family_id stays NULL until a dedicated
migrate_families.py exists.

Source files (read first, never mutated):
    data/sku_category_map.csv         sku, name, cat
    data/sku_subcat_map.csv           sku, name, sub_cat, grup
    data/sku_plan_list.csv            sku, name, cat, oznaka, vpc,
                                       ws_xyz, ws_cv, ws_nz_weeks,
                                       total_xyz, total_cv,
                                       total_nz_weeks, ws_share_26w
    data/sales_detailed.csv           31 cols — used for suppliers
                                       (proizvodjac), stores (jedinica,
                                       naziv_jedinice, source_country),
                                       partners (partner, naziv_partnera,
                                       drzava)
    data/kam_cm_config.json           kam_cm_config.<username> blocks

Usage from project root:
    python -m db.migrate_dimensions
"""
from __future__ import annotations

import json
from pathlib import Path

import pandas as pd
import psycopg2.extras

from db.connection import get_connection, get_engine

PROJECT_ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = PROJECT_ROOT / "data"

# sales_detailed.csv `source_country` tag → ISO 3166-1 alpha-2 for dim_stores.country.
# `drzava` on partners is already ISO ("SI", "HR", "AT") — no mapping needed.
_COUNTRY_MAP = {"cro": "HR", "slo": "SI", "at": "AT"}


def _truncate_all(conn) -> None:
    """Wipe every table this migrator owns. CASCADE clears downstream
    fact tables (forecasts, sku_planning, on_top_inputs, ...) that would
    hold dangling FK refs after id reassignment."""
    tables = [
        "sku_planning",
        "dim_partners",
        "dim_stores",
        "dim_products",
        "product_families",
        "dim_categories",
        "dim_suppliers",
        "users",
    ]
    with conn.cursor() as cur:
        cur.execute(
            "TRUNCATE TABLE " + ", ".join(tables) + " RESTART IDENTITY CASCADE"
        )


def _load_suppliers(conn) -> int:
    """Unique `proizvodjac` strings from sales_detailed.csv. No proper
    supplier code exists in source data — use the manufacturer name as
    both the unique business code (truncated to 64 chars) and the name.
    Truncation can produce duplicates only if two manufacturers share
    the first 64 chars (highly unlikely); deduped defensively below."""
    fp = DATA_DIR / "sales_detailed.csv"
    if not fp.exists():
        print("  [warn] sales_detailed.csv missing — dim_suppliers left empty")
        return 0
    df = pd.read_csv(fp, usecols=["proizvodjac"], dtype=str, encoding="utf-8")
    df["proizvodjac"] = df["proizvodjac"].fillna("").str.strip()
    names = sorted({n for n in df["proizvodjac"] if n})
    seen_codes: set[str] = set()
    rows: list[tuple] = []
    for n in names:
        code = n[:64]
        if code in seen_codes:
            continue
        seen_codes.add(code)
        rows.append((code, n))
    with conn.cursor() as cur:
        cur.executemany(
            "INSERT INTO dim_suppliers (code, name) VALUES (%s, %s)",
            rows,
        )
    return len(rows)


def _load_categories(conn) -> int:
    fp = DATA_DIR / "sku_category_map.csv"
    df = pd.read_csv(fp, usecols=["cat"], dtype=str, encoding="utf-8")
    df["cat"] = df["cat"].fillna("").str.strip()
    names = sorted({c for c in df["cat"] if c})
    with conn.cursor() as cur:
        cur.executemany(
            "INSERT INTO dim_categories (name) VALUES (%s)",
            [(n,) for n in names],
        )
    return len(names)


def _load_users(conn) -> int:
    """KAM/CM rows from kam_cm_config.json + two hardcoded DP admins
    (monika, lovro). password_hash NULL — auth comes later."""
    fp = DATA_DIR / "kam_cm_config.json"
    if fp.exists():
        cfg_users = json.loads(fp.read_text(encoding="utf-8")).get(
            "kam_cm_config", {}
        )
    else:
        print("  [warn] kam_cm_config.json missing — only hardcoded DPs loaded")
        cfg_users = {}

    rows: list[tuple] = []
    for username, info in cfg_users.items():
        cats = info.get("categories")
        buyers = info.get("buyers")
        rows.append((
            username,
            info.get("display_name"),
            None,                                      # password_hash
            info.get("role"),
            info.get("type"),
            info.get("channel"),
            psycopg2.extras.Json(cats) if cats is not None else None,
            psycopg2.extras.Json(buyers) if buyers else None,
            info.get("slack_user_id"),
            True,                                       # active
        ))
    # Hardcoded DP admins. Same shape, no Slack id, NULL JSONB cols.
    for username, display in (("monika", "Monika"), ("lovro", "Lovro")):
        rows.append((
            username, display, None, "admin", "DP",
            None, None, None, None, True,
        ))

    with conn.cursor() as cur:
        cur.executemany(
            """
            INSERT INTO users (
                username, display_name, password_hash, role, type,
                channel, categories, buyers, slack_user_id, active
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """,
            rows,
        )
    return len(rows)


def _load_products(engine) -> int:
    """All ~7,600 SKUs from sku_category_map.csv, enriched with sub_cat
    from sku_subcat_map.csv. category_id resolved via name lookup against
    the freshly-populated dim_categories. family_id stays NULL."""
    cat_map_fp = DATA_DIR / "sku_category_map.csv"
    subcat_fp = DATA_DIR / "sku_subcat_map.csv"

    df = pd.read_csv(cat_map_fp, dtype=str, encoding="utf-8")
    df["sku"] = df["sku"].fillna("").str.strip()
    df["name"] = df["name"].fillna("").str.strip()
    df["cat"] = df["cat"].fillna("").str.strip()
    df = df[df["sku"] != ""].drop_duplicates(subset=["sku"], keep="first")

    if subcat_fp.exists():
        sub = pd.read_csv(subcat_fp, dtype=str, encoding="utf-8")[
            ["sku", "sub_cat"]
        ]
        sub["sku"] = sub["sku"].fillna("").str.strip()
        sub["sub_cat"] = sub["sub_cat"].fillna("").str.strip()
        sub = sub[sub["sku"] != ""].drop_duplicates(
            subset=["sku"], keep="first"
        )
        df = df.merge(sub, on="sku", how="left")
    else:
        print("  [warn] sku_subcat_map.csv missing — subcategory left NULL")
        df["sub_cat"] = ""

    cat_rows = pd.read_sql("SELECT id, name FROM dim_categories", engine)
    cat_to_id = dict(zip(cat_rows["name"], cat_rows["id"]))
    df["category_id"] = df["cat"].map(cat_to_id)
    df.loc[df["cat"] == "", "category_id"] = pd.NA

    out = pd.DataFrame({
        "sku":            df["sku"],
        "name":           df["name"].where(df["name"] != "", None),
        "category_id":    df["category_id"].astype("Int64"),
        "subcategory":    df["sub_cat"].where(df["sub_cat"] != "", None),
        "brand":          None,
        "supplier_id":    pd.array([pd.NA] * len(df), dtype="Int64"),
        "family_id":      pd.array([pd.NA] * len(df), dtype="Int64"),
        "alternate_code": None,
        "flavor_color":   None,
        "size":           None,
        "active":         True,
    })
    out.to_sql(
        "dim_products", engine, if_exists="append",
        index=False, chunksize=2000, method="multi",
    )
    return len(out)


def _load_stores(conn) -> int:
    fp = DATA_DIR / "sales_detailed.csv"
    if not fp.exists():
        print("  [warn] sales_detailed.csv missing — dim_stores left empty")
        return 0
    df = pd.read_csv(
        fp,
        usecols=["jedinica", "naziv_jedinice", "source_country"],
        dtype=str, encoding="utf-8",
    )
    df["jedinica"] = df["jedinica"].fillna("").str.strip()
    df["naziv_jedinice"] = df["naziv_jedinice"].fillna("").str.strip()
    df["source_country"] = df["source_country"].fillna("").str.strip().str.lower()
    df = df[df["jedinica"] != ""]
    # Pick the most common (name, country) per unit_code — guards against
    # one-off rows with a typoed name.
    grouped = (
        df.groupby("jedinica")
          .agg(
              name=("naziv_jedinice",
                    lambda s: s.mode().iat[0] if len(s.mode()) else ""),
              source_country=("source_country",
                              lambda s: s.mode().iat[0] if len(s.mode()) else ""),
          )
          .reset_index()
    )
    rows = []
    for _, r in grouped.iterrows():
        unit = r["jedinica"]
        rows.append((
            unit,
            (r["name"] or None),
            _COUNTRY_MAP.get(r["source_country"]),
            unit == "01",   # is_warehouse — unit_code '01' is the central warehouse
        ))
    with conn.cursor() as cur:
        cur.executemany(
            """
            INSERT INTO dim_stores (unit_code, name, country, is_warehouse)
            VALUES (%s, %s, %s, %s)
            """,
            rows,
        )
    return len(rows)


def _load_partners(conn) -> int:
    fp = DATA_DIR / "sales_detailed.csv"
    if not fp.exists():
        print("  [warn] sales_detailed.csv missing — dim_partners left empty")
        return 0
    df = pd.read_csv(
        fp,
        usecols=["partner", "naziv_partnera", "drzava"],
        dtype=str, encoding="utf-8",
    )
    df["partner"] = df["partner"].fillna("").str.strip()
    df["naziv_partnera"] = df["naziv_partnera"].fillna("").str.strip()
    df["drzava"] = df["drzava"].fillna("").str.strip()
    df = df[df["partner"] != ""]
    grouped = (
        df.groupby("partner")
          .agg(
              name=("naziv_partnera",
                    lambda s: s.mode().iat[0] if len(s.mode()) else ""),
              country=("drzava",
                       lambda s: s.mode().iat[0] if len(s.mode()) else ""),
          )
          .reset_index()
    )
    rows = [
        (
            r["partner"],
            (r["name"] or None),
            None,                         # partner_type — not in source
            (r["country"] or None),
        )
        for _, r in grouped.iterrows()
    ]
    with conn.cursor() as cur:
        cur.executemany(
            """
            INSERT INTO dim_partners (code, name, partner_type, country)
            VALUES (%s, %s, %s, %s)
            """,
            rows,
        )
    return len(rows)


def _load_sku_planning(engine) -> int:
    """sku_plan_list.csv → sku_planning. SKUs missing from dim_products
    are dropped with a warning (a row in sku_plan_list with no master
    product would have nowhere to point its FK).

    Note: the spec said `vpc → cost_price`, but the schema has a `vpc`
    column and no `cost_price` column on sku_planning. Mapped to `vpc`
    here; flag if cost_price is wanted instead.
    """
    fp = DATA_DIR / "sku_plan_list.csv"
    if not fp.exists():
        print("  [warn] sku_plan_list.csv missing — sku_planning left empty")
        return 0

    df = pd.read_csv(fp, encoding="utf-8")
    df["sku"] = df["sku"].astype(str).str.strip()

    prod = pd.read_sql("SELECT id, sku FROM dim_products", engine)
    sku_to_id = dict(zip(prod["sku"], prod["id"]))
    df["product_id"] = df["sku"].map(sku_to_id)

    missing = int(df["product_id"].isna().sum())
    if missing:
        print(
            f"  [warn] {missing} sku_plan_list rows had no matching "
            f"dim_products row — skipped"
        )
    df = df[df["product_id"].notna()].copy()
    df["product_id"] = df["product_id"].astype(int)

    out = pd.DataFrame({
        "product_id":     df["product_id"],
        "tier":           df.get("oznaka"),
        "ws_xyz":         df.get("ws_xyz"),
        "total_xyz":      df.get("total_xyz"),
        "ws_cv":          pd.to_numeric(df.get("ws_cv"), errors="coerce"),
        "total_cv":       pd.to_numeric(df.get("total_cv"), errors="coerce"),
        "ws_nz_weeks":    pd.to_numeric(df.get("ws_nz_weeks"),
                                          errors="coerce").astype("Int64"),
        "total_nz_weeks": pd.to_numeric(df.get("total_nz_weeks"),
                                          errors="coerce").astype("Int64"),
        "ws_share_26w":   pd.to_numeric(df.get("ws_share_26w"),
                                          errors="coerce"),
        "vpc":            pd.to_numeric(df.get("vpc"), errors="coerce"),
    })
    out.to_sql(
        "sku_planning", engine, if_exists="append",
        index=False, chunksize=1000, method="multi",
    )
    return len(out)


def migrate_dimensions() -> None:
    print("Connecting to Postgres...")
    conn = get_connection()
    conn.autocommit = False
    engine = get_engine()
    try:
        print("Truncating dimension tables (CASCADE)...")
        _truncate_all(conn)
        conn.commit()

        # Each loader gets its own try/commit/rollback so a partial run
        # leaves prior steps in place — easier to iterate while developing.
        steps: list[tuple[str, callable]] = [
            ("dim_suppliers",   lambda: _load_suppliers(conn)),
            ("dim_categories",  lambda: _load_categories(conn)),
            ("users",           lambda: _load_users(conn)),
            # product_families intentionally skipped — see module docstring.
            ("dim_products",    lambda: _load_products(engine)),
            ("dim_stores",      lambda: _load_stores(conn)),
            ("dim_partners",    lambda: _load_partners(conn)),
            ("sku_planning",    lambda: _load_sku_planning(engine)),
        ]
        for label, fn in steps:
            try:
                n = fn()
                conn.commit()
                print(f"  {label:18s} {n:>7,} rows")
            except Exception:
                conn.rollback()
                raise
    finally:
        conn.close()
    print()
    print("Dimensions loaded.")


if __name__ == "__main__":
    migrate_dimensions()