tests/test_extensions.py

"""Tests for the backend extensions — plain asserts, run:
python tests/test_extensions.py

The real local model is never loaded: engine._llm is replaced with a fake so
generation-shaped tests stay fast and offline.
"""

import os
import sys
import tempfile

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import hashlib
import re
from datetime import date, timedelta

import numpy as np

import engine
import memory
import rag
from features import websearch
from features.health import plan_workout_slots


def _fake_embed(text: str, is_query: bool = False) -> np.ndarray:
    """Deterministic hashed bag-of-words vector so RAG tests run without
    downloading/loading the real nomic-embed model."""
    vec = np.zeros(256, dtype=np.float32)
    for tok in re.findall(r"[a-z0-9]+", text.lower()):
        vec[int(hashlib.md5(tok.encode()).hexdigest(), 16) % 256] += 1.0
    norm = np.linalg.norm(vec)
    return vec / norm if norm > 0 else vec


rag.embed = _fake_embed  # never touch the real embedder in tests

# Redirect the memory store to a temp file so tests don't clobber demo data.
_tmpdir = tempfile.mkdtemp()
memory.DATA_DIR = _tmpdir
memory.MEMORY_PATH = os.path.join(_tmpdir, "memory.json")


class _FakeLLM:
    """Stand-in for the llama.cpp model so engine tests need no GGUF download.
    Mentions 'chicken' and '$8,000' so existing content assertions hold."""

    TEXT = "Recommend chicken thighs this week; your used car goal of $8,000 is on track."

    def create_chat_completion(self, messages, max_tokens=384, temperature=0.4, stream=False):
        if stream:
            return (
                {"choices": [{"delta": {"content": w + " "}}]}
                for w in self.TEXT.split(" ")
            )
        return {"choices": [{"message": {"content": self.TEXT}}]}


# Inject the fake so get_llm() returns it instead of loading the real model.
engine._llm = _FakeLLM()
engine.MODEL_STATE = "ready"
engine.ACTIVE_BACKEND = "fake"


def test_seed_schema() -> None:
    mem = memory.reset_to_seed()
    assert mem["user_profile"]["city"] == "Toronto"
    assert mem["user_profile"]["address"] == ""
    assert len(mem["calendar"]) >= 3
    for ev in mem["calendar"]:
        assert set(ev) == {"id", "title", "date", "start", "end", "kind"}
        assert len(ev["id"]) == 8
        assert ev["kind"] in ("work", "class", "social", "other")
    assert mem["workout_schedule"] == {"days": ["mon", "wed", "sat"], "time": "07:00"}
    pays = mem["finances"]["monthly_payments"]
    assert sorted(p["amount"] for p in pays) == [45, 128, 650]
    g = mem["goals"][0]
    assert g["title"] == "Used car" and g["target_amount"] == 8000 and g["saved"] == 1200


def test_event_helpers_roundtrip() -> None:
    memory.reset_to_seed()
    mem = memory.add_event({"title": "Study group", "date": date.today().isoformat(), "start": "14:00", "end": "16:00", "kind": "social"})
    added = [e for e in mem["calendar"] if e["title"] == "Study group"]
    assert len(added) == 1 and len(added[0]["id"]) == 8
    n = len(mem["calendar"])
    mem = memory.delete_event(added[0]["id"])
    assert len(mem["calendar"]) == n - 1
    assert not [e for e in mem["calendar"] if e["title"] == "Study group"]


def test_schedule_payments_goal_profile_helpers() -> None:
    memory.reset_to_seed()
    mem = memory.set_workout_schedule(["Tuesday", "thu"], "18:30")
    assert mem["workout_schedule"] == {"days": ["tue", "thu"], "time": "18:30"}

    mem = memory.set_monthly_payments([{"name": "Gym", "amount": 35, "due_day": 5}])
    assert mem["finances"]["monthly_payments"] == [{"name": "Gym", "amount": 35, "due_day": 5}]

    mem = memory.upsert_goal({"title": "Trip", "target_amount": 2000, "saved": 100, "deadline": "2027-01"})
    trip = [g for g in mem["goals"] if g["title"] == "Trip"][0]
    mem = memory.upsert_goal({"id": trip["id"], "title": "Trip", "saved": 500})
    trip2 = [g for g in mem["goals"] if g["id"] == trip["id"]][0]
    assert trip2["saved"] == 500 and trip2["target_amount"] == 2000
    assert len(mem["goals"]) == 2  # seed goal + Trip, no dup

    mem = memory.set_profile({"city": "Ottawa", "income_monthly": 1600})
    assert mem["user_profile"]["city"] == "Ottawa"
    assert mem["finances"]["income_monthly"] == 1600
    assert mem["user_profile"]["name"] == "Awais Aziz"  # merge, not replace
    assert "income_monthly" not in mem["user_profile"]


def test_websearch_format() -> None:
    sample = [
        {"title": "FreshMart flyer", "snippet": "chicken $2.49/lb", "url": "https://example.com/a"},
        {"title": "City Grocer", "snippet": "beef $3.99/lb", "url": "https://example.com/b"},
    ]
    text = websearch.format_results(sample)
    assert "WEB SEARCH RESULTS:" in text and sample[0]["title"] in text
    assert websearch.format_results([]) == ""
    # Real network call: returns a list (possibly empty offline) and never raises.
    out = websearch.search_web("toronto grocery deals", max_results=2)
    assert isinstance(out, list) and len(out) <= 2
    assert all(set(r) == {"title", "snippet", "url"} for r in out)


def test_health_plan_slots() -> None:
    # Fixed Monday so weekday math is deterministic.
    monday = date(2026, 6, 8)
    schedule = {"days": ["mon", "wed"], "time": "07:00"}
    # Monday event overlaps 07:00-08:00 -> clash; Wednesday event is later -> free.
    events = [
        {"id": "aaaaaaaa", "title": "Early shift", "date": monday.isoformat(), "start": "06:30", "end": "09:00"},
        {"id": "bbbbbbbb", "title": "Lecture", "date": (monday + timedelta(days=2)).isoformat(), "start": "10:00", "end": "12:00"},
    ]
    slots = plan_workout_slots(schedule, events, today=monday)
    assert [s["day"] for s in slots] == ["mon", "wed"]
    mon_slot, wed_slot = slots
    assert mon_slot["free"] is False and mon_slot["clashes"][0]["title"] == "Early shift"
    assert wed_slot["free"] is True and wed_slot["clashes"] == []

    # Boundary: event ending exactly at slot start does not clash.
    events2 = [{"id": "c", "title": "X", "date": monday.isoformat(), "start": "06:00", "end": "07:00"}]
    assert plan_workout_slots(schedule, events2, today=monday)[0]["free"] is True
    # Event starting exactly at slot end does not clash.
    events3 = [{"id": "d", "title": "Y", "date": monday.isoformat(), "start": "08:00", "end": "09:00"}]
    assert plan_workout_slots(schedule, events3, today=monday)[0]["free"] is True
    # No preferred days -> no slots.
    assert plan_workout_slots({"days": [], "time": "07:00"}, events, today=monday) == []


def test_goal_and_meal_prompts() -> None:
    mem = memory.reset_to_seed()
    msgs = engine.build_prompt("goal", mem, "I want to buy a used car")
    assert msgs[0]["role"] == "system" and "Socratic" in msgs[0]["content"]
    assert "monthly_payments" in msgs[1]["content"] or "Rent" in msgs[1]["content"]

    msgs = engine.build_prompt("meal_photo", mem, "FOOD ITEMS:\nchicken thighs")
    assert "FOOD ITEMS" in msgs[1]["content"]

    # generation runs through the injected fake LLM
    out = list(engine.generate_stream(msgs, domain="meal_photo"))
    assert out and "chicken" in out[-1].lower()
    out = list(engine.generate_stream(engine.build_prompt("goal", mem, "car"), domain="goal"))
    assert out and "8,000" in out[-1]

    # extra_context is appended to the user message and still streams
    msgs = engine.build_prompt("chat", mem, "any deals?")
    extra = "WEB SEARCH RESULTS:\n1. test"
    chunks = list(engine.generate_stream(msgs, domain="chat", extra_context=extra))
    assert chunks


def test_blank_slate_prompt() -> None:
    """A brand-new user (no name set) gets a graceful, possessive-safe prompt."""
    blank = memory.empty_data()
    msgs = engine.build_prompt("chat", blank, "hello")
    sys_msg = msgs[0]["content"]
    assert "your own machine" in sys_msg and "'s own machine" not in sys_msg
    assert "=== YOUR MEMORY ===" in msgs[1]["content"]


def test_events_in_window() -> None:
    memory.reset_to_seed()
    memory.add_event({"title": "Future", "date": (date.today() + timedelta(days=3)).isoformat(), "start": "09:00", "end": "10:00"})
    memory.add_event({"title": "Far", "date": (date.today() + timedelta(days=30)).isoformat(), "start": "09:00", "end": "10:00"})
    evs = memory.events_in_window(7)
    titles = [e["title"] for e in evs]
    assert "Future" in titles and "Far" not in titles
    assert evs == sorted(evs, key=lambda e: (e["date"], e["start"]))


def test_profile_name_split_and_backfill() -> None:
    import json

    mem = memory.reset_to_seed()
    p = mem["user_profile"]
    assert p["first_name"] == "Awais" and p["last_name"] == "Aziz" and p["name"] == "Awais Aziz"
    assert p["address"] == "" and p["city"] == "Toronto" and p["postal_code"] == "" and p["country"] == "Canada"

    # Legacy file with only `name` upgrades on load. Missing fields backfill
    # to blank defaults (never demo persona data) — country stays "".
    p.pop("first_name"); p.pop("last_name"); p.pop("postal_code"); p.pop("country")
    p["name"] = "Jane Q Doe"
    with open(memory.MEMORY_PATH, "w", encoding="utf-8") as f:
        json.dump(mem, f)
    mem2 = memory.load()
    p2 = mem2["user_profile"]
    assert p2["first_name"] == "Jane" and p2["last_name"] == "Q Doe"
    assert p2["name"] == "Jane Q Doe" and p2["country"] == ""

    # set_profile recomputes name.
    mem3 = memory.set_profile({"first_name": "Ali", "last_name": "Khan", "postal_code": "M5V 1A1"})
    assert mem3["user_profile"]["name"] == "Ali Khan"
    assert mem3["user_profile"]["postal_code"] == "M5V 1A1"


def test_set_section() -> None:
    memory.reset_to_seed()
    meals = [{"date": "2026-01-01", "dish": "Toast", "ingredients": ["bread"]}]
    mem = memory.set_section("meals", meals)
    assert mem["meals"] == meals
    subs = [{"name": "OnlySub", "cost": 5, "last_used": "2026-01-01"}]
    mem = memory.set_section("subscriptions", subs)
    assert mem["finances"]["subscriptions"] == subs
    mem = memory.set_section("monthly_payments", [])
    assert mem["finances"]["monthly_payments"] == []
    try:
        memory.set_section("goals", [])
        raise AssertionError("expected ValueError")
    except ValueError:
        pass


def test_goal_id_backfill_and_delete() -> None:
    import json

    mem = memory.reset_to_seed()
    assert all(len(g["id"]) == 8 for g in mem["goals"])
    # Strip ids -> load backfills.
    for g in mem["goals"]:
        g.pop("id")
    with open(memory.MEMORY_PATH, "w", encoding="utf-8") as f:
        json.dump(mem, f)
    mem = memory.load()
    assert all(g.get("id") for g in mem["goals"])

    mem = memory.upsert_goal({"title": "Laptop", "target_amount": 1500})
    gid = [g for g in mem["goals"] if g["title"] == "Laptop"][0]["id"]
    mem = memory.upsert_goal({"id": gid, "saved": 300})
    assert [g for g in mem["goals"] if g["id"] == gid][0]["saved"] == 300
    n = len(mem["goals"])
    mem = memory.delete_goal(gid)
    assert len(mem["goals"]) == n - 1 and not [g for g in mem["goals"] if g["id"] == gid]


def test_rag_note_management() -> None:
    import json

    import rag

    rag.DATA_DIR = _tmpdir
    rag.STORE_PATH = os.path.join(_tmpdir, "longterm.json")
    if os.path.exists(rag.STORE_PATH):
        os.remove(rag.STORE_PATH)
    rag.ensure_seeded()

    notes = rag.list_notes()
    assert notes and all(set(n) == {"id", "text", "kind"} for n in notes)

    # id backfill for old stores
    raw = json.load(open(rag.STORE_PATH, encoding="utf-8"))
    for n in raw:
        n.pop("id", None)
    json.dump(raw, open(rag.STORE_PATH, "w", encoding="utf-8"))
    notes = rag.list_notes()
    assert all(n["id"] for n in notes)

    nid = notes[0]["id"]
    assert rag.update_note(nid, "Awais now loves cilantro") is True
    updated = [n for n in rag.list_notes() if n["id"] == nid][0]
    assert updated["text"] == "Awais now loves cilantro"
    assert updated["kind"] == notes[0]["kind"]
    assert rag.delete_note(nid) is True
    assert not [n for n in rag.list_notes() if n["id"] == nid]
    assert rag.delete_note("nope") is False and rag.update_note("nope", "x") is False


def test_build_prompt_domains() -> None:
    import engine

    mem = memory.reset_to_seed()
    msgs = engine.build_prompt("chat", mem, "what should I cook?", domains=["food"])
    body = msgs[1]["content"]
    assert "recent_meals" in body
    assert "subscriptions" not in body and "income_monthly" not in body

    sl = engine.slice_for_domains(mem, ["kitchen", "health"])
    assert "recent_meals" in sl and "workouts_last_14_days" in sl and "user_profile" in sl
    assert "finances" not in sl

    # None = full context unchanged
    full = engine.build_prompt("chat", mem, "hi")[1]["content"]
    assert "finances" in full and "recent_meals" in full


def test_chat_endpoint_refs() -> None:
    import app

    memory.reset_to_seed()
    out = list(app.chat("what should I cook?", "[]", "", '["kitchen"]'))
    assert out and len(out[-1]) > 20
    # backward compat: no refs arg
    out = list(app.chat("hello", "[]"))
    assert out


if __name__ == "__main__":
    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    for fn in fns:
        fn()
        print(f"PASS {fn.__name__}")
    print(f"All {len(fns)} tests passed.")