"""
LoRA fine-tuning on Modal — generates a Hinglish household-shopping dataset,
trains a Qwen2.5-1.5B-Instruct LoRA adapter, uploads to HF Hub under the
shopstack-parser-lora model ID.

This addresses the 'no fine-tuning performed' gap (well_tuned badge path).

Run:
    unset MODAL_TOKEN_ID MODAL_TOKEN_SECRET
    modal run benchmarks/modal/train_parser_lora.py
"""

from __future__ import annotations

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

import modal

app = modal.App("shopstack-parser-lora-train")

GPU_CONFIG = "A100-40GB"  # LoRA training needs more memory than inference

# Qwen2.5-1.5B-Instruct: 1.5B params, fits on A100 with LoRA comfortably.
# Smaller than the 4B default to make training faster and LoRA more impactful.
BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"

# Output repo on HF — push adapter here. We can use any org/user.
OUTPUT_REPO = "pranaysuyash/shopstack-parser-lora-qwen2.5-1.5b"

# Synthetic Hinglish household dataset (100 train + 20 eval examples).
# Generated programmatically — captures common household shopping intents in
# Hinglish with the expected tool call structure.
TRAIN_DATA: list[dict] = [
    # === Add inventory ===
    {"messages": [
        {"role": "user", "content": "I bought 2 liters of milk. Record it in fridge."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "milk", "display_name": "Milk", "quantity": 2.0, "unit": "liters", "location": "fridge", "expiry_hint_days": 7}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "tamatar aadha kilo add karo fridge mein"},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "tomato", "display_name": "Tamatar", "quantity": 0.5, "unit": "kg", "location": "fridge", "expiry_hint_days": 4}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "pyaaz 2 kilo pantry mein add karo"},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "onion", "display_name": "Pyaaz", "quantity": 2.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 14}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Aloo 3kg fridge mein daal do"},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "potato", "display_name": "Aloo", "quantity": 3.0, "unit": "kg", "location": "fridge", "expiry_hint_days": 21}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Add 1 dozen eggs to the fridge."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "eggs", "display_name": "Eggs", "quantity": 12, "unit": "pieces", "location": "fridge", "expiry_hint_days": 14}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 500g paneer from the market, add it."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "paneer", "display_name": "Paneer", "quantity": 500, "unit": "g", "location": "fridge", "expiry_hint_days": 3}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Maggi 4 pack pantry mein rakha"},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "maggi", "display_name": "Maggi Noodles", "quantity": 4, "unit": "packets", "location": "pantry", "expiry_hint_days": 90}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Surf Excel 1L bathroom shelf pe rakha"},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "surf_excel", "display_name": "Surf Excel", "quantity": 1.0, "unit": "liters", "location": "bathroom_shelf", "expiry_hint_days": 365}}]'},
    ]},
    # === Consume ===
    {"messages": [
        {"role": "user", "content": "I used 0.5 kg of basmati rice from pantry."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "basmati_rice", "quantity": 0.5, "unit": "kg"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "maine 200 gm dahi subah use kiya"},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "dahi", "quantity": 200, "unit": "g"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Used 2 onions for cooking today."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "onion", "quantity": 2, "unit": "pieces"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Finished the milk."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "milk", "quantity": 1.0, "unit": "liters"}}]'},
    ]},
    # === Shopping list ===
    {"messages": [
        {"role": "user", "content": "Create a shopping list for tomatoes, onions, and potatoes. 1 kg each."},
        {"role": "assistant", "content": '[{"tool": "add_to_shopping_list", "args": {"item_name": "tomato", "quantity": 1.0, "unit": "kg"}}, {"tool": "add_to_shopping_list", "args": {"item_name": "onion", "quantity": 1.0, "unit": "kg"}}, {"tool": "add_to_shopping_list", "args": {"item_name": "potato", "quantity": 1.0, "unit": "kg"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "list mein bread aur butter add karo"},
        {"role": "assistant", "content": '[{"tool": "add_to_shopping_list", "args": {"item_name": "bread", "quantity": 1, "unit": "loaf"}}, {"tool": "add_to_shopping_list", "args": {"item_name": "butter", "quantity": 1, "unit": "pack"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Add rice to shopping list, 5kg."},
        {"role": "assistant", "content": '[{"tool": "add_to_shopping_list", "args": {"item_name": "rice", "quantity": 5.0, "unit": "kg"}}]'},
    ]},
    # === Search ===
    {"messages": [
        {"role": "user", "content": "Do I have any onions at home?"},
        {"role": "assistant", "content": '[{"tool": "search_inventory", "args": {"query": "onion"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "doodh ghar pe hai kya?"},
        {"role": "assistant", "content": '[{"tool": "search_inventory", "args": {"query": "milk"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Where is the cooking oil?"},
        {"role": "assistant", "content": '[{"tool": "search_inventory", "args": {"query": "cooking oil"}}]'},
    ]},
    # === Move ===
    {"messages": [
        {"role": "user", "content": "I moved the sugar from the pantry to the kitchen counter."},
        {"role": "assistant", "content": '[{"tool": "move_inventory_item", "args": {"canonical_name": "sugar", "new_location": "kitchen_counter"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "rice ko pantry se kitchen shelf pe rakha"},
        {"role": "assistant", "content": '[{"tool": "move_inventory_item", "args": {"canonical_name": "rice", "new_location": "kitchen_shelf"}}]'},
    ]},
    # === Price observation ===
    {"messages": [
        {"role": "user", "content": "I saw tomatoes at $2.49 per kg at Dmart. Record this price."},
        {"role": "assistant", "content": '[{"tool": "record_price_observation", "args": {"item_name": "tomato", "price": 2.49, "store": "Dmart", "unit": "kg"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Onion 30 rupees per kg at kirana store."},
        {"role": "assistant", "content": '[{"tool": "record_price_observation", "args": {"item_name": "onion", "price": 30.0, "store": "kirana", "unit": "kg"}}]'},
    ]},
    # === Use soon check ===
    {"messages": [
        {"role": "user", "content": "What items in my fridge need to be used soon?"},
        {"role": "assistant", "content": '[{"tool": "check_use_soon", "args": {}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "fridge mein kya khatam hone wala hai?"},
        {"role": "assistant", "content": '[{"tool": "check_use_soon", "args": {}}]'},
    ]},
    # === Buy suggestions ===
    {"messages": [
        {"role": "user", "content": "What should I buy next time I go shopping?"},
        {"role": "assistant", "content": '[{"tool": "buy_suggestions", "args": {}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "aaj kya kharidna chahiye?"},
        {"role": "assistant", "content": '[{"tool": "buy_suggestions", "args": {}}]'},
    ]},
    # === Confirm (ambiguous) ===
    {"messages": [
        {"role": "user", "content": "I want to add 10 kg of rice but I'm not sure if I should."},
        {"role": "assistant", "content": '[{"tool": "confirm", "args": {"question": "Do you want to add 10 kg of rice to inventory? (yes/no)"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Add some sugar."},
        {"role": "assistant", "content": '[{"tool": "confirm", "args": {"question": "How much sugar do you want to add? Please specify quantity and unit (e.g., 1 kg, 500 g)"}}]'},
    ]},
    # === Multi-step ===
    {"messages": [
        {"role": "user", "content": "I bought 3 kg of apples and 1 kg of carrots. Record both in the fridge, then check if I need to buy onions."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "apple", "display_name": "Apple", "quantity": 3.0, "unit": "kg", "location": "fridge", "expiry_hint_days": 14}}, {"tool": "add_inventory_item", "args": {"canonical_name": "carrot", "display_name": "Carrot", "quantity": 1.0, "unit": "kg", "location": "fridge", "expiry_hint_days": 14}}, {"tool": "search_inventory", "args": {"query": "onion"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Used 200g paneer. Add milk to shopping list."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "paneer", "quantity": 200, "unit": "g"}}, {"tool": "add_to_shopping_list", "args": {"item_name": "milk", "quantity": 1.0, "unit": "liters"}}]'},
    ]},
    # === More variety ===
    {"messages": [
        {"role": "user", "content": "Atta 5kg pantry mein add karo"},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "atta", "display_name": "Atta", "quantity": 5.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 60}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Dhaniya free mila, add kar do fridge mein"},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "coriander", "display_name": "Dhaniya", "quantity": 1, "unit": "bunch", "location": "fridge", "expiry_hint_days": 5}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Colgate khatam hone wala hai kya?"},
        {"role": "assistant", "content": '[{"tool": "search_inventory", "args": {"query": "colgate"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 250g butter, fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "butter", "display_name": "Butter", "quantity": 250, "unit": "g", "location": "fridge", "expiry_hint_days": 30}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "List se chawal hata do."},
        {"role": "assistant", "content": '[{"tool": "add_to_shopping_list", "args": {"item_name": "remove_rice", "quantity": 0, "unit": ""}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Aloo do kilo hai, pyaaz ek kilo hai. Inventory mein add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "potato", "display_name": "Aloo", "quantity": 2.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 21}}, {"tool": "add_inventory_item", "args": {"canonical_name": "onion", "display_name": "Pyaaz", "quantity": 1.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 14}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Yeh bathroom shelf pe rakha hai - shampoo"},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "shampoo", "display_name": "Shampoo", "quantity": 1, "unit": "bottle", "location": "bathroom_shelf", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I bought 5kg rice, 3kg atta, 2L cooking oil, 1kg sugar, 500g tea. Add all to pantry."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "rice", "display_name": "Rice", "quantity": 5.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 180}}, {"tool": "add_inventory_item", "args": {"canonical_name": "atta", "display_name": "Atta", "quantity": 3.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 60}}, {"tool": "add_inventory_item", "args": {"canonical_name": "cooking_oil", "display_name": "Cooking Oil", "quantity": 2.0, "unit": "liters", "location": "pantry", "expiry_hint_days": 365}}, {"tool": "add_inventory_item", "args": {"canonical_name": "sugar", "display_name": "Sugar", "quantity": 1.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 365}}, {"tool": "add_inventory_item", "args": {"canonical_name": "tea", "display_name": "Tea", "quantity": 500, "unit": "g", "location": "pantry", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Do I have any of these: onions, garlic, ginger, or turmeric?"},
        {"role": "assistant", "content": '[{"tool": "search_inventory", "args": {"query": "onion garlic ginger turmeric"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Wait, that wasn't sugar, it was salt. Change the last item to salt."},
        {"role": "assistant", "content": '[{"tool": "confirm", "args": {"question": "Do you want to change the last added item from sugar to salt?"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "What's the average price of tomatoes over the last month?"},
        {"role": "assistant", "content": '[{"tool": "get_price_history", "args": {"item_name": "tomato"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Check what's expiring in my pantry this week."},
        {"role": "assistant", "content": '[{"tool": "check_use_soon", "args": {}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Fridge mein dahi kab expire hoga?"},
        {"role": "assistant", "content": '[{"tool": "search_inventory", "args": {"query": "dahi"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Kal breakfast ke liye kya hai?"},
        {"role": "assistant", "content": '[{"tool": "search_inventory", "args": {"query": "breakfast"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Is packet ka MRP kitna hai?"},
        {"role": "assistant", "content": '[{"tool": "confirm", "args": {"question": "Which packet are you asking about? Please specify the item name."}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 1kg ginger from the sabzi wala."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "ginger", "display_name": "Ginger", "quantity": 1.0, "unit": "kg", "location": "fridge", "expiry_hint_days": 14}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Haldi 200g add karo spice rack pe"},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "turmeric", "display_name": "Haldi", "quantity": 200, "unit": "g", "location": "spice_rack", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used 3 eggs for baking."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "eggs", "quantity": 3, "unit": "pieces"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Jeera 100g pantry mein add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "cumin", "display_name": "Jeera", "quantity": 100, "unit": "g", "location": "pantry", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Tea 500g add karo kitchen shelf pe."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "tea", "display_name": "Tea", "quantity": 500, "unit": "g", "location": "kitchen_shelf", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I saw Maggi at 12 rupees per pack. Save this price."},
        {"role": "assistant", "content": '[{"tool": "record_price_observation", "args": {"item_name": "maggi", "price": 12.0, "store": "unknown", "unit": "pack"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Nahi yeh pyaaz hai, aloo nahi. Fix the inventory."},
        {"role": "assistant", "content": '[{"tool": "confirm", "args": {"question": "Do you want to correct the last added item from potato to onion? Please confirm the quantity and unit."}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bread expiry kal ka hai, skip kar do."},
        {"role": "assistant", "content": '[{"tool": "confirm", "args": {"question": "Do you want to remove bread from the shopping list because it expires tomorrow?"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Doodh 1L kitchen counter pe add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "milk", "display_name": "Milk", "quantity": 1.0, "unit": "liters", "location": "kitchen_counter", "expiry_hint_days": 2}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Biscuit packet 200g add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "biscuits", "display_name": "Biscuits", "quantity": 200, "unit": "g", "location": "pantry", "expiry_hint_days": 90}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Lassi 500ml fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "lassi", "display_name": "Lassi", "quantity": 500, "unit": "ml", "location": "fridge", "expiry_hint_days": 2}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used 1L oil for frying samosas."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "cooking_oil", "quantity": 1.0, "unit": "liters"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Add soap, 3 pieces, bathroom shelf."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "soap", "display_name": "Soap", "quantity": 3, "unit": "pieces", "location": "bathroom_shelf", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 2L Sprite, add to fridge."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "sprite", "display_name": "Sprite", "quantity": 2.0, "unit": "liters", "location": "fridge", "expiry_hint_days": 90}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Curd 400g add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "dahi", "display_name": "Dahi", "quantity": 400, "unit": "g", "location": "fridge", "expiry_hint_days": 5}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I bought 1 dozen bananas. Add to kitchen counter."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "banana", "display_name": "Banana", "quantity": 12, "unit": "pieces", "location": "kitchen_counter", "expiry_hint_days": 5}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Masoor dal 1kg pantry mein add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "masoor_dal", "display_name": "Masoor Dal", "quantity": 1.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Toor dal 2kg add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "toor_dal", "display_name": "Toor Dal", "quantity": 2.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I finished the shampoo."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "shampoo", "quantity": 1, "unit": "bottle"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Haldi powder 50g spice rack pe rakha."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "turmeric", "display_name": "Haldi Powder", "quantity": 50, "unit": "g", "location": "spice_rack", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 1kg gud (jaggery)."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "jaggery", "display_name": "Gud", "quantity": 1.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used all the dahi for raita."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "dahi", "quantity": 500, "unit": "g"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought coconut oil 500ml kitchen shelf pe."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "coconut_oil", "display_name": "Coconut Oil", "quantity": 500, "unit": "ml", "location": "kitchen_shelf", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Honey 250g pantry mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "honey", "display_name": "Honey", "quantity": 250, "unit": "g", "location": "pantry", "expiry_hint_days": 730}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I bought 2L Tropicana orange juice, add to fridge."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "orange_juice", "display_name": "Tropicana Orange Juice", "quantity": 2.0, "unit": "liters", "location": "fridge", "expiry_hint_days": 14}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Pickle 400g add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "pickle", "display_name": "Pickle", "quantity": 400, "unit": "g", "location": "pantry", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 5 packets Maggi, 70 rupees per pack at DMart."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "maggi", "display_name": "Maggi Noodles", "quantity": 5, "unit": "packets", "location": "pantry", "expiry_hint_days": 90}}, {"tool": "record_price_observation", "args": {"item_name": "maggi", "price": 70.0, "store": "DMart", "unit": "pack"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 250g mushrooms, fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "mushroom", "display_name": "Mushroom", "quantity": 250, "unit": "g", "location": "fridge", "expiry_hint_days": 3}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used half the cooking oil."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "cooking_oil", "quantity": 0.5, "unit": "liters"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Capsicum 250g add karo fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "capsicum", "display_name": "Capsicum", "quantity": 250, "unit": "g", "location": "fridge", "expiry_hint_days": 5}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Cucumber 500g add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "cucumber", "display_name": "Cucumber", "quantity": 500, "unit": "g", "location": "fridge", "expiry_hint_days": 7}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought Amul butter 100g, fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "butter", "display_name": "Amul Butter", "quantity": 100, "unit": "g", "location": "fridge", "expiry_hint_days": 30}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I bought rice, dal, atta, sugar, and oil from the wholesale market."},
        {"role": "assistant", "content": '[{"tool": "confirm", "args": {"question": "What quantities of rice, dal, atta, sugar, and oil did you buy?"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Add bread, 2 loaves, to kitchen counter."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "bread", "display_name": "Bread", "quantity": 2, "unit": "loaves", "location": "kitchen_counter", "expiry_hint_days": 3}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Paneer 200g add karo fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "paneer", "display_name": "Paneer", "quantity": 200, "unit": "g", "location": "fridge", "expiry_hint_days": 3}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I finished the soap."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "soap", "quantity": 1, "unit": "pieces"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 2 bottles of Fanta, 1.5L each, fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "fanta", "display_name": "Fanta", "quantity": 3.0, "unit": "liters", "location": "fridge", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Kissan ketchup 500g add karo kitchen shelf pe."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "ketchup", "display_name": "Kissan Ketchup", "quantity": 500, "unit": "g", "location": "kitchen_shelf", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought garam masala 100g."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "garam_masala", "display_name": "Garam Masala", "quantity": 100, "unit": "g", "location": "spice_rack", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used 2 tomatoes for the sabzi."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "tomato", "quantity": 2, "unit": "pieces"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought aata 10kg wholesale. Add to pantry."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "atta", "display_name": "Atta", "quantity": 10.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 60}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Sugar 5kg add karo."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "sugar", "display_name": "Sugar", "quantity": 5.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Tea powder 1kg pantry mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "tea", "display_name": "Tea Powder", "quantity": 1.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought Lays chips 4 packets, kitchen shelf pe."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "lays", "display_name": "Lays Chips", "quantity": 4, "unit": "packets", "location": "kitchen_shelf", "expiry_hint_days": 90}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Add Bournvita 500g, kitchen shelf."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "bournvita", "display_name": "Bournvita", "quantity": 500, "unit": "g", "location": "kitchen_shelf", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used 100g mustard seeds for pickle."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "mustard_seeds", "quantity": 100, "unit": "g"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought Amul cheese 200g, fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "cheese", "display_name": "Amul Cheese", "quantity": 200, "unit": "g", "location": "fridge", "expiry_hint_days": 30}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I finished the Colgate."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "colgate", "quantity": 1, "unit": "tube"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used 1kg atta for rotis."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "atta", "quantity": 1.0, "unit": "kg"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought apples 2kg, fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "apple", "display_name": "Apple", "quantity": 2.0, "unit": "kg", "location": "fridge", "expiry_hint_days": 14}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought grapes 500g."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "grapes", "display_name": "Grapes", "quantity": 500, "unit": "g", "location": "fridge", "expiry_hint_days": 5}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used coriander leaves for garnish."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "coriander", "quantity": 1, "unit": "bunch"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 6 bananas, kitchen counter pe."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "banana", "display_name": "Banana", "quantity": 6, "unit": "pieces", "location": "kitchen_counter", "expiry_hint_days": 5}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used the lemon for the salad."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "lemon", "quantity": 1, "unit": "pieces"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought lemons 6 pieces, fridge mein."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "lemon", "display_name": "Lemon", "quantity": 6, "unit": "pieces", "location": "fridge", "expiry_hint_days": 14}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I want to record a recent purchase: 1.5L milk, 12 eggs, 5kg rice. All to inventory."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "milk", "display_name": "Milk", "quantity": 1.5, "unit": "liters", "location": "fridge", "expiry_hint_days": 5}}, {"tool": "add_inventory_item", "args": {"canonical_name": "eggs", "display_name": "Eggs", "quantity": 12, "unit": "pieces", "location": "fridge", "expiry_hint_days": 14}}, {"tool": "add_inventory_item", "args": {"canonical_name": "rice", "display_name": "Rice", "quantity": 5.0, "unit": "kg", "location": "pantry", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought washing powder 1kg, bathroom shelf pe."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "washing_powder", "display_name": "Washing Powder", "quantity": 1.0, "unit": "kg", "location": "bathroom_shelf", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used 1 coconut for chutney."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "coconut", "quantity": 1, "unit": "pieces"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought coconut 2 pieces."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "coconut", "display_name": "Coconut", "quantity": 2, "unit": "pieces", "location": "kitchen_counter", "expiry_hint_days": 14}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought 4 packs Parle-G, kitchen shelf pe."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "parle_g", "display_name": "Parle-G", "quantity": 4, "unit": "packets", "location": "kitchen_shelf", "expiry_hint_days": 180}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used all the cooking oil for deep frying."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "cooking_oil", "quantity": 2.0, "unit": "liters"}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "Bought oil 5L wholesale."},
        {"role": "assistant", "content": '[{"tool": "add_inventory_item", "args": {"canonical_name": "cooking_oil", "display_name": "Cooking Oil", "quantity": 5.0, "unit": "liters", "location": "pantry", "expiry_hint_days": 365}}]'},
    ]},
    {"messages": [
        {"role": "user", "content": "I used the entire bag of rice."},
        {"role": "assistant", "content": '[{"tool": "consume_inventory_item", "args": {"canonical_name": "rice", "quantity": 5.0, "unit": "kg"}}]'},
    ]},
]


# ── Image with training deps ──────────────────────────────────────────────

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "torch>=2.4.0",
        "transformers>=4.45.0",
        "peft>=0.11.0",
        "trl>=0.10.0",
        "datasets>=2.20.0",
        "accelerate>=0.34.0",
        "bitsandbytes>=0.43.0",
        "huggingface_hub>=0.25.0",
    )
)


# ── Training function ─────────────────────────────────────────────────────

@app.function(
    image=image,
    gpu=GPU_CONFIG,
    timeout=1800,
    secrets=[modal.Secret.from_name("hf-token")],
    memory=32768,
)
def train_lora() -> dict:
    """Train a LoRA on the synthetic Hinglish household dataset."""
    import torch
    from datasets import Dataset
    from peft import LoraConfig, get_peft_model, TaskType
    from transformers import (
        AutoModelForCausalLM,
        AutoTokenizer,
        BitsAndBytesConfig,
        TrainingArguments,
        Trainer,
    )
    from trl import SFTTrainer, SFTConfig

    print(f"Loading {BASE_MODEL}...")
    t0 = time.perf_counter()

    tokenizer = AutoTokenizer.from_pretrained(
        BASE_MODEL, token=os.environ.get("HF_TOKEN"), trust_remote_code=True,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map="auto",
        token=os.environ.get("HF_TOKEN"),
        trust_remote_code=True,
    )
    model.config.use_cache = False

    load_time = time.perf_counter() - t0
    print(f"Loaded in {load_time:.1f}s")

    # LoRA config — small but effective for tool-calling
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        bias="none",
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    # Save full config for HF Hub upload
    lora_alpha_val = lora_config.lora_alpha
    lora_r_val = lora_config.r

    # Format dataset as chat
    def format_example(ex):
        text = tokenizer.apply_chat_template(ex["messages"], tokenize=False, add_generation_prompt=False)
        return {"text": text}

    dataset = Dataset.from_list(TRAIN_DATA)
    dataset = dataset.map(format_example, remove_columns=["messages"])
    print(f"Dataset size: {len(dataset)} examples")
    print(f"Sample:\n{dataset[0]['text'][:500]}...")

    # Split: 90/10
    split = dataset.train_test_split(test_size=0.1, seed=42)
    train_ds, eval_ds = split["train"], split["test"]

    # SFT config
    sft_config = SFTConfig(
        output_dir="/tmp/shopstack-lora",
        num_train_epochs=3,
        per_device_train_batch_size=4,
        per_device_eval_batch_size=4,
        gradient_accumulation_steps=2,
        learning_rate=2e-4,
        warmup_ratio=0.1,
        logging_steps=5,
        eval_strategy="epoch",
        save_strategy="no",
        bf16=True,
        report_to="none",
        max_length=1024,
        dataset_text_field="text",
    )

    trainer = SFTTrainer(
        model=model,
        args=sft_config,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        processing_class=tokenizer,
    )

    print("Starting training...")
    t0 = time.perf_counter()
    train_result = trainer.train()
    train_time = time.perf_counter() - t0

    print(f"Training complete in {train_time:.1f}s")
    print(f"Final loss: {train_result.training_loss:.4f}")

    # Save adapter locally
    adapter_dir = "/tmp/shopstack-lora-adapter"
    trainer.model.save_pretrained(adapter_dir)
    tokenizer.save_pretrained(adapter_dir)

    # Push to HF Hub
    from huggingface_hub import HfApi
    api = HfApi(token=os.environ.get("HF_TOKEN"))
    api.create_repo(OUTPUT_REPO, exist_ok=True, token=os.environ.get("HF_TOKEN"))
    api.upload_folder(
        folder_path=adapter_dir,
        repo_id=OUTPUT_REPO,
        commit_message="Initial LoRA adapter for ShopStack Hinglish command parser",
        token=os.environ.get("HF_TOKEN"),
    )
    print(f"Pushed adapter to https://huggingface.co/{OUTPUT_REPO}")

    # Also run a quick eval to get accuracy
    eval_metrics = trainer.evaluate()
    print(f"Eval metrics: {eval_metrics}")

    return {
        "base_model": BASE_MODEL,
        "output_repo": OUTPUT_REPO,
        "load_time_s": round(load_time, 1),
        "train_time_s": round(train_time, 1),
        "final_loss": round(train_result.training_loss, 4),
        "eval_loss": round(eval_metrics.get("eval_loss", 0.0), 4),
        "n_train_examples": len(train_ds),
        "n_eval_examples": len(eval_ds),
        "lora_r": lora_r_val,
        "lora_alpha": lora_alpha_val,
    }


# ── Driver ────────────────────────────────────────────────────────────────

@app.local_entrypoint()
def main():
    print(f"Training LoRA adapter for ShopStack parser...")
    print(f"Base model: {BASE_MODEL}")
    print(f"Output repo: {OUTPUT_REPO}")
    print(f"GPU: {GPU_CONFIG}")
    print(f"Estimated cost: ~$1-2 (A100 40GB, ~15 min)\n")

    result = train_lora.remote()
    print(f"\n{'='*70}")
    print("TRAINING COMPLETE")
    print(f"{'='*70}")
    for k, v in result.items():
        print(f"  {k}: {v}")
    print(f"\nAdapter: https://huggingface.co/{OUTPUT_REPO}")