app.py

# app.py
import os, io, json
from datetime import datetime, timezone
from dateutil import tz
import pandas as pd
import gradio as gr

from huggingface_hub import (
    HfApi,
    HfFolder,
    hf_hub_download,  # anonymous reads for public datasets
)

# --------------------
# Config
# --------------------
DATA_REPO = os.environ.get("DATA_REPO", "jeffrey1963/Farm_Sim_Data")
DATA_REPO_TYPE = "dataset"  # don't change

# --------------------
# HF helpers
# --------------------
def _api_readonly() -> HfApi:
    """OK without token for public dataset reads."""
    tok = HfFolder.get_token()
    return HfApi(token=tok) if tok else HfApi()

def _api_rw() -> HfApi:
    """Require token for uploads."""
    tok = HfFolder.get_token()
    if not tok:
        raise RuntimeError("Missing HF_TOKEN (add a Space secret named HF_TOKEN).")
    return HfApi(token=tok)

def _candidate_purchase_paths(student_id: str) -> list[str]:
    """
    Return all jsonl paths for this student.
    Supports both:
      - purchases/{id}.jsonl
      - purchases/{id}-TIMESTAMP.jsonl  (your current store behavior)
    """
    api = _api_readonly()
    try:
        files = api.list_repo_files(repo_id=DATA_REPO, repo_type=DATA_REPO_TYPE)
    except Exception:
        return []
    prefix = f"purchases/{student_id}"
    out = []
    for f in files:
        if not f.startswith("purchases/") or not f.endswith(".jsonl"):
            continue
        if f == f"{prefix}.jsonl" or f.startswith(prefix + "-"):
            out.append(f)
    return sorted(out)  # oldest first

def _download_to_local(path_in_repo: str) -> str | None:
    try:
        return hf_hub_download(
            repo_id=DATA_REPO,
            repo_type=DATA_REPO_TYPE,
            filename=path_in_repo,
        )
    except Exception:
        return None

# --------------------
# Core
# --------------------
def fetch_student_assets(student_id: str):
    """Return list of assets purchased by this student (aggregate all files)."""
    sid = (student_id or "").strip()
    if not sid:
        return []

    paths = _candidate_purchase_paths(sid)
    if not paths:
        return []

    rows = []
    for p in paths:
        local = _download_to_local(p)
        if not local:
            continue
        with open(local, "r", encoding="utf-8") as f:
            for line in f:
                s = line.strip()
                if not s:
                    continue
                try:
                    rows.append(json.loads(s))
                except Exception:
                    pass

    assets = []
    for r in rows:
        try:
            ts = int(r["ts"])
            dt = datetime.fromtimestamp(ts, tz=timezone.utc).astimezone(tz.tzlocal())
            asset_id = f'{r.get("catalog_id","unknown")}-{ts}'
            snap = r.get("snapshot", {}) or {}
            assets.append({
                "asset_id": asset_id,
                "label": r.get("label", "(unknown)"),
                "catalog_id": r.get("catalog_id", ""),
                "purchase_ts": ts,
                "purchase_date": dt.strftime("%Y-%m-%d"),
                "cost": float(r.get("price", 0.0)),
                "expected_life_years": snap.get("expected_life_years") or 10,
                "salvage_value_pct": snap.get("salvage_value_pct") or 0.2,
                "dep_method_default": (snap.get("dep_method_default") or "SL").upper(),
            })
        except Exception:
            continue

    return assets

def straight_line_schedule(cost: float, life_years: int, salvage_value: float, start_year: int):
    dep = (cost - salvage_value) / life_years
    rows = []
    beg = cost
    acc = 0.0
    for i in range(1, life_years + 1):
        year = start_year + (i - 1)
        end = max(salvage_value, beg - dep)
        d = beg - end
        acc += d
        rows.append({
            "Year": year,
            "Begin BV": round(beg, 2),
            "Depreciation": round(d, 2),
            "Accum Dep": round(acc, 2),
            "End BV": round(end, 2),
        })
        beg = end
    return pd.DataFrame(rows)

def build_schedule(student_id, asset_pick, life_years, salvage_pct):
    if not student_id:
        return None, "Enter Student ID."

    assets = fetch_student_assets(student_id)
    idx = {a["asset_id"]: a for a in assets}
    if asset_pick not in idx:
        if not assets:
            return None, f"No purchases found for '{student_id}'."
        return None, "Pick an asset."

    a = idx[asset_pick]
    life = int(life_years) if life_years else int(a["expected_life_years"])
    spct = float(salvage_pct) if salvage_pct not in (None, "") else float(a["salvage_value_pct"])
    salvage_val = spct * a["cost"]

    df = straight_line_schedule(
        cost=a["cost"],
        life_years=life,
        salvage_value=salvage_val,
        start_year=int(a["purchase_date"][:4]),
    )

    # write schedule CSV back to dataset
    csv_bytes = df.to_csv(index=False).encode()
    api = _api_rw()
    path_in_repo = f"schedules/{student_id}/{asset_pick}.csv"
    api.upload_file(
        path_or_fileobj=io.BytesIO(csv_bytes),
        path_in_repo=path_in_repo,
        repo_id=DATA_REPO,
        repo_type=DATA_REPO_TYPE,
    )
    note = f"📄 Wrote schedule to dataset at {path_in_repo}"
    return df, f"✅ Schedule generated for {a['label']} — {note}"

def load_assets_for_ui(student_id):
    if not student_id:
        return gr.update(choices=[], value=None), "Enter Student ID then click “🔍 Load my assets”."
    assets = fetch_student_assets(student_id)
    choices = [(f"{a['label']} ({a['purchase_date']}) — ${a['cost']:,.0f}", a["asset_id"]) for a in assets]
    if not choices:
        # helpful hint: show what we looked for
        paths = _candidate_purchase_paths(student_id)
        hint = " none" if not paths else ("\n".join(f"- {p}" for p in paths))
        return gr.update(choices=[], value=None), f"No purchases found. Looked for:\n{hint}"
    return gr.update(choices=choices, value=choices[0][1]), f"Loaded {len(choices)} asset(s)."

# --------------------
# UI
# --------------------
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("## 🧮 Jerry — Depreciation Helper")

    with gr.Row():
        student_id = gr.Textbox(label="Student ID", placeholder="e.g., jeff")
        load_btn = gr.Button("🔍 Load my assets")

    asset_dd = gr.Dropdown(label="Select an asset")

    with gr.Row():
        life_years = gr.Number(label="Service life (years)", value=None, precision=0)
        salvage_pct = gr.Slider(label="Salvage % of cost", minimum=0, maximum=0.9, step=0.01, value=None)

    build_btn = gr.Button("Build Straight-Line Schedule")
    out_tbl = gr.Dataframe(label="Depreciation Schedule", interactive=False)
    status = gr.Markdown()

    load_btn.click(load_assets_for_ui, inputs=[student_id], outputs=[asset_dd, status])
    build_btn.click(build_schedule, inputs=[student_id, asset_dd, life_years, salvage_pct],
                    outputs=[out_tbl, status])

    # allow prefill via URL like ?student_id=jeff
    def prefill(req: gr.Request):
        sid = req.query_params.get("student_id")
        return sid or ""

    app.load(prefill, inputs=None, outputs=[student_id])

if __name__ == "__main__":
    app.launch()