"""Dev-time verification of LifeOS model stack on Modal (Spaces-like CPU env)."""
import modal

app = modal.App("lifeos-check")

# Spaces-like image: plain pip from PyPI with apt build tools (source build).
image = (
    modal.Image.debian_slim(python_version="3.13")
    .apt_install("build-essential", "cmake", "ninja-build", "git")
    .pip_install("llama-cpp-python==0.3.28", "huggingface-hub", "numpy")
)

CHAT_REPO = "nvidia/NVIDIA-Nemotron-3-Nano-4B-GGUF"
CHAT_FILE = "NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
FALLBACK_REPO = "bartowski/nvidia_Llama-3.1-Nemotron-Nano-4B-v1.1-GGUF"
FALLBACK_FILE = "nvidia_Llama-3.1-Nemotron-Nano-4B-v1.1-Q4_K_M.gguf"
EMB_REPO = "nomic-ai/nomic-embed-text-v1.5-GGUF"
EMB_FILE = "nomic-embed-text-v1.5.Q8_0.gguf"


@app.function(image=image, cpu=2.0, memory=16384, timeout=1800)
def check():
    import time
    import traceback

    import numpy as np
    from huggingface_hub import hf_hub_download
    from llama_cpp import Llama

    results = {}

    # ---------- Chat model ----------
    used_model = None
    llm = None
    try:
        print(f"Downloading {CHAT_REPO}/{CHAT_FILE} ...")
        path = hf_hub_download(repo_id=CHAT_REPO, filename=CHAT_FILE)
        print("Loading Nemotron-3-Nano-4B ...")
        llm = Llama(model_path=path, n_ctx=8192, n_threads=2, verbose=True)
        used_model = "NEMOTRON3 (primary)"
    except Exception:
        print("PRIMARY MODEL FAILED:")
        traceback.print_exc()
        print(f"Falling back to {FALLBACK_REPO}/{FALLBACK_FILE} ...")
        path = hf_hub_download(repo_id=FALLBACK_REPO, filename=FALLBACK_FILE)
        llm = Llama(model_path=path, n_ctx=8192, n_threads=2, verbose=True)
        used_model = "FALLBACK (Llama-3.1-Nemotron-Nano-4B-v1.1)"

    print(f"\n=== MODEL LOADED: {used_model} ===\n")

    messages = [
        {"role": "system", "content": "/no_think\nYou are LifeOS, a concise local assistant."},
        {"role": "user", "content": (
            "Recommend tomorrow's workout. Last 7 days: run 30min (Mon), "
            "push 50min (Tue), pull 45min (Thu), run 25min (Fri), "
            "legs 55min (Sat=yesterday). Goal: build muscle + 10K in September."
        )},
    ]

    t0 = time.time()
    out_text = ""
    n_tokens = 0
    for chunk in llm.create_chat_completion(messages, max_tokens=256, stream=True):
        delta = chunk["choices"][0]["delta"]
        if "content" in delta and delta["content"]:
            out_text += delta["content"]
            n_tokens += 1
    dt = time.time() - t0
    tps = n_tokens / dt if dt > 0 else 0.0

    print("\n=== GENERATED OUTPUT ===")
    print(out_text)
    print("=== END OUTPUT ===")
    print(f"Tokens generated: {n_tokens}")
    print(f"Time: {dt:.1f}s -> {tps:.2f} tokens/sec")

    del llm

    # ---------- Embedding model ----------
    print("\nDownloading nomic-embed-text-v1.5 Q8_0 ...")
    emb_path = hf_hub_download(repo_id=EMB_REPO, filename=EMB_FILE)
    emb = Llama(model_path=emb_path, embedding=True, n_threads=2, verbose=False)

    def vec(s):
        v = np.array(emb.create_embedding(s)["data"][0]["embedding"], dtype=np.float32)
        if v.ndim > 1:
            v = v.mean(axis=0)
        return v / np.linalg.norm(v)

    a = vec("I went for a 5km run this morning")
    b = vec("Morning jog around the park, about five kilometers")
    c = vec("The stock market dropped sharply today")

    sim_ab = float(np.dot(a, b))
    sim_ac = float(np.dot(a, c))
    print(f"cos(similar pair)   = {sim_ab:.4f}")
    print(f"cos(dissimilar pair)= {sim_ac:.4f}")
    sanity = "PASS" if sim_ab > sim_ac else "FAIL"
    print(f"Embedding sanity (similar > dissimilar): {sanity}")

    print("\n=== SUMMARY ===")
    print(f"Chat model used : {used_model}")
    print(f"Tokens/sec      : {tps:.2f}")
    print(f"Embedding sanity: {sanity}")
    return {"model": used_model, "tps": tps, "tokens": n_tokens,
            "sanity": sanity, "sim_ab": sim_ab, "sim_ac": sim_ac, "text": out_text}


@app.local_entrypoint()
def main():
    res = check.remote()
    print("\nLOCAL SUMMARY:", {k: v for k, v in res.items() if k != "text"})