# (24_ask.py 의 위쪽을 그대로 복사해 온다 — data, model, search)
import json
import time
import urllib.request

SOFT = """너는 사내 문서 안내원이다. 주어진 자료만 근거로 답한다.
자료에 없는 내용은 절대 쓰지 않는다. 자료로 답할 수 없으면 「자료에 없다」고만 답한다.
반드시 한국어로만 답한다. 세 문장 이내로 답한다."""


def timed(prompt, num_gpu=None):
    """네이티브 창구로 부른다. 걸린 시간을 잘게 돌려주기 때문이다."""
    options = {"num_predict": 120}
    if num_gpu is not None:
        options["num_gpu"] = num_gpu        # 0 이면 CPU 로만 돌린다

    body = json.dumps({
        "model": "qwen2.5:3b",
        "messages": [
            {"role": "system", "content": SOFT},
            {"role": "user", "content": prompt},
        ],
        "stream": False,
        "options": options,
    }).encode()

    req = urllib.request.Request(
        "http://localhost:11434/api/chat", data=body,
        headers={"Content-Type": "application/json"},
    )
    r = json.loads(urllib.request.urlopen(req, timeout=900).read().decode())

    # 나노초로 온다. 1e9 로 나누면 초가 된다
    return {
        "in_tok": r["prompt_eval_count"],
        "prefill": r["prompt_eval_duration"] / 1e9,
        "total": r["total_duration"] / 1e9,
    }


question = "조각을 겹치게 자르는 이유가 뭔가"

for k in [1, 3, 5, 10, 20, 30]:
    context = "\n\n".join(text for score, slug, text in search(question, k))
    prompt = f"[자료]\n{context}\n\n[질문]\n{question}\n"

    g = timed(prompt)               # 되는 대로 (우리 기준 GPU)
    c = timed(prompt, num_gpu=0)    # CPU 강제 (학원 노트북 흉내)

    print(f"k={k:<3} {len(context):>5}자 {g['in_tok']:>5}tok | "
          f"GPU {g['prefill']:>4.1f}/{g['total']:>4.1f} | "
          f"CPU {c['prefill']:>4.1f}/{c['total']:>4.1f}")

# (28_grade.py 의 CASES · judge · grade 를 그대로 가져온다)

for k in [1, 3, 5, 10, 20, 30]:
    total = 0
    for question, kind, good, bad in CASES:
        context = "\n\n".join(t for s, sl, t in search(question, k))   # 여기만 바뀐다
        prompt = f"[자료]\n{context}\n\n[질문]\n{question}\n"

        ok = 0
        for _ in range(N):
            res = client.chat.completions.create(
                model="qwen2.5:3b",
                messages=[
                    {"role": "system", "content": SOFT},
                    {"role": "user", "content": prompt},
                ],
            )
            if judge(res.choices[0].message.content, good, bad):
                ok += 1
        total += ok
    print(f"k={k:<3} {total}/18")

question = "우리가 자른 조각의 크기는 몇 자인가"

for k in [3, 10, 30]:
    hits = search(question, k)
    inside = any("size=200" in text for score, slug, text in hits)
    context = "\n\n".join(text for score, slug, text in hits)

    print(f"--- k={k}  정답 조각 들어있나: {inside}")
    for _ in range(3):
        answer, sources = ask_with(context, question)
        print("   >", answer[:60])