import json, random
from collections import deque
import numpy as np
import faiss, gradio as gr
from sentence_transformers import SentenceTransformer
from huggingface_hub import hf_hub_download

REPO_ID = "Elevi7/actionmatch-microactions-en"

index_path = hf_hub_download(repo_id=REPO_ID, filename="index/index.faiss", repo_type="dataset")
actions_path = hf_hub_download(repo_id=REPO_ID, filename="actions.jsonl", repo_type="dataset")

with open(actions_path, "r", encoding="utf-8") as f:
    actions = [json.loads(l) for l in f]

index = faiss.read_index(index_path)
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
recent = deque(maxlen=120)

def render(a):
    ctx = ", ".join(a["context"]) if isinstance(a.get("context"), list) else (a.get("context") or "")
    return f"**{a['title']}**  \n{a['instruction']}  \nGoal: {a['goal']} • Duration: {a['duration_min']} min • Energy: {a['energy']} • Context: {ctx}"

def pick_unique(idxs, goal, energy, max_minutes, ignore_goal=False, ignore_energy=False, ignore_minutes=False, need=3, used=None):
    if used is None:
        used = set()
    out = []
    for i in idxs:
        a = actions[i]
        if a["title"] in used or a["title"] in recent:
            continue
        if not ignore_goal and goal and a["goal"] != goal:
            continue
        if not ignore_energy and energy and a["energy"] != energy:
            continue
        if not ignore_minutes and max_minutes and a["duration_min"] > int(max_minutes):
            continue
        used.add(a["title"])
        out.append(a)
        if len(out) == need:
            break
    return out, used

def fill_random(need, used, goal, energy, max_minutes):
    pool = [a for a in actions if a["title"] not in used and a["title"] not in recent and (not goal or a["goal"]==goal) and (not energy or a["energy"]==energy) and (not max_minutes or a["duration_min"]<=int(max_minutes))]
    if len(pool) < need:
        pool = [a for a in actions if a["title"] not in used and a["title"] not in recent]
    random.shuffle(pool)
    return pool[:need]

def search(query, goal, energy, max_minutes):
    try:
        q = (query or "").strip()
        qx = f"{q} Goal:{goal or 'any'} Energy:{energy or 'any'} Max:{int(max_minutes) if max_minutes else ''} minutes"
        v = model.encode([qx], normalize_embeddings=True)
        D, I = index.search(np.asarray(v, dtype="float32"), 800)
        idxs = list(I[0]); random.shuffle(idxs)
        res, used = [], set()
        step, used = pick_unique(idxs, goal, energy, max_minutes, False, False, False, 3, used); res += step
        if len(res) < 3:
            step, used = pick_unique(idxs, goal, energy, max_minutes, False, True, False, 3-len(res), used); res += step
        if len(res) < 3:
            step, used = pick_unique(idxs, goal, energy, max_minutes, False, True, True, 3-len(res), used); res += step
        if len(res) < 3:
            step, used = pick_unique(idxs, goal, energy, max_minutes, True, True, True, 3-len(res), used); res += step
        if len(res) < 3:
            res += fill_random(3-len(res), used, goal, energy, max_minutes)
        recent.extend([a["title"] for a in res[:3]])
        return "\n\n---\n\n".join(render(a) for a in res[:3])
    except Exception:
        pool = [a for a in actions if (not goal or a["goal"]==goal) and (not energy or a["energy"]==energy) and (not max_minutes or a["duration_min"]<=int(max_minutes))]
        if len(pool) < 3:
            pool = actions[:]
        random.shuffle(pool)
        return "\n\n---\n\n".join(render(a) for a in pool[:3])

goals = ["","calm","focus","productivity","wellbeing"]
energies = ["","low","medium","high"]

with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")) as demo:
    gr.Markdown("# ActionMatch\nTop-3 micro-actions based on your situation, goal, energy and time.")
    with gr.Row():
        with gr.Column(scale=3):
            q = gr.Textbox(lines=2, label="Your situation", placeholder="e.g., Stressed before exam")
            btn = gr.Button("Recommend")
        with gr.Column(scale=2):
            g = gr.Dropdown(goals, label="Goal")
            e = gr.Radio(energies, label="Energy")
            m = gr.Slider(1, 15, step=1, value=5, label="Max minutes")
    out = gr.Markdown()
    btn.click(search, [q, g, e, m], out)
    gr.Examples(
        [["Stressed before exam","calm","low",5],
         ["No energy but need to start studying","focus","low",7],
         ["Keep switching tabs while writing essay","focus","medium",10]],
        inputs=[q, g, e, m],
        outputs=out,
        fn=search,
        label="One-click examples",
        cache_examples=False
    )

demo.queue()
demo.launch()