#!/usr/bin/env python
"""The Web Research Agent (#13) — on-demand, cited, multi-step research.

Modes (same primitives, more orchestration each step):
  quick : search → read top pages → cited answer            (fast; the public-demo mode)
  deep  : plan sub-questions → research each → reflect on gaps → one more round → cited report
  swarm : deep, but sub-questions are gathered in PARALLEL (bounded, fleet-friendly)

Grounding is enforced like our #10 RAG app: the writer may use ONLY the numbered sources it actually
read, must cite inline [n], and must say what's missing rather than invent. Citations are validated
back against the sources that were fetched.

CLI:  python agent.py "your question"  [--mode deep] [--area technical]
"""
import concurrent.futures as cf
import os
import re

import llm
import search as search_mod
import fetch as fetch_mod
import sources as sources_mod
import memory

CONCURRENCY   = int(os.environ.get("RESEARCH_CONCURRENCY", "4"))   # parallel gather workers
READ_PER_SUBQ = int(os.environ.get("RESEARCH_READ_PER_SUBQ", "3")) # pages read per sub-question
SOURCE_CHARS  = 1200   # excerpt per source handed to the synthesizer


# ---------------------------------------------------------------- planning
def plan(question: str, breadth: int = 4) -> list:
    sys = ("You are a research planner. Break the QUESTION into distinct web-search queries that together "
           "cover it. Write them the way a person types into a search engine: plain keywords, 3-8 words "
           "each, NO quotation marks, NO website/domain names. Keep the key nouns so meaning isn't lost. "
           "Return a JSON array of query strings.")
    qs = llm.chat_json(sys, f"QUESTION: {question}\n\nGive {breadth} search queries as a JSON array.")
    out = [q.strip().strip('"') for q in qs if isinstance(q, str) and q.strip()] if isinstance(qs, list) else []
    # only echo the original phrasing when it's short enough to be a decent query itself
    if len(question) <= 140 and question not in out:
        out = [question] + out
    return out[:breadth] or [question[:140]]


# ---------------------------------------------------------------- gather (a "sub-agent")
def gather(query: str, area: str = "web", n_read: int = READ_PER_SUBQ) -> list:
    """One unit of research: search a query, read the top readable pages. Returns source dicts
    {title,url,text,query}. Network-bound → safe to run many in parallel."""
    results = sources_mod.search_area(query, area, n=n_read + 3)
    out = []
    for r in results:
        if len(out) >= n_read:
            break
        f = fetch_mod.fetch(r["url"], max_chars=SOURCE_CHARS)
        if not f["error"] and f["text"]:
            out.append({"title": f["title"] or r["title"], "url": f["final_url"],
                        "text": f["text"], "query": query})
    return out


def _gather_many(queries, area, trace):
    """Run gather() across queries in parallel (bounded). Dedup by URL."""
    sources, seen = [], set()
    with cf.ThreadPoolExecutor(max_workers=CONCURRENCY) as ex:
        futs = {ex.submit(gather, q, area): q for q in queries}
        for fut in cf.as_completed(futs):
            q = futs[fut]
            try:
                got = fut.result()
            except Exception:
                got = []
            trace.append(f"searched: {q}  →  read {len(got)} page(s)")
            for s in got:
                key = s["url"].split("#")[0].rstrip("/")
                if key not in seen:
                    seen.add(key)
                    sources.append(s)
    return sources


# ---------------------------------------------------------------- synthesis (grounded + cited)
_SYNTH_SYS = (
    "You are a careful research analyst. Write a briefing that answers the QUESTION using ONLY the "
    "numbered SOURCES provided — sources you were actually given. Rules:\n"
    "- Cite every factual claim inline as [n] using the source numbers. Use only numbers that exist.\n"
    "- Lead with a direct 2-3 sentence answer, then supporting detail in short sections or bullets.\n"
    "- If sources disagree, say so and cite both. If the sources don't fully answer the question, state "
    "plainly what's missing — never invent facts or cite a source that doesn't support the claim.\n"
    "- Be concise and concrete. No preamble like 'Here is'."
)


def _domain(url: str) -> str:
    from urllib.parse import urlparse
    d = urlparse(url).netloc.lower()
    return d[4:] if d.startswith("www.") else d


def _used_from(report: str, n_sources: int) -> list:
    """Parse cited source numbers, handling [n], [n, m] and ranges [n-m]."""
    used = set()
    for grp in re.findall(r"\[([\d,\s–-]+)\]", report):
        for part in grp.split(","):
            part = part.strip().replace("–", "-")
            if "-" in part:
                try:
                    a, b = (int(x) for x in part.split("-", 1))
                    used.update(range(a, b + 1))
                except Exception:
                    pass
            elif part.isdigit():
                used.add(int(part))
    return sorted(n for n in used if 1 <= n <= n_sources)


def synthesize(question: str, sources: list) -> dict:
    if not sources:
        return {"report": "_No readable sources were found for this question._", "used": [], "sources": []}
    block = "\n\n".join(f"[{i+1}] {s['title']} — {s['url']}\n{s['text'][:SOURCE_CHARS]}"
                        for i, s in enumerate(sources))
    user = f"QUESTION: {question}\n\nSOURCES:\n{block}\n\nWrite the cited briefing."
    report = llm.chat(_SYNTH_SYS, user, temperature=0.3, max_tokens=1100)
    used = _used_from(report, len(sources))
    all_sources = [{"n": i + 1, "title": s["title"], "url": s["url"], "domain": _domain(s["url"]),
                    "excerpt": s["text"][:400], "query": s.get("query", "")}
                   for i, s in enumerate(sources)]
    return {"report": report, "used": used, "sources": all_sources}


# ---------------------------------------------------------------- reflect (the "deep" step)
def reflect(question: str, report: str, breadth: int = 2) -> list:
    sys = ("You are a research critic. Given a QUESTION and a DRAFT briefing, list the most important "
           "aspects still unanswered or weakly supported. Return a JSON array of NEW, specific web-search "
           "queries to close those gaps. If the draft is already thorough, return an empty array [].")
    qs = llm.chat_json(sys, f"QUESTION: {question}\n\nDRAFT:\n{report[:2500]}\n\nGaps as JSON array of queries.")
    return [q.strip() for q in qs if isinstance(q, str) and q.strip()][:breadth] if isinstance(qs, list) else []


# ---------------------------------------------------------------- orchestration
def research(question: str, mode: str = "deep", area: str = "web", breadth: int = 4) -> dict:
    trace = []
    if mode == "quick" and len(question) <= 140:
        queries = [question]
    else:
        # long questions (e.g. a scoped 'dig deeper') make terrible single search queries —
        # plan concise ones with the key nouns preserved.
        queries = plan(question, breadth=(3 if mode == "quick" else breadth))
        trace.append(f"planned {len(queries)} sub-question(s)")

    sources = _gather_many(queries, area, trace)

    if mode in ("deep", "swarm") and sources:
        draft = synthesize(question, sources)
        gaps = reflect(question, draft["report"])
        if gaps:
            trace.append(f"gap-fill round: {len(gaps)} follow-up quer(ies)")
            more = _gather_many(gaps, area, trace)
            seen = {s["url"].split("#")[0].rstrip("/") for s in sources}
            sources += [s for s in more if s["url"].split("#")[0].rstrip("/") not in seen]

    trace.append(f"synthesizing from {len(sources)} source(s)")
    result = synthesize(question, sources)
    result["trace"] = trace
    result["backend"] = llm.backend_label()
    result["mode"] = mode
    try:
        if memory.save(question, result["report"], result.get("sources", []), area):
            trace.append("saved to vault (compounding memory)")
    except Exception:
        pass
    return result


def dig(focus: str, parent_question: str, instruction: str = "", area: str = "web", mode: str = "deep") -> dict:
    """Scoped follow-up research on a claim/selection/source — powers the UI's 'dig deeper' branches.
    `focus` = the selected text / source the user is drilling into; `instruction` = the user's own
    words for exactly what they want explored further (optional but this is the steering wheel)."""
    scoped = f'In the context of the research question: "{parent_question}"\n'
    if focus:
        scoped += f"Focus specifically on this part: {focus}\n"
    scoped += ("Investigate what the user asked for and cite sources: "
               + (instruction.strip() or f"go deeper on: {focus}"))
    res = research(scoped, mode=mode, area=area, breadth=3)
    res["focus"] = focus
    res["instruction"] = instruction
    return res


if __name__ == "__main__":
    import sys, argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("question", nargs="*")
    ap.add_argument("--mode", default="deep", choices=["quick", "deep", "swarm"])
    ap.add_argument("--area", default="web")
    ap.add_argument("--save", action="store_true", help="bank the brief to the Hermes vault")
    a = ap.parse_args()
    if a.save:
        memory.ENABLED = True
    q = " ".join(a.question) or "How are small language models changing enterprise AI in 2026?"
    print(f"Q: {q}\nmode={a.mode} area={a.area} backend={llm.backend_label()}\n")
    res = research(q, mode=a.mode, area=a.area)
    for t in res["trace"]:
        print(f"  · {t}")
    print("\n" + "=" * 70 + "\n" + res["report"] + "\n" + "=" * 70)
    print("\nSOURCES:")
    for s in res["sources"]:
        print(f"  [{s['n']}] {s['title'][:70]} — {s['url']}")