"""Build the SFT dataset for the MiniCPM4.1-8B recipe planner.

Reads the Kaggle "better-recipes-for-a-better-life" dataset and produces
supervised fine-tuning pairs for BOTH planner tasks, matching the exact
prompt formats the app uses (src/prompts/planner_propose.txt and
planner_recipe.txt):

  1. propose  : ingredients -> {"options": [{name, why} x3]}
  2. recipe   : dish + ingredients -> {"name", "cuisine", "servings",
                "total_time_minutes", "final_dish_visual", "steps":[...]}

Run locally (once) before fine-tuning:
    python scripts/build_recipe_dataset.py

Requires:
    pip install kagglehub pandas pyarrow datasets huggingface_hub tqdm
    ~/.kaggle/kaggle.json with your credentials
"""
from __future__ import annotations

import json
import random
import re
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

import pandas as pd
from tqdm import tqdm

from src import config

random.seed(42)

HF_DATASET_REPO = "eldinosaur/cook-with-me-recipes-sft"

# ---------------------------------------------------------------------------
# 1. Download (use ONLY recipes.csv — test_recipes.csv has a different schema
#    whose capitalized columns shadowed the real data in the old version)
# ---------------------------------------------------------------------------

print("Pulling Kaggle dataset…")
import kagglehub

raw_path = Path(kagglehub.dataset_download(config.KAGGLE_DATASET))
main_csv = raw_path / "recipes.csv"
print(f"Reading {main_csv}")

# cp1252 decodes the fraction/symbol bytes that show up as � under utf-8
try:
    raw_df = pd.read_csv(main_csv, encoding="cp1252", on_bad_lines="skip")
except Exception:
    raw_df = pd.read_csv(main_csv, encoding="utf-8", on_bad_lines="skip")

print(f"Rows: {len(raw_df)}  columns: {list(raw_df.columns)}")


# ---------------------------------------------------------------------------
# 2. Cleaning helpers
# ---------------------------------------------------------------------------

_UNIT = (
    r"(cups?|tablespoons?|tbsps?|teaspoons?|tsps?|pounds?|lbs?|ounces?|ozs?|"
    r"grams?|kgs?|mls?|liters?|pinch(?:es)?|dash(?:es)?|cloves?|cans?|"
    r"packages?|pkgs?|sheets?|slices?|sticks?|quarts?|pints?|jars?|bunch(?:es)?|"
    r"heads?|stalks?|sprigs?|pieces?|fillets?)"
)
_PREP_WORDS = {
    "peeled", "chopped", "diced", "sliced", "minced", "cored", "thawed",
    "drained", "rinsed", "softened", "melted", "beaten", "divided", "cubed",
    "to taste", "optional", "or more", "plus more", "for garnish", "for serving",
    "lightly beaten", "room temperature", "at room temperature", "finely chopped",
    "thinly sliced", "cut into", "more", "and", "or other", "such as",
}


def _clean_text(val: str) -> str:
    if not isinstance(val, str):
        return ""
    # drop any remaining replacement chars and collapse whitespace
    val = val.replace("�", " ")
    return re.sub(r"[ \t]+", " ", val).strip()


def _simplify_ingredient(raw: str) -> str:
    s = re.sub(r"\([^)]*\)", "", raw)             # remove parentheticals
    s = _clean_text(s).lower()
    s = re.sub(r"^[\d\s./¼½¾⅓⅔⅛+-]+", "", s)       # leading quantities
    s = re.sub(rf"^{_UNIT}\b\.?\s*", "", s)         # leading unit word
    s = re.sub(r"^(of|the|a|an)\s+", "", s)
    s = s.split(",")[0]                              # drop trailing prep clause
    s = re.sub(r"[^a-z\s-]", "", s)                  # keep letters only
    s = re.sub(r"\s+", " ", s).strip()
    return s


def _ingredient_list(raw: str) -> list[str]:
    if not isinstance(raw, str):
        return []
    out, seen = [], set()
    for part in raw.split(","):
        name = _simplify_ingredient(part)
        if not name or len(name) < 3 or len(name.split()) > 4:
            continue
        if name in _PREP_WORDS or name in seen:
            continue
        seen.add(name)
        out.append(name)
    return out


def _steps_from_directions(raw: str) -> list[str]:
    if not isinstance(raw, str):
        return []
    raw = _clean_text(raw.replace("\r", "\n"))
    # Prefer explicit newlines; otherwise split into sentences.
    parts = [p.strip() for p in raw.split("\n") if p.strip()]
    if len(parts) < 2:
        parts = [p.strip() for p in re.split(r"(?<=[.!?])\s+(?=[A-Z])", raw) if p.strip()]
    # merge very short fragments into the previous step
    steps: list[str] = []
    for p in parts:
        if steps and len(p) < 25:
            steps[-1] = steps[-1] + " " + p
        else:
            steps.append(p)
    return [s for s in steps if len(s) > 15]


def _minutes(row) -> int:
    for col in ("total_time", "cook_time", "prep_time"):
        v = row.get(col)
        if isinstance(v, str):
            h = re.search(r"(\d+)\s*hr", v)
            m = re.search(r"(\d+)\s*min", v)
            total = (int(h.group(1)) * 60 if h else 0) + (int(m.group(1)) if m else 0)
            if total:
                return total
    return 0


def _cuisine(row) -> str:
    cp = row.get("cuisine_path")
    if isinstance(cp, str):
        segs = [s for s in cp.split("/") if s]
        if segs:
            return segs[0].replace("-", " ").strip().title()
    return "International"


def _distribute(total: int, n: int) -> list[int]:
    if n <= 0:
        return []
    if total <= 0:
        total = n * 6
    base = max(2, total // n)
    durs = [base] * n
    durs[-1] = max(2, total - base * (n - 1))
    return durs


# ---------------------------------------------------------------------------
# 3. Normalize into clean recipe records
# ---------------------------------------------------------------------------

recipes: list[dict] = []
for _, r in tqdm(raw_df.iterrows(), total=len(raw_df), desc="Normalizing"):
    name = _clean_text(r.get("recipe_name", ""))
    ings = _ingredient_list(r.get("ingredients", ""))
    steps = _steps_from_directions(r.get("directions", ""))
    if not name or len(ings) < 3 or len(steps) < 2:
        continue
    steps = steps[:7]
    if len(steps) < 4 and len(steps) >= 2:
        pass  # keep short recipes too, 2-3 steps is fine
    minutes = _minutes(r) or len(steps) * 6
    try:
        servings = int(float(str(r.get("servings", "2")).split()[0]))
    except Exception:
        servings = 2
    servings = min(max(servings, 1), 12)
    recipes.append({
        "name": name,
        "ingredients": ings[:14],
        "steps": steps,
        "cuisine": _cuisine(r),
        "minutes": int(minutes),
        "servings": servings,
    })

print(f"\nClean recipes: {len(recipes)}")

config.DATA_DIR.mkdir(parents=True, exist_ok=True)
pd.DataFrame(recipes).to_parquet(config.RECIPES_PARQUET, index=False)
print(f"Saved -> {config.RECIPES_PARQUET}")


# ---------------------------------------------------------------------------
# 4. Build SFT pairs matching the app's exact prompt formats
# ---------------------------------------------------------------------------

PROPOSE_TMPL = (config.PROMPTS_DIR / "planner_propose.txt").read_text(encoding="utf-8")
RECIPE_TMPL = (config.PROMPTS_DIR / "planner_recipe.txt").read_text(encoding="utf-8")

_WHY = [
    "Uses your {a} and {b} for a quick, satisfying result.",
    "A fresh way to combine {a} with {b}.",
    "Turns {a} and {b} into a comforting classic.",
    "Light and flavorful, built around {a} and {b}.",
    "Makes the most of {a}, {b} and a few pantry staples.",
]


def _recipe_json(rec: dict) -> str:
    durs = _distribute(rec["minutes"], len(rec["steps"]))
    steps = [
        {"n": i + 1, "instruction": s, "duration": f"{d} min", "tip": None}
        for i, (s, d) in enumerate(zip(rec["steps"], durs))
    ]
    obj = {
        "name": rec["name"],
        "cuisine": rec["cuisine"],
        "servings": rec["servings"],
        "total_time_minutes": rec["minutes"],
        "final_dish_visual": f"A beautifully plated {rec['name'].lower()}, ready to serve.",
        "steps": steps,
    }
    return json.dumps(obj, ensure_ascii=False)


def _propose_json(rec: dict, others: list[dict]) -> str:
    a = rec["ingredients"][0] if rec["ingredients"] else "your ingredients"
    b = rec["ingredients"][1] if len(rec["ingredients"]) > 1 else "pantry staples"
    options = [{"name": rec["name"], "why": random.choice(_WHY).format(a=a, b=b)}]
    for o in others:
        oa = o["ingredients"][0] if o["ingredients"] else a
        ob = o["ingredients"][1] if len(o["ingredients"]) > 1 else b
        options.append({"name": o["name"], "why": random.choice(_WHY).format(a=oa, b=ob)})
    return json.dumps({"options": options}, ensure_ascii=False)


sft_path = config.DATA_DIR / "recipes_sft.jsonl"
n_recipe = n_propose = 0
with open(sft_path, "w", encoding="utf-8") as f:
    for idx, rec in enumerate(tqdm(recipes, desc="Building SFT")):
        ing_str = ", ".join(rec["ingredients"])

        # --- recipe task ---
        user_recipe = RECIPE_TMPL.replace("{dish_name}", rec["name"]).replace("{ingredients}", ing_str)
        f.write(json.dumps({"messages": [
            {"role": "user", "content": user_recipe},
            {"role": "assistant", "content": _recipe_json(rec)},
        ]}, ensure_ascii=False) + "\n")
        n_recipe += 1

        # --- propose task (use two other recipes as alternative options) ---
        others = [recipes[(idx + 7) % len(recipes)], recipes[(idx + 53) % len(recipes)]]
        user_propose = PROPOSE_TMPL.replace("{ingredients}", ing_str)
        f.write(json.dumps({"messages": [
            {"role": "user", "content": user_propose},
            {"role": "assistant", "content": _propose_json(rec, others)},
        ]}, ensure_ascii=False) + "\n")
        n_propose += 1

print(f"\nSFT pairs: {n_recipe} recipe + {n_propose} propose = {n_recipe + n_propose} -> {sft_path}")


# ---------------------------------------------------------------------------
# 5. Push to HF Hub
# ---------------------------------------------------------------------------

if HF_DATASET_REPO:
    from datasets import load_dataset

    ds = load_dataset("json", data_files=str(sft_path), split="train")
    ds.push_to_hub(HF_DATASET_REPO)
    print(f"Pushed {len(ds)} rows to {HF_DATASET_REPO}")

print("\nDone.")