import argparse
import time

import jev
from generate import generate
from rerank import rerank
from decompose import decompose
from retrieve import retrieve_many
from route import route
from verify import verify

CANNED = {
    "needs_live_web": "That needs live web data, and I only search my own notes.",
    "ambiguous_ask_back": "Could you say a bit more about what you want to know?",
}


def timed(record, stage, fn, *args):
    start, calls = time.perf_counter(), len(jev.CALLS)
    result = fn(*args)
    new = jev.CALLS[calls:]
    record["stages"][stage] = {"seconds": time.perf_counter() - start,
                               "jev_tokens": sum(c["input_tokens"] for c in new),
                               "jev_dollars": sum(c["dollars"] for c in new)}
    return result


def run(query, use_jev=True, split=False):
    record = {"query": query, "stages": {}, "route": None, "candidates": [], "chunks": [], "verify": None,
              "subqueries": [query]}
    started = time.perf_counter()
    retrieve_needed = True
    if use_jev:
        record["route"] = timed(record, "route", route, query)
        r = record["route"]
        retrieve_needed = r["intent"] == "answer_from_my_notes" or r["confidence"] < 0.5
    if retrieve_needed:
        if split:
            record["subqueries"] = timed(record, "decompose", decompose, query)
        record["candidates"] = timed(record, "retrieve", retrieve_many, record["subqueries"])
        if use_jev:
            ranked = timed(record, "rerank", rerank, query, record["candidates"], None)
            record["reranked_ids"] = [c["id"] for c in ranked]
            record["chunks"] = ranked[:5]
        else:
            record["chunks"] = record["candidates"][:5]
    if retrieve_needed or record["route"]["intent"] == "small_talk":
        record["answer"], record["llm_usage"] = timed(record, "generate", generate, query, record["chunks"])
    else:
        record["answer"] = CANNED[record["route"]["intent"]]
    if use_jev and record["chunks"]:
        record["verify"] = timed(record, "verify", verify, record["answer"], record["chunks"])
    record["seconds"] = time.perf_counter() - started
    record["jev_tokens"] = sum(s["jev_tokens"] for s in record["stages"].values())
    record["jev_dollars"] = sum(s["jev_dollars"] for s in record["stages"].values())
    return record


def show(record):
    for stage, s in record["stages"].items():
        cost = f"  {s['jev_tokens']:>6} tok  ${s['jev_dollars']:.6f}" if s["jev_tokens"] else ""
        print(f"{stage:<9}{s['seconds'] * 1000:>8.0f} ms{cost}")
    if record["route"]:
        r = record["route"]
        print(f"\nroute: {r['intent']} (confidence {r['confidence']:.2f}), needs_retrieval {r['needs_retrieval']:.2f}")
    if len(record["subqueries"]) > 1:
        print("subqueries: " + " | ".join(record["subqueries"]))
    for i, c in enumerate(record["chunks"]):
        score = f"jev {c['relevance']:.2f}" if "relevance" in c else f"cos {c['similarity']:.3f}"
        print(f"  [c{i + 1}] {score}  {c['title'][:50]} / {c['section'][:40]}  {c['url']}")
    print(f"\n{record['answer']}\n")
    v = record["verify"]
    if v:
        status = "ESCALATE" if v["escalate"] else "supported"
        print(f"verify: {status}, {v['checked']} cited sentences checked")
        for claim in v["claims"]:
            print(f"  {claim['support']:.2f}  {claim['sentence'][:100]}")
    print(f"\ntotal {record['seconds'] * 1000:.0f} ms, Jev {record['jev_tokens']} tokens, ${record['jev_dollars']:.6f}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("question")
    parser.add_argument("--no-jev", action="store_true")
    parser.add_argument("--decompose", action="store_true")
    args = parser.parse_args()
    show(run(args.question, use_jev=not args.no_jev and jev.has_key(), split=args.decompose))