"""The companion's tools over the knowledge base.

Typed per question-shape rather than one generic `search()`, because the useful
*return fields* differ: an office needs a room and a phone number, an event needs
a date and a contact, a degree program needs the catalog's verbatim rule text. A
single tool returning a blob would push all of that back onto the model.

Every tool returns compact dicts, hard-capped in size. Tool payloads are the main
cost driver of a turn and a 4,000-character course description helps nobody — so
snippets are truncated and result counts are bounded, here rather than by asking
the model nicely.

Each call also records the documents it returned into a per-turn collector, which
is how the server later proves which of the model's links actually came from
retrieval instead of trusting it to cite honestly.
"""
from __future__ import annotations

import contextvars

from langchain_core.tools import tool

from .. import kb
from ..kb import offerings, queries

MAX_RESULTS = 10
SNIPPET = 320

# Per-turn scratch space. ContextVars rather than globals so concurrent students
# never see each other's retrieval.
_retrieved: contextvars.ContextVar[dict] = contextvars.ContextVar("retrieved")
_suggestion: contextvars.ContextVar[dict] = contextvars.ContextVar("suggestion")


def start_turn() -> tuple[dict, dict]:
    """(retrieved documents, suggestion slot) for the turn about to run."""
    retrieved: dict = {}
    suggestion: dict = {}
    _retrieved.set(retrieved)
    _suggestion.set(suggestion)
    return retrieved, suggestion


def _record(docs) -> None:
    try:
        store = _retrieved.get()
    except LookupError:
        return
    for d in docs:
        store[d.id] = d


def _clamp(limit: int | None, default: int = 6) -> int:
    return max(1, min(int(limit or default), MAX_RESULTS))


def _hit(doc, query: str = "") -> dict:
    out = {
        "id": doc.id,
        "title": doc.title,
        "kind": doc.kind,
        "source": kb.SOURCE_LABELS.get(doc.source, doc.source),
        "text": kb.snippet(doc, query, SNIPPET),
    }
    if doc.url:
        out["url"] = doc.url
    if doc.domains:
        out["domains"] = list(doc.domains)
    return out


def _search(query, *, limit, **filters):
    hits = kb.get_index().search(query or "", limit=_clamp(limit), **filters)
    docs = [d for d, _ in hits]
    _record(docs)
    return docs


def _empty(what: str) -> dict:
    return {"results": [], "note": f"Nothing in the knowledge base matched. {what}"}


# --- general ----------------------------------------------------------------
@tool
def search_campus(query: str, domains: list[str] | None = None,
                  kinds: list[str] | None = None, limit: int = 6) -> dict:
    """Search all Vanderbilt campus information: student organizations, offices,
    resources, housing, study abroad, immersion, scholarships, and guidance pages.

    This is the general-purpose tool — use it whenever a more specific tool
    doesn't obviously fit.

    Args:
        query: What to look for, in the student's own words.
        domains: Optionally narrow to First View domains: strengths, crew, future, vu.
        kinds: Optionally narrow by record type, e.g. organization, office, page,
            resource, house, policy, award.
        limit: How many results to return (max 10).
    """
    docs = _search(query, limit=limit, domains=domains, kinds=kinds)
    if not docs:
        return _empty("Try different words, or use web_search if it may not be campus information.")
    return {"results": [_hit(d, query) for d in docs]}


# --- events -----------------------------------------------------------------
@tool
def find_events(query: str = "", from_date: str = "", to_date: str = "",
                domains: list[str] | None = None, limit: int = 8) -> dict:
    """Find upcoming campus events — student org events and the university calendar,
    already merged and deduplicated.

    Leave `query` empty to simply list what's happening in a date range.

    Args:
        query: Optional topic, e.g. "running", "career fair", "free food".
        from_date: ISO date (YYYY-MM-DD). Defaults to today.
        to_date: ISO date (YYYY-MM-DD). Defaults to 30 days out.
        domains: Optionally narrow to strengths, crew, future, vu.
        limit: How many events to return (max 10).
    """
    start, end = queries.window(from_date, to_date, days=queries.EVENT_WINDOW_DAYS)
    hits = queries.events_between(kb.get_index(), start, end, query=query,
                                  domains=domains, limit=_clamp(limit, 8))
    _record([h.doc for h in hits])
    if not hits:
        return _empty(f"No events found between {start} and {end}.")

    out = []
    for h in hits:
        d = h.doc
        item = _hit(d, query)
        # The occurrence that matched the range, not the day the series began.
        item["starts"] = h.when
        for key in ("location", "contact_email", "cost", "has_registration"):
            if d.extra.get(key) is not None:
                item[key] = d.extra[key]
        hosts = d.extra.get("hosts")
        if hosts:
            item["hosts"] = hosts
        # A recurring event is one record with many dates; show that it repeats
        # rather than listing it once per occurrence.
        upcoming = queries.upcoming_occurrences(d, start)
        if len(upcoming) > 1:
            item["repeats"] = f"{len(upcoming)} more dates"
            item["next_dates"] = upcoming[:5]
        out.append(item)
    return {"range": {"from": start, "to": end}, "results": out}


# --- places -----------------------------------------------------------------
@tool
def find_place(query: str, limit: int = 5) -> dict:
    """Find a named campus office, department, or building — where it is and how
    to reach it. Returns room, address, phone, email and map link when known.

    Best with a name you already have. If the student described a *need* rather
    than a place ("help with a paper", "I can't afford my books"), call
    `search_campus` first to find out which office handles it, then call this with
    that office's name.

    Args:
        query: An office, department or building name, e.g. "Writing Studio",
            "Student Care Network", "Featheringill Hall".
        limit: How many places to return (max 10).
    """
    index = kb.get_index()
    docs = _search(query, limit=limit, kinds=["office", "building", "contact"])
    if not docs:
        return _empty("Try the building name, or search_campus for the service itself.")

    out = []
    for d in docs:
        item = _hit(d, query)
        for key in ("building_name", "room", "address", "phone", "email",
                    "hours", "office", "map_url", "lat", "lng"):
            if d.extra.get(key):
                item[key] = d.extra[key]
        # An office record names its building; pull the building's address and map
        # link across so the student gets one complete answer.
        slug = d.extra.get("building_slug")
        if slug:
            building = index.get(f"buildings:building:{slug}")
            if building:
                item.setdefault("address", building.extra.get("address"))
                item["map_url"] = building.extra.get("map_url")
                item["building_name"] = building.title
                _record([building])
        elif d.extra.get("location_unknown"):
            # 447 offices have contact details but no confidently resolved
            # building. Say so rather than letting the model fill the gap.
            item["location_note"] = ("No building or room is published for this "
                                     "office — give the contact details and say the "
                                     "location isn't listed.")
        out.append(item)
    return {"results": out}


# --- academics --------------------------------------------------------------
@tool
def lookup_courses(query: str = "", subject: str = "", course_id: str = "",
                   limit: int = 6) -> dict:
    """Look up Vanderbilt courses — descriptions, credit hours, AXLE/CORE tags,
    and prerequisites, from the undergraduate catalog.

    Note this is the catalog of every course that exists, not the list of what is
    actually offered in a given term — use `find_offered_classes` for that.

    Args:
        query: Topic or title, e.g. "intro psychology", "data science".
        subject: Optional subject code, e.g. "BSCI", "CS", "ECON".
        course_id: An exact course id, e.g. "CORE1010", for a direct lookup.
        limit: How many courses to return (max 10).
    """
    index = kb.get_index()
    if course_id:
        doc = index.get(f"academics:course:{course_id.upper().replace(' ', '')}")
        if doc:
            _record([doc])
            item = _hit(doc, query)
            item.update({k: v for k, v in doc.extra.items()
                         if k in ("credit_hours", "axle", "core", "relations",
                                  "subject_name", "description") and v})
            return {"results": [item], "note": "Catalog listing — use find_offered_classes for the term schedule."}

    text = " ".join(p for p in (query, subject) if p)
    docs = _search(text, limit=limit, kinds=["course"])
    if not docs:
        return _empty("Try a subject code, or a broader topic.")
    out = []
    for d in docs:
        item = _hit(d, text)
        for key in ("credit_hours", "axle", "core", "subject_name"):
            if d.extra.get(key):
                item[key] = d.extra[key]
        out.append(item)
    return {"results": out,
            "note": "Catalog listings — use find_offered_classes for what's actually offered this term."}


@tool
def find_offered_classes(query: str = "", subject: str = "", term: str = "",
                         level: str = "", limit: int = 6) -> dict:
    """What is actually being offered in a term — sections, meeting times,
    instructors, and enrollment status, from Vanderbilt's public class schedule.

    Use this for "what can I take this fall", "is X offered", "who's teaching X",
    "when does X meet". The knowledge base holds the current and upcoming terms;
    `lookup_courses` is the full catalog of what *exists* instead.

    Enrollment counts are a snapshot from the last refresh, not live — say "as of
    the last check" and remind the student that registration itself happens in YES.

    Args:
        query: Topic, title words, or a course id, e.g. "machine learning", "DS 1000".
        subject: Optional subject code to narrow, e.g. "DS", "CS", "MATH".
        term: Optional term words, e.g. "fall", "2027 spring". Empty = any term on file.
        level: Optional level: "undergraduate" or "graduate".
        limit: How many courses to return (max 10).
    """
    # The selection logic lives in kb/offerings.py, shared with the My VU
    # page's /api/classes. The tool does NOT exclude the med-school "Year"
    # pseudo-term — Ask Foresight can still answer a med student; the
    # student-facing endpoint does.
    items, docs, note = offerings.search(
        kb.get_index(), query=query, subject=subject, term=term,
        level=level, limit=limit)
    if not items:
        return _empty("Try a subject code (e.g. 'DS'), fewer filters, or "
                      "lookup_courses if it may not run this term.")
    _record(docs)
    return {"results": items, "note": note}


@tool
def program_requirements(program: str) -> dict:
    """Get the degree requirements for a Vanderbilt major, minor, or program.

    Returns the catalog's own wording for each requirement group. Quote it rather
    than tidying it up, and tell the student to confirm with their adviser —
    requirements depend on the year they matriculated.

    Args:
        program: The major, minor or program name, e.g. "economics major",
            "computer science", "medicine health and society minor".
    """
    docs = _search(program, limit=3, kinds=["program"])
    if not docs:
        return _empty("Try the major's name on its own.")

    top = docs[0]
    groups = top.extra.get("requirements") or []
    out = {
        "program": top.title,
        "schools": top.extra.get("schools") or [],
        "kind": top.extra.get("program_kind"),
        "total_hours": top.extra.get("total_hours"),
        "catalog": "Undergraduate Catalog 2025-26",
        "other_matches": [d.title for d in docs[1:]],
    }
    if top.url:
        out["url"] = top.url
    if not groups:
        out["note"] = ("No structured requirements on file for this program — "
                       "point the student to their school's academic office.")
        return out
    out["requirement_groups"] = [{
        "label": g.get("label"),
        "select": g.get("select"),
        "hours": g.get("hours"),
        # Verbatim, deliberately. The catalog's own defects are preserved at build
        # time; smoothing them here would invent requirements.
        "rule_text": g.get("rule_text"),
        "courses": (g.get("courses") or [])[:40],
        "subgroups": [{"label": s.get("label"), "courses": (s.get("courses") or [])[:20]}
                      for s in (g.get("subgroups") or [])],
    } for g in groups[:12]]
    return out


# --- deadlines --------------------------------------------------------------
@tool
def find_deadlines(query: str = "", from_date: str = "", to_date: str = "") -> dict:
    """Find deadlines and key dates — scholarships and fellowships, academic
    calendar milestones like add/drop and registration, and housing dates.

    Use this for "when is X due" and "am I too late for Y".

    Args:
        query: Optional topic, e.g. "fulbright", "add drop", "housing".
        from_date: ISO date (YYYY-MM-DD) for calendar milestones. Defaults to today.
        to_date: ISO date (YYYY-MM-DD). Defaults to one year out.
    """
    start, end = queries.window(from_date, to_date, days=queries.DEADLINE_WINDOW_DAYS)

    dated = queries.key_dates_between(kb.get_index(), start, end, query=query, limit=8)
    _record([h.doc for h in dated])
    awards = _search(query, limit=6, kinds=["award", "research_program"])
    processes = _search(query, limit=3, kinds=["process"]) if query else []

    if not (dated or awards or processes):
        return _empty("Try the name of the award or the process.")

    out: dict = {"range": {"from": start, "to": end}}
    if dated:
        out["academic_calendar"] = [
            {"title": h.doc.title,
             "date": h.doc.extra.get("date") or (h.when or "")[:10],
             "iso_date": (h.when or "")[:10],
             "terms": h.doc.extra.get("terms") or [], "url": h.doc.url}
            for h in dated]
    if awards:
        out["awards"] = [
            {"name": d.title,
             "campus_deadline": d.extra.get("campus_deadline"),
             "priority_deadline": d.extra.get("priority_deadline"),
             "eligibility": d.extra.get("eligibility"),
             "applicant_level": d.extra.get("applicant_level"),
             "tier": d.extra.get("tier"),
             "url": d.url}
            for d in awards]
        out["deadline_note"] = ("Award deadlines are stated as months/days without a "
                                "year — say which cycle you mean and tell the student "
                                "to confirm on the linked page.")
    if processes:
        out["processes"] = [
            {"process": d.title, "summary": d.extra.get("summary"),
             "key_dates": d.extra.get("key_dates"), "url": d.url}
            for d in processes]
    return out


# --- how to engage / routing ------------------------------------------------
@tool
def how_to_engage(topic: str, limit: int = 6) -> dict:
    """Find the actual next action for something: application portals, advising
    appointment links, forms, and where to submit things.

    Also the right tool for anything behind a Vanderbilt login — YES, Brightspace,
    Handshake, the Housing Portal, InfoReady. You cannot see inside those, but this
    returns the entry point so you can tell the student where to go and what to
    look for once they're there.

    Args:
        topic: What the student is trying to do, e.g. "register for classes",
            "apply for housing", "submit a fellowship application", "see an adviser".
        limit: How many entry points to return (max 10).
    """
    docs = _search(topic, limit=limit, kinds=["resource", "contact", "process"])
    if not docs:
        return _empty("Use find_place to name the office that owns this instead.")
    out = []
    for d in docs:
        item = {"what": d.title, "url": d.url,
                "kind": d.extra.get("resource_kind") or d.kind,
                "source": kb.SOURCE_LABELS.get(d.source, d.source)}
        for key in ("phone", "email", "hours", "summary", "steps"):
            value = d.extra.get(key)
            if value:
                item[key] = value[:4] if key == "steps" else value
        out.append(item)
    return {"results": out,
            "reminder": ("You cannot log in to any of these on the student's behalf. "
                         "Tell them what to look for once they're signed in.")}


# --- the proactive nudge ----------------------------------------------------
# Every route here exists in the app's hash router (see app/app.js). A route the
# router doesn't know would render a "not found" screen, so this is an enum.
ROUTES = {
    "#/today": "the Today dashboard",
    "#/calendar": "the calendar",
    "#/campus": "the campus map",
    "#/prep": "Email Prep — help drafting emails and conversations",
    "#/story": "My Story — their record of what they've done",
    "#/domain/strengths": "My Strengths",
    "#/domain/crew": "My Crew",
    "#/domain/future": "My Future",
    "#/domain/vu": "My VU",
}


@tool
def suggest_next_step(label: str, route: str) -> dict:
    """Attach one suggested next action to your answer, shown as a button.

    Use it when there's a genuinely useful next move inside Foresight — not on
    every turn, and at most once per answer. Skip it if the answer is complete on
    its own.

    Args:
        label: The button text, written as an invitation, e.g. "Draft that email"
            or "See this week's events". Keep it under about six words.
        route: Where the button goes. One of: #/today, #/calendar, #/campus,
            #/prep, #/story, #/domain/strengths, #/domain/crew, #/domain/future,
            #/domain/vu.
    """
    if route not in ROUTES:
        return {"ok": False,
                "error": f"'{route}' isn't a real screen. Choose one of: {', '.join(ROUTES)}"}
    suggestion = {"label": label.strip()[:60], "route": route}
    # Hand it to the turn directly rather than making the caller re-parse this
    # result out of a stringified ToolMessage.
    try:
        _suggestion.get().update(suggestion)
    except LookupError:
        pass
    return {"ok": True, "suggestion": suggestion}


KB_TOOLS = [search_campus, find_events, find_place, lookup_courses,
            find_offered_classes, program_requirements, find_deadlines,
            how_to_engage, suggest_next_step]

# Student-legible status text. The wait is several seconds of tool calls, and
# "Thinking…" tells a student nothing about whether it's working.
TOOL_LABELS = {
    "search_campus": "Searching campus resources…",
    "find_events": "Checking what's happening on campus…",
    "find_place": "Finding the place and who to ask…",
    "lookup_courses": "Looking through the course catalog…",
    "find_offered_classes": "Checking the class schedule…",
    "program_requirements": "Pulling up the degree requirements…",
    "find_deadlines": "Checking deadlines and key dates…",
    "how_to_engage": "Finding where to actually do this…",
    "suggest_next_step": "",
    "web_search": "Looking this up on the web…",
    "web_search_call": "Looking this up on the web…",
}