#!/usr/bin/python3
# ask.py: ask questions, on device, about your own recorded days and conversations. The meeting-notes
# "everything you captured, at your fingertips" feature, but local: it never sends a word off this Mac.
#
# Reads the organizer's diary.json (engine/organize.py) for conversations + summaries and the per-clip
# sidecars (.organized/<file>.json) for full transcripts, scopes them to the time the question implies
# ("today", "yesterday", "last night", "this week", a weekday, or everything), assembles the relevant
# notes to a budget, and asks a LOCAL model (Ollama, via local_ai.py) to answer using ONLY those notes.
# If no local model is installed it falls back to an on-device extractive answer (no network either way).
#
#   ask.py --rec-dir DIR --q "what did I talk about last night?" [--scope auto] [--with-transcripts]
#          [--now 2026-06-21T15:00:00] [--model qwen2.5]
#
# Returns JSON: {ok, answer, tier, model, scope:{label,start,end}, used:[...], considered}
# Pure stdlib. Cloud is never used here, by design.
import sys, os, json, re, datetime

HOME = os.path.expanduser("~")
DEFAULT_REC = os.path.join(HOME, "Recordings/Epiphani")
CONTEXT_BUDGET = 12000     # chars of notes to hand the local model (fits a small/medium model's window)

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
try:
    import local_ai          # on-device Ollama client (same engine dir)
except Exception:
    local_ai = None


# ----------------------------------------------------------------------------- diary access
def load_diary(rec_dir):
    try:
        with open(os.path.join(rec_dir, "diary.json"), encoding="utf-8") as f:
            return json.load(f)
    except (OSError, ValueError):
        return None


def _clip_dt(fname):
    m = re.match(r"(\d{4})-(\d{2})-(\d{2})_(\d{2})-(\d{2})-(\d{2})", fname or "")
    return datetime.datetime(*[int(x) for x in m.groups()]) if m else None


def conv_dt(day, conv):
    """Best-effort datetime for a conversation: first clip's filename, else the day's date."""
    for c in conv.get("clips") or []:
        dt = _clip_dt(c.get("file", ""))
        if dt:
            return dt
    try:
        return datetime.datetime.fromisoformat((day.get("date") or "") + "T00:00:00")
    except ValueError:
        return None


def sidecar_transcript(rec_dir, fname):
    try:
        with open(os.path.join(rec_dir, ".organized", fname + ".json"), encoding="utf-8") as f:
            return (json.load(f).get("transcript") or "").strip()
    except (OSError, ValueError):
        return ""


# ----------------------------------------------------------------------------- time scope
WEEKDAYS = ["monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"]


def parse_scope(question, now, explicit=None):
    """Return (label, start_dt, end_dt) from an explicit scope or words in the question.

    'last night' = yesterday 6pm -> today 6am. 'today'/'this morning' = since midnight.
    A weekday name = the most recent occurrence of that day. Default = the last 7 days.
    """
    q = (question or "").lower()
    sel = (explicit or "auto").lower()
    midnight = now.replace(hour=0, minute=0, second=0, microsecond=0)

    def has(*words):
        return any(w in q for w in words)

    if sel == "all" or has("all time", "ever", "everything i", "all my"):
        return ("everything", datetime.datetime(2000, 1, 1), now + datetime.timedelta(days=1))
    if sel == "today" or has("today", "this morning", "this afternoon", "earlier today", "so far today"):
        return ("today", midnight, now + datetime.timedelta(minutes=1))
    if sel == "last night" or has("last night", "yesterday evening", "last evening", "overnight"):
        start = (midnight - datetime.timedelta(days=1)).replace(hour=18)
        return ("last night", start, midnight.replace(hour=6))
    if sel == "yesterday" or has("yesterday"):
        y = midnight - datetime.timedelta(days=1)
        return ("yesterday", y, midnight)
    if sel == "this week" or has("this week"):
        start = midnight - datetime.timedelta(days=midnight.weekday())
        return ("this week", start, now + datetime.timedelta(minutes=1))
    if sel == "last week" or has("last week"):
        this_mon = midnight - datetime.timedelta(days=midnight.weekday())
        return ("last week", this_mon - datetime.timedelta(days=7), this_mon)
    for i, wd in enumerate(WEEKDAYS):
        if sel == wd or re.search(r"\b%s\b" % wd, q):
            delta = (midnight.weekday() - i) % 7 or 7      # most recent past occurrence
            day0 = midnight - datetime.timedelta(days=delta)
            return (wd.capitalize(), day0, day0 + datetime.timedelta(days=1))
    # default: a useful recent window rather than the entire archive
    return ("the last 7 days", midnight - datetime.timedelta(days=6), now + datetime.timedelta(minutes=1))


# ----------------------------------------------------------------------------- gather
def gather(rec_dir, diary, start, end, with_transcripts):
    """Conversations whose time falls in [start, end), newest first, with notes assembled to a budget."""
    items = []
    for day in (diary or {}).get("days", []):
        for conv in day.get("conversations", []):
            dt = conv_dt(day, conv)
            if dt is None or not (start <= dt < end):
                continue
            items.append((dt, day, conv))
    items.sort(key=lambda x: x[0], reverse=True)

    notes, used, total = [], [], 0
    for dt, day, conv in items:
        title = conv.get("title", "Conversation")
        head = "## %s — %s%s" % (dt.strftime("%a %b %d, %I:%M%p").replace(" 0", " "), title,
                                 (" (with %s)" % conv["who_with"]) if conv.get("who_with") else "")
        body = (conv.get("summary") or "").strip()
        actions = [a for a in (conv.get("action_items") or []) if a]
        if actions:
            body += "\nAction items: " + "; ".join(actions)
        if with_transcripts:
            tx = " ".join(sidecar_transcript(rec_dir, c.get("file", "")) for c in (conv.get("clips") or []))
            if tx.strip():
                body += "\nTranscript: " + tx.strip()
        chunk = head + "\n" + body
        if total + len(chunk) > CONTEXT_BUDGET and notes:
            break                                          # keep within the model's budget, newest first
        notes.append(chunk[:CONTEXT_BUDGET])
        total += len(chunk)
        # `key` is the conversation's first clip filename — the same identifier citations.py and
        # main.swift's openConvo() already use to resolve a conversation (App Atlas P0-2: this is
        # what lets the UI turn a source row into a real click-through into the conversation).
        clips = conv.get("clips") or []
        key = clips[0].get("file", "") if clips else ""
        used.append({"title": title, "when": dt.strftime("%Y-%m-%d %H:%M"),
                     "who_with": conv.get("who_with", ""), "key": key})
    return "\n\n".join(notes), used, len(items)


# ----------------------------------------------------------------------------- answer
ASK_SYS = (
    "You answer questions about the user's OWN recorded conversations and days. Use ONLY the notes "
    "provided; do not invent anything. Be concise, specific, and concrete (names, times, decisions, "
    "action items). If the notes do not contain the answer, say so plainly. When useful, mention which "
    "conversation (by its title) the answer comes from."
)


def _answer_local(question, scope_label, notes, model_req):
    """Ask the on-device model. Returns (answer, model) or (None, None) if no local model."""
    if local_ai is None:
        return None, None
    model = local_ai.pick_model(model_req)
    if not model:
        return None, None
    user = "Notes from %s:\n\n%s\n\nQuestion: %s" % (scope_label, notes, question)
    try:
        resp = local_ai._post("/api/chat", {
            "model": model,
            "messages": [{"role": "system", "content": ASK_SYS}, {"role": "user", "content": user}],
            "stream": False, "options": {"temperature": 0.2},
        })
        return (resp.get("message") or {}).get("content", "").strip(), model
    except Exception:
        return None, model


def _answer_extractive(question, notes, used):
    """No-LLM fallback: surface the most keyword-relevant notes. Still 100% on device."""
    if not used:
        return "I don't have any recorded conversations in that time range to answer from."
    words = set(re.findall(r"[a-z0-9]{3,}", question.lower())) - {
        "what", "did", "the", "about", "talk", "talked", "say", "said", "tell", "have", "with", "who",
        "when", "discuss", "discussed", "you", "your", "and", "for", "was", "were", "any"}
    blocks = notes.split("\n\n")
    scored = sorted(blocks, key=lambda b: sum(b.lower().count(w) for w in words), reverse=True)
    top = "\n\n".join(scored[:3]) if scored else notes[:1500]
    return ("No on-device language model is installed, so here are the most relevant notes "
            "(install one via the Models tab for a written answer):\n\n" + top)


def ask(rec_dir, question, scope=None, with_transcripts=False, now=None, model_req=None):
    now = now or datetime.datetime.now()
    diary = load_diary(rec_dir)
    if diary is None:
        return {"ok": False, "error": "No diary yet. Record and organize something first."}
    label, start, end = parse_scope(question, now, scope)
    notes, used, considered = gather(rec_dir, diary, start, end, with_transcripts)
    scope_obj = {"label": label, "start": start.isoformat(timespec="minutes"),
                 "end": end.isoformat(timespec="minutes")}
    if not used:
        return {"ok": True, "tier": "on-device", "answer":
                "I don't have any conversations recorded for %s." % label,
                "scope": scope_obj, "used": [], "considered": considered, "model": ""}

    answer, model = _answer_local(question, label, notes, model_req)
    if answer:
        tier = "on-device (Ollama, no network)"
    else:
        answer = _answer_extractive(question, notes, used)
        tier, model = "on-device (extractive, no model installed)", ""
    return {"ok": True, "answer": answer, "tier": tier, "model": model or "",
            "scope": scope_obj, "used": used, "considered": considered}


# ----------------------------------------------------------------------------- cli
def _opt(argv, flag, default=None):
    return argv[argv.index(flag) + 1] if flag in argv and argv.index(flag) + 1 < len(argv) else default


def main(argv):
    rec_dir = os.path.expanduser(_opt(argv, "--rec-dir", DEFAULT_REC))
    question = _opt(argv, "--q", "")
    scope = _opt(argv, "--scope", None)
    model_req = _opt(argv, "--model", None)
    now = None
    if _opt(argv, "--now"):
        try:
            now = datetime.datetime.fromisoformat(_opt(argv, "--now"))
        except ValueError:
            pass
    if not question:
        sys.stdout.write(json.dumps({"ok": False, "error": "no question (--q)"}))
        return
    rep = ask(rec_dir, question, scope=scope, with_transcripts=("--with-transcripts" in argv),
              now=now, model_req=model_req)
    sys.stdout.write(json.dumps(rep, ensure_ascii=False))


if __name__ == "__main__":
    main(sys.argv[1:])
