app.py

# app.py — “Jerry” NLP Agent (per‑acre‑per‑year; uses salvage internally for Depr & OCC)
# ---------------------------------------------------------------------------------
# Students ALWAYS type: equipment_name, acres (per year), speed_acph (ac/hr), life_years, cost
# Optional (only when needed): pto_hp, diesel_price, wage, rate
#
# Jerry computes EXACTLY ONE cost per prompt (NLP → strict schema), with unit basis aligned:
#   • repair_per_acre_year       = (cost × repair_fraction@AccumHours) / (life_years × acres)
#   • fuel_per_acre              = (0.044 × pto_hp × diesel_price) / speed_acph
#   • lube_per_acre              = 0.15 × fuel_per_acre
#   • labor_per_acre             = wage / speed_acph
#   • depreciation_per_acre_year = (C − S) / (life_years × acres)   ← S from salvage table (hidden from students)
#   • occ_per_acre_year          = rate × ((C + S)/2) / (life_years × acres)   ← S from salvage table
#
# Shared formula:
#   AccumHours = life_years × (acres / speed_acph)

from __future__ import annotations
from typing import Optional, Dict, Any
import os, re, json
import gradio as gr

from repair_table_data import get_repair_fraction  # repair fraction by accumulated hours
from salvage_lookup import classify_machine_for_salvage, get_salvage_fraction  # salvage fraction (hidden)

from salvage_lookup import classify_machine_for_salvage, get_salvage_fraction  # keep
# add this helper inside app.py:
def salvage_parts(equip: str, cost: float, hp, acres: float, speed: float, life: float):
    """Return (salv_frac, salvage_value, class_used). Fallback frac = 0.20."""
    try:
        cls = classify_machine_for_salvage(equip, hp)
    except Exception:
        cls = None
    annual_hours = float(acres) / max(float(speed), 1e-9)
    try:
        frac = get_salvage_fraction(cls, annual_hours, float(life)) if cls else None
    except Exception:
        frac = None
    if frac is None:
        frac = 0.20
    return float(frac), float(cost) * float(frac), cls


import copy

# Which fields we track in the student HUD
HUD_FIELDS = [
    "equipment_name", "acres", "speed_acph", "life_years", "cost",
    "which_cost", "pto_hp", "diesel_price", "wage", "rate"
]

# Requirements to be able to compute each cost
REQUIREMENTS = {
    "repair_per_acre_year":   ["equipment_name", "acres", "speed_acph", "life_years", "cost"],
    "fuel_per_acre":          ["acres", "speed_acph", "life_years", "cost", "pto_hp", "diesel_price"],
    "lube_per_acre":          ["acres", "speed_acph", "life_years", "cost", "pto_hp", "diesel_price"],
    "labor_per_acre":         ["acres", "speed_acph", "life_years", "cost", "wage"],
    "depreciation_per_acre_year": ["equipment_name", "acres", "speed_acph", "life_years", "cost"],
    "occ_per_acre_year":      ["equipment_name", "acres", "speed_acph", "life_years", "cost", "rate"],
    "tax":                    ["cost"],
    "insurance":              ["cost"],
    "housing":                ["cost"],
}

REQUIREMENTS.update({
    "tax":       ["cost", "acres"],
    "insurance": ["cost", "acres"],
    "housing":   ["cost", "acres"],
})

def _pp(obj):
    try:
        return json.dumps(obj, indent=2, ensure_ascii=False, default=str)
    except Exception:
        return str(obj)

def merge_into_hud(hud: dict, new_data: dict) -> tuple[dict, list]:
    """Overwrite HUD with any non-null values from new_data. Return (updated_hud, changed_keys)."""
    hud = copy.deepcopy(hud)
    changed = []
    for k in HUD_FIELDS:
        if k in new_data and new_data[k] is not None:
            if hud.get(k) != new_data[k]:
                hud[k] = new_data[k]
                changed.append(k)
    return hud, changed

def missing_for(which_cost: str, hud: dict) -> list:
    req = REQUIREMENTS.get(which_cost or "", [])
    miss = [k for k in req if hud.get(k) in (None, "", float("nan"))]
    return miss

#def _pp(obj):
#    try:
#        return json.dumps(obj, indent=2, ensure_ascii=False, default=str)
#    except Exception:
#        return str(obj)

# ---- Optional LLM ----
LLM_OK = False
client = None
try:
    from openai import OpenAI
    if os.getenv("OPENAI_API_KEY"):
        client = OpenAI()
        LLM_OK = True
except Exception:
    LLM_OK = False

JERRY_SYSTEM_PROMPT = (
    "You are JERRY, a farm machinery cost coach **in training**. The student is the boss; "
    "you do exactly what they ask and nothing extra. Your job is to READ the student's sentence "
    "and OUTPUT ONLY a **minified JSON** object that our app will compute with.\n"

    "SCHEMA (always output these keys when present):\n"
    "equipment_name (string), acres (number), speed_acph (number), life_years (number), cost (number), "
    "which_cost (one of: 'repair_per_acre_year','fuel_per_acre','lube_per_acre','labor_per_acre',"
    "'depreciation_per_acre_year','occ_per_acre_year'), "
    "OPTIONAL: pto_hp (number), diesel_price (number), wage (number), rate (number), "
    "OPTIONAL diagnostics: missing (array of strings listing required-but-missing fields).\n"

    "TEACHING NOTES (formulas are for validation only—DO NOT COMPUTE VALUES):\n"
    "- Fuel_per_acre  = (0.044 * pto_hp * diesel_price) / speed_acph\n"
    "- Lube_per_acre  = 0.15 * Fuel_per_acre  (needs same inputs as fuel)\n"
    "- Labor_per_acre = wage / speed_acph\n"
    "- Repair_per_acre_year uses a table fraction at AccumHours, where AccumHours = life_years * (acres / speed_acph)\n"
    "- Depreciation_per_acre_year = (Cost - Salvage) / (life_years * acres)  (salvage is internal; you never output it)\n"
    "- OCC_per_acre_year = rate * ((Cost + Salvage)/2) / (life_years * acres) (salvage internal)\n"

    "CHECKS & BALANCES (training Jerry):\n"
    "1) The student ALWAYS supplies the five base items: equipment_name, acres, speed_acph, life_years, cost.\n"
    "   - If any base item is missing, still output JSON with that field = null and list it in 'missing'.\n"
    "2) For the requested cost, ensure the extra inputs exist:\n"
    "   - fuel_per_acre: require pto_hp and diesel_price; if absent set nulls and add to 'missing'.\n"
    "   - lube_per_acre: require pto_hp and diesel_price; if absent set nulls and add to 'missing'.\n"
    "   - labor_per_acre: require wage; if absent set null and add to 'missing'.\n"
    "   - depreciation_per_acre_year, occ_per_acre_year: DO NOT ask for salvage; it is internal.\n"
    "3) Parse numbers robustly—accept units and symbols and strip them: '$', commas, '%', 'ac/hr', 'acph', 'acres per hour', '/gal'. "
    "Examples: '$350,000' -> 350000; '3.80/gal' -> 3.80; '12 ac/hr' -> 12; '8%' -> 0.08.\n"
    "4) Never invent numbers. If unsure, use null and list the field in 'missing'.\n"
    "5) Map the student’s words to exactly one which_cost (priority by explicit keyword in the text):\n"
    "   'repair' -> 'repair_per_acre_year';\n"
    "   'fuel' or 'fuel cost' -> 'fuel_per_acre';\n"
    "   'lube' or 'lubrication' -> 'lube_per_acre';\n"
    "   'labor' or 'wage' -> 'labor_per_acre';\n"
    "   'depreciation' or 'depr' or 'straight line' -> 'depreciation_per_acre_year';\n"
    "   'occ' or 'opportunity cost' or 'interest' -> 'occ_per_acre_year'.\n"
    "   If none appears, leave which_cost null.\n"
    "6) Accept em dashes or punctuation as separators (e.g., '— fuel').\n"

    "OUTPUT RULES:\n"
    "- Output ONLY a single minified JSON object (no prose, no code block, no comments).\n"
    "- Keys must match the schema exactly. Use nulls for unknown numerics. Include 'missing' only when nonempty.\n"

    "EXAMPLES (do not echo back; just follow the pattern):\n"
    "Input: '4WD tractor, acres=1200, speed=12 ac/hr, life=8, cost=$350000, hp=300, diesel=3.80 — fuel'\n"
    "Output: {\"equipment_name\":\"4WD tractor\",\"acres\":1200,\"speed_acph\":12,\"life_years\":8,\"cost\":350000,"
    "\"pto_hp\":300,\"diesel_price\":3.8,\"which_cost\":\"fuel_per_acre\"}\n"

    "Input: 'Planter acres 900 speed 9 ac/hr life 12 cost 180000 lube'\n"
    "Output: {\"equipment_name\":\"Planter\",\"acres\":900,\"speed_acph\":9,\"life_years\":12,\"cost\":180000,"
    "\"which_cost\":\"lube_per_acre\"}\n"

    "Input: 'Tractor acres=1000 speed=10 ac/hr life=10 cost=$200,000 wage $22 — labor'\n"
    "Output: {\"equipment_name\":\"Tractor\",\"acres\":1000,\"speed_acph\":10,\"life_years\":10,\"cost\":200000,"
    "\"wage\":22,\"which_cost\":\"labor_per_acre\"}\n"
)



# ---- Computations ----

def accumulated_hours(life_years: float, acres: float, speed_acph: float) -> float:
    return float(life_years) * (float(acres)/max(speed_acph,1e-9))


def _salvage_value_hidden(equipment_name: str, cost: float, acres: float, speed_acph: float,
                          life_years: float, pto_hp: Optional[float]) -> float:
    """Hidden salvage value S = cost × salv_frac (from table). Students never see S directly."""
    machine_class = classify_machine_for_salvage(equipment_name, pto_hp)
    annual_hours = float(acres)/max(speed_acph,1e-9)
    salv_frac = get_salvage_fraction(machine_class, annual_hours, float(life_years)) if machine_class else None
    if salv_frac is None:
        salv_frac = 0.20  # safety fallback
    return float(cost) * float(salv_frac)


# --- Salvage helper (table-driven; fallback = 20%) ---
def salvage_value(equip: str, cost: float, hp, acres: float, speed: float, life: float) -> float:
    """
    Return salvage $ using your salvage table.
    - Classify machine with (equip, hp)
    - annual_hours = acres / speed
    - age_years = life
    - salvage = cost * salvage_fraction
    Fallback to 0.20 if the table can’t classify or returns None.
    """
    try:
        cls = classify_machine_for_salvage(equip, hp)
    except Exception:
        cls = None

    try:
        annual_hours = float(acres) / max(float(speed), 1e-9)
    except Exception:
        annual_hours = 0.0

    try:
        frac = get_salvage_fraction(cls, annual_hours, float(life)) if cls else None
    except Exception:
        frac = None

    if frac is None:
        frac = 0.20
    return float(cost) * float(frac)


def compute_repair_per_acre_year(equip: str, cost: float, life: float, acres: float, speed: float) -> Dict[str, Any]:
    hrs = accumulated_hours(life, acres, speed)
    frac, canon, lo, hi = get_repair_fraction(equip, hrs)
    total_repair = cost * frac
    val = total_repair / max(life*acres,1e-9)
    return {"canon":canon,"hours":hrs,"fraction":frac,"repair_per_acre_year":val}


def compute_fuel_per_acre(hp: float, diesel: float, speed: float) -> float:
    return (0.044*hp*diesel)/max(speed,1e-9)

def compute_lube_per_acre(fuel: float) -> float:
    return 0.15*fuel

def compute_labor_per_acre(wage: float, speed: float) -> float:
    return wage/max(speed,1e-9)

def compute_depr_per_acre_year(cost: float, salvage_value: float, life: float, acres: float) -> float:
    return (float(cost) - float(salvage_value)) / max(life * acres, 1e-9)

def compute_occ_per_acre_year(cost: float, salvage: float, rate: float, life: float, acres: float) -> float:
    avg_invest = 0.5 * (cost + salvage)
    return float(rate) * avg_invest / max(acres, 1e-9)

def compute_tax_per_acre(cost: float, acres: float) -> float:
    # Tax = 1% of cost, allocated per acre
    return 0.01 * float(cost) / max(float(acres), 1e-9)

def compute_insurance_per_acre(cost: float, acres: float) -> float:
    # Insurance = 0.5% of cost, allocated per acre
    return 0.005 * float(cost) / max(float(acres), 1e-9)

def compute_housing_per_acre(cost: float, acres: float) -> float:
    # Housing = 0.5% of cost, allocated per acre
    return 0.005 * float(cost) / max(float(acres), 1e-9)


# ---- Keyword intent + Regex fallback ----
NUM=r"[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?"

def _find(pat,s): m=re.search(pat,s,re.I); return float(m.group(1)) if m else None

def _name_guess(s: str) -> Optional[str]:
    m=re.search(r"^\s*([^,\n]+?)\s*(?:,|acres|$)",s,flags=re.I)
    return m.group(1).strip() if m else None


def which_from_keywords(s: str) -> Optional[str]:
    t=s.lower()
    if "repair" in t or "maint" in t: return "repair_per_acre_year"
    if "fuel" in t: return "fuel_per_acre"
    if "lube" in t or "lubric" in t: return "lube_per_acre"
    if "labor" in t or "wage" in t: return "labor_per_acre"
    if "depr" in t or "straight line" in t: return "depreciation_per_acre_year"
    if any(k in t for k in ["occ","interest","opportunity cost"]): return "occ_per_acre_year"
    return None


def regex_parse(user:str)->Dict[str,Any]:
    s=user
    d:Dict[str,Any]={}
    d["equipment_name"]=_name_guess(s)
    d["acres"]=_find(r"acres\s*(?:=|:)?\s*("+NUM+")",s) or _find(r"("+NUM+")\s*acres",s)
    d["speed_acph"]=_find(r"(speed|acph)\s*(?:=|:)?\s*("+NUM+")",s) or _find(r"("+NUM+")\s*(?:acph|ac/hr|ac\/hr|acres per hour)",s)
    d["life_years"]=_find(r"(life|years?)\s*(?:=|:)?\s*("+NUM+")",s)
    d["cost"]=_find(r"(cost|price|purchase)\s*(?:=|:|\$)?\s*("+NUM+")",s) or _find(r"\$("+NUM+")",s)
    d["pto_hp"]=_find(r"(pto\s*hp|hp|horsepower)\s*(?:=|:)?\s*("+NUM+")",s)
    d["diesel_price"]=_find(r"(diesel|fuel)\s*(price)?\s*(?:=|:|\$)?\s*("+NUM+")",s)
    d["wage"]=_find(r"(wage|labor\s*rate)\s*(?:=|:|\$)?\s*("+NUM+")",s)
    d["rate"]=_find(r"(rate|interest|occ)\s*(?:=|:)?\s*("+NUM+")\s*%?",s)
    w=which_from_keywords(s)
    if w: d["which_cost"]=w
    return {k:v for k,v in d.items() if v is not None}

# ---- LLM parse (same schema) ----

def llm_parse(user:str)->Dict[str,Any]:
    if not LLM_OK: return {}
    try:
        resp=client.chat.completions.create(
            model="gpt-4o-mini",temperature=0,
            messages=[{"role":"system","content":JERRY_SYSTEM_PROMPT},{"role":"user","content":user}])
        txt=(resp.choices[0].message.content or "").strip()
        if txt.startswith("{"): return json.loads(txt)
    except: pass
    return {}

# ---- Orchestrator ----
DEFAULT_RATE=0.08
def _to_float(x):
    if x is None: return None
    if isinstance(x, (int, float)): return float(x)
    s = str(x).strip().lower()
    # strip common decorations
    s = s.replace("$","").replace(",","")
    s = s.replace("ac/hr","").replace("acph","").replace("acres per hour","")
    s = s.replace("%","")
    # keep only number-ish chars
    s = re.sub(r"[^0-9eE\.\-\+]", "", s)
    try:
        return float(s)
    except Exception:
        return None

def which_from_keywords(s: str) -> Optional[str]:
    t = s.lower()
    if "fuel" in t: return "fuel_per_acre"
    if "lube" in t or "lubric" in t: return "lube_per_acre"
    if "labor" in t or "wage" in t: return "labor_per_acre"
    if "depr" in t or "depreciation" in t or "straight line" in t: return "depreciation_per_acre_year"
    if "occ" in t or "opportunity cost" in t or "interest" in t: return "occ_per_acre_year"
    if "repair" in t or "maint" in t: return "repair_per_acre_year"
    return None

def normalize_llm_data(data: Dict[str, Any], user_text: str) -> Dict[str, Any]:
    d = dict(data or {})

    # unify key names the LLM might choose
    if "hp" in d and "pto_hp" not in d:
        d["pto_hp"] = d["hp"]
    if "diesel" in d and "diesel_price" not in d:
        d["diesel_price"] = d["diesel"]
    if "life" in d and "life_years" not in d:
        d["life_years"] = d["life"]
    if "speed" in d and "speed_acph" not in d:
        d["speed_acph"] = d["speed"]

    # normalize which_cost
    wc = str(d.get("which_cost","")).strip().lower()
    wc_map = {
        "fuel":"fuel_per_acre",
        "fuel per acre":"fuel_per_acre",
        "fuel_per_acre":"fuel_per_acre",
        "repair":"repair_per_acre_year",
        "repair_per_acre_year":"repair_per_acre_year",
        "lube":"lube_per_acre",
        "lube_per_acre":"lube_per_acre",
        "labor":"labor_per_acre",
        "labor_per_acre":"labor_per_acre",
        "depreciation":"depreciation_per_acre_year",
        "depr":"depreciation_per_acre_year",
        "depreciation_sl":"depreciation_per_acre_year",
        "occ":"occ_per_acre_year",
        "opportunity cost":"occ_per_acre_year",
        "interest":"occ_per_acre_year",
        "occ_per_acre_year":"occ_per_acre_year",
    }
    d["which_cost"] = wc_map.get(wc) or which_from_keywords(user_text) or d.get("which_cost")

    # cast numerics (accept strings like "$350,000", "8%")
    for k in ["acres","speed_acph","life_years","cost","pto_hp","diesel_price","wage","rate"]:
        if k in d:
            d[k] = _to_float(d[k])

    # rate may be typed as 8 meaning 8% → 0.08
    if d.get("rate") is not None and d["rate"] > 1.0:
        d["rate"] = d["rate"] / 100.0

    return d

def jerry_agent(user: str):
    # ---------- Parse (LLM first, then regex) ----------
    raw_llm = llm_parse(user)            # dict or {}
    if raw_llm:
        mode = "LLM"
        raw_regex = {}
        raw = raw_llm
    else:
        mode = "regex"
        raw_regex = regex_parse(user)     # dict or {}
        raw = raw_regex

    # Fallback: infer which_cost from the text if parser left it empty
    def _which_from_text(txt: str):
        t = (txt or "").lower()
        if "fuel" in t: return "fuel_per_acre"
        if "lube" in t or "lubric" in t: return "lube_per_acre"
        if "labor" in t or "wage" in t: return "labor_per_acre"
        if "depr" in t or "depreciation" in t or "straight line" in t: return "depreciation_per_acre_year"
        if "occ" in t or "opportunity cost" in t or "interest" in t: return "occ_per_acre_year"
        if "tax" in t: return "tax"
        if "insurance" in t or "ins " in t: return "insurance"
        if "housing" in t or "house " in t: return "housing"
        if "repair" in t or "maint" in t: return "repair_per_acre_year"
        if "tax" in t or "taxes" in t: return "tax"
        if "insurance" in t or "ins " in t or "insur" in t: return "insurance"
        if "housing" in t or "house " in t or "storage" in t or "shed" in t: return "housing"
        return None

    data = dict(raw)
    if not data.get("which_cost"):
        guess = _which_from_text(user)
        if guess: data["which_cost"] = guess

    # ---------- Pull fields with defensive casting ----------
    def _f(x, default=None):
        if x is None: return default
        try: return float(x)
        except: return default

    equip = str(data.get("equipment_name", "tractor"))
    acres = _f(data.get("acres"), 0.0)
    speed = _f(data.get("speed_acph"), 1.0)
    life  = _f(data.get("life_years"), 1.0)
    cost  = _f(data.get("cost"), 0.0)
    which = str(data.get("which_cost") or "")

    hp     = _f(data.get("hp") if data.get("hp") is not None else data.get("pto_hp"))
    diesel = _f(data.get("diesel_price"))
    wage   = _f(data.get("wage"))
    rate   = _f(data.get("rate"), 0.08)
    if rate is not None and rate > 1.0:  # accept 8 or 8% as 0.08
        rate = rate / 100.0

    # ---------- Derived diagnostics (for the teacher JSON box) ----------
    annual_hours  = acres / max(speed, 1e-9)
    accum_hours   = life * annual_hours

    # Salvage diagnostics
    machine_class = classify_machine_for_salvage(equip, hp)
    salv_frac = None
    if machine_class:
        try:
            salv_frac = get_salvage_fraction(machine_class, annual_hours, life)
        except Exception:
            salv_frac = None
    salvage = cost * (salv_frac if salv_frac is not None else 0.20)

    # Repair fraction diagnostics (safe try)
    repair_frac = None
    repair_bracket = None
    canon_name = equip
    try:
        rf, canon, lo, hi = get_repair_fraction(equip, accum_hours)
        repair_frac = rf
        repair_bracket = {"low": {"hours": lo[0], "frac": lo[1]},
                          "high":{"hours": hi[0], "frac": hi[1]}}
        canon_name = canon
    except Exception:
        pass

    # ---------- Compute chosen cost ----------
    ans = None
    if which == "repair_per_acre_year":
        # Use the same computation you already rely on
        r = compute_repair_per_acre_year(equip, cost, life, acres, speed)
        ans = f"Jerry → REPAIR per acre-year: fraction={r['fraction']:.3f}, value=${r['repair_per_acre_year']:.2f}/ac/yr (mode={mode})"
    elif which == "fuel_per_acre":
        if hp is None or diesel is None:
            ans = f"Jerry → Need hp and diesel. (mode={mode})"
        else:
            f = compute_fuel_per_acre(hp, diesel, speed)
            ans = f"Jerry → FUEL per acre=${f:.2f} (mode={mode})"
    elif which == "lube_per_acre":
        if hp is None or diesel is None:
            ans = f"Jerry → Need hp and diesel. (mode={mode})"
        else:
            f = compute_fuel_per_acre(hp, diesel, speed)
            l = compute_lube_per_acre(f)
            ans = f"Jerry → LUBE per acre=${l:.2f} (mode={mode})"
    elif which == "labor_per_acre":
        if wage is None:
            ans = f"Jerry → Need wage. (mode={mode})"
        else:
            l = compute_labor_per_acre(wage, speed)
            ans = f"Jerry → LABOR per acre=${l:.2f} (mode={mode})"
    elif which == "depreciation_per_acre_year":
        salv_frac, salvage, salvage_class = salvage_parts(equip, cost, hp, acres, speed, life)
        d = compute_depr_per_acre_year(cost, salvage, life, acres)
        ans = (
            f"Jerry → DEPRECIATION per acre-year=${d:.2f} "
            f"(class={salvage_class}, salvage fraction={salv_frac:.2f}, salvage=${salvage:,.0f}; mode={mode})"
        )
    elif which == "occ_per_acre_year":
        salv_frac, salvage, salvage_class = salvage_parts(equip, cost, hp, acres, speed, life)
        o = compute_occ_per_acre_year(cost, salvage, rate, life, acres)
        ans = (
            f"Jerry → OCC per acre-year=${o:.2f} "
            f"(class={salvage_class}, salvage fraction={salvage_frac:.2f}, salvage=${salvage:,.0f}; mode={mode})"
        )
    elif which == "tax":
        tval = compute_tax_per_acre(cost, acres)
        ans = f"Jerry → TAX per acre=${tval:.2f} (mode={mode})"
    elif which == "insurance":
        ival = compute_insurance_per_acre(cost, acres)
        ans = f"Jerry → INSURANCE per acre=${ival:.2f} (mode={mode})"
    elif which == "housing":
        hval = compute_housing_per_acre(cost, acres)
        ans = f"Jerry → HOUSING per acre=${hval:.2f} (mode={mode})"
    else:
        ans = f"Jerry → Unknown cost. (mode={mode})"

    # ---------- Build teacher JSON (third box) ----------
    teacher_json = _pp({
        "mode": mode,
        "raw_llm": raw_llm,
        "raw_regex": raw_regex,
        "final_data": data,
        "derived": {
            "machine_canonical": canon_name,
            "machine_class_for_salvage": machine_class,
            "annual_hours": annual_hours,
            "accum_hours": accum_hours,
            "salvage_fraction": salv_frac,
            "salvage_value": salvage,
            "repair_fraction": repair_frac,
            "repair_bracket": repair_bracket,
        }
    })

    # Return BOTH: (student answer, teacher JSON)
    return ans, teacher_json


# ---- UI ----
with gr.Blocks(css="footer {visibility: hidden}") as demo:
    gr.Markdown("## Jerry — NLP Farm Machinery Cost Coach (per-acre per-year, with salvage)")

    # --- Session HUD state (persists across prompts within the browser session) ---
    hud_state = gr.State({k: None for k in HUD_FIELDS})

    with gr.Row():
        with gr.Column():
            q = gr.Textbox(label="Talk to Jerry", lines=4, placeholder="e.g., Tractor, acres=1000, speed=10 ac/hr, life=10, cost=$200000 — repair")
            ask = gr.Button("Ask Jerry")
        with gr.Column():
            out = gr.Textbox(label="Jerry says", lines=8)
        with gr.Column():
            gr.Markdown("### Student HUD — what Jerry sees")
            hud_box = gr.Textbox(label="Final Data (persists this session)", lines=18, value=_pp({k: None for k in HUD_FIELDS}))
            changed_box = gr.Textbox(label="Fields updated by last prompt", lines=3, value="(none)")

    # Optional: teacher debug accordion if you kept the teacher JSON earlier
    # (If you added a teacher debug already, you can leave this out or keep your version.)
    with gr.Accordion("Teacher debug — Parser & JSON (optional)", open=False):
        dbg = gr.Textbox(label="LLM/Regex + Raw/Final JSON + Diagnostics", lines=22)

    def _ask_with_hud(user_text, hud):
        # ---- Parse (LLM first, then regex) ----
        raw_llm = llm_parse(user_text)
        if raw_llm:
            mode = "LLM"; raw_regex = {}; raw = raw_llm
        else:
            mode = "regex"; raw_regex = regex_parse(user_text); raw = raw_regex

        # Fallback which_cost if missing
        def _which_from_text(txt: str):
            t = (txt or "").lower()
            if "fuel" in t: return "fuel_per_acre"
            if "lube" in t or "lubric" in t: return "lube_per_acre"
            if "labor" in t or "wage" in t: return "labor_per_acre"
            if "depr" in t or "depreciation" in t or "straight line" in t: return "depreciation_per_acre_year"
            if "occ" in t or "opportunity cost" in t or "interest" in t: return "occ_per_acre_year"
            if "tax" in t: return "tax"
            if "insurance" in t or "ins " in t: return "insurance"
            if "housing" in t or "house " in t: return "housing"
            if "repair" in t or "maint" in t: return "repair_per_acre_year"
            return None

        data = dict(raw)
        if not data.get("which_cost"):
            guessed = _which_from_text(user_text)
            if guessed: data["which_cost"] = guessed

        # ---- Merge into HUD memory ----
        new_hud, changed = merge_into_hud(hud, data)

        # Pull fields safely from HUD (not just this prompt)
        equip = str(new_hud.get("equipment_name") or "tractor")
        acres = float(new_hud.get("acres") or 0.0)
        speed = float(new_hud.get("speed_acph") or 1.0)
        life  = float(new_hud.get("life_years") or 1.0)
        cost  = float(new_hud.get("cost") or 0.0)
        which = str(new_hud.get("which_cost") or "")

        hp     = new_hud.get("hp") if new_hud.get("hp") is not None else new_hud.get("pto_hp")
        hp     = float(hp) if hp is not None else None
        diesel = float(new_hud.get("diesel_price")) if new_hud.get("diesel_price") is not None else None
        wage   = float(new_hud.get("wage")) if new_hud.get("wage") is not None else None
        rate   = new_hud.get("rate")
        if rate is not None:
            rate = float(rate)
            if rate > 1.0: rate = rate / 100.0  # accept 8 or 8% as 0.08
        else:
            rate = 0.08

        # ---- Readiness check for chosen cost ----
        if not which:
            msg = "Jerry → Tell me which cost (e.g., fuel, repair, depreciation). Check the HUD to see what I have so far."
            teacher_dbg = _pp({"mode": mode, "raw_llm": raw_llm, "raw_regex": raw_regex, "final_data": new_hud})
            return msg, _pp(new_hud), ", ".join(changed) or "(none)", teacher_dbg, new_hud

        missing = missing_for(which, new_hud)
        if missing:
            msg = f"Jerry → Not enough data for {which}. Missing: {', '.join(missing)}. Check the HUD and add what's missing."
            teacher_dbg = _pp({"mode": mode, "raw_llm": raw_llm, "raw_regex": raw_regex, "final_data": new_hud, "missing": missing})
            return msg, _pp(new_hud), ", ".join(changed) or "(none)", teacher_dbg, new_hud

        # ---- Compute (now that HUD is complete for this cost) ----
        salv_frac, salvage, salvage_class = salvage_parts(equip, cost, hp, acres, speed, life)
        if which == "repair_per_acre_year":
            r = compute_repair_per_acre_year(equip, cost, life, acres, speed)
            ans = f"Jerry → REPAIR per acre-year: fraction={r['fraction']:.3f}, value=${r['repair_per_acre_year']:.2f}/ac/yr (mode={mode})"
        elif which == "fuel_per_acre":
            f = compute_fuel_per_acre(hp, diesel, speed)
            ans = f"Jerry → FUEL per acre=${f:.2f} (mode={mode})"
        elif which == "lube_per_acre":
            f = compute_fuel_per_acre(hp, diesel, speed)
            l = compute_lube_per_acre(f)
            ans = f"Jerry → LUBE per acre=${l:.2f} (mode={mode})"
        elif which == "labor_per_acre":
            l = compute_labor_per_acre(wage, speed)
            ans = f"Jerry → LABOR per acre=${l:.2f} (mode={mode})"
        
        elif which == "depreciation_per_acre_year":
            d = compute_depr_per_acre_year(cost, salvage, life, acres)
            ans = (
                f"Jerry → DEPRECIATION per acre-year=${d:.2f} "
                f"(salvage fraction={salv_frac:.2f}, salvage=${salvage:,.0f}; mode={mode})"
            )
        elif which == "occ_per_acre_year":
            o = compute_occ_per_acre_year(cost, salvage, rate, life, acres)
            ans = (
                f"Jerry → OCC per acre-year=${o:.2f} "
                f"(salvage fraction={salv_frac:.2f}, salvage=${salvage:,.0f}; mode={mode})"
            )
        elif which == "tax":
            tval = compute_tax_per_acre(cost, acres)
            ans = f"Jerry → TAX per acre=${tval:.2f} (mode={mode})"
        elif which == "insurance":
           ival = compute_insurance_per_acre(cost, acres)
           ans = f"Jerry → INSURANCE per acre=${ival:.2f} (mode={mode})"
        elif which == "housing":
           hval = compute_housing_per_acre(cost, acres)
           ans = f"Jerry → HOUSING per acre=${hval:.2f} (mode={mode})"
        teacher_dbg = _pp({
            "mode": mode,
            "raw_llm": raw_llm,
            "raw_regex": raw_regex,
            "final_data": new_hud
        })
        return ans, _pp(new_hud), ", ".join(changed) or "(none)", teacher_dbg, new_hud

    # Wire it up; we update 5 outputs: student answer, HUD JSON, changed fields, teacher debug, and the state
    ask.click(_ask_with_hud, [q, hud_state], [out, hud_box, changed_box, dbg, hud_state])

if __name__=="__main__":
    demo.queue().launch(server_name="0.0.0.0",server_port=int(os.getenv("PORT","7860")))