import sqlite3

con = sqlite3.connect("shop.db")
cur = con.cursor()

print("골든셋:", cur.execute("SELECT COUNT(*) FROM purchases WHERE is_holdout = 1").fetchone()[0], "건")
print("이력  :", cur.execute("SELECT COUNT(*) FROM purchases WHERE is_holdout = 0").fetchone()[0], "건")

# dict 로 감싸면 [(고객, 상품), ...] 이 {고객: 상품} 이 된다.
# 고객 번호로 바로 정답을 꺼낼 수 있어 채점할 때 편하다
answers = dict(cur.execute(
    "SELECT customer_id, product_id FROM purchases WHERE is_holdout = 1").fetchall())

print("고객 수:", len(answers))
print("C001 의 정답:", answers["C001"])

# 정답 상품이 몇 종류인가. 한 상품에 몰려 있으면 그것만 찍어도 점수가 나온다
print("정답 상품 종류:", cur.execute(
    "SELECT COUNT(DISTINCT product_id) FROM purchases WHERE is_holdout = 1").fetchone()[0])

duplicated = 0

for customer_id, answer in answers.items():
    bought = {r[0] for r in cur.execute(
        "SELECT product_id FROM purchases WHERE customer_id = ? AND is_holdout = 0",
        (customer_id,))}
    if answer in bought:
        duplicated += 1

print("이력에도 있는 정답:", duplicated, "건")

"""데이터베이스에 닿는 자리를 여기 하나로 모은다.

다른 파일은 전부 이렇게 쓴다 ─
    from db import query, one
"""

import sqlite3

DB_PATH = "shop.db"

con = sqlite3.connect(DB_PATH)


def query(sql, params=()):
    """여러 줄을 꺼낸다."""
    return con.execute(sql, params).fetchall()


def one(sql, params=()):
    """한 줄만 꺼낸다. 없으면 None 이 온다."""
    return con.execute(sql, params).fetchone()

import json
import time

import numpy as np

from config import client, MODEL
from db import query, one          # 이제 sqlite3 를 여기서 안 부른다

N = 30                  # 채점할 고객 수
N_CANDIDATES = 10       # 모델에게 줄 후보 개수
N_PICK = 5              # 그중 고르게 할 개수

# 저장해둔 벡터를 꺼낸다 (4·5편)
products = query("SELECT product_id, vector FROM product_vectors")
product_ids = [p for p, _ in products]
V = np.array([json.loads(v) for _, v in products], dtype="float32")

customers = query("SELECT customer_id, vector FROM customer_vectors")
customer_ids = [c for c, _ in customers]
CV = np.array([json.loads(v) for _, v in customers], dtype="float32")

# 답안지 (골든셋)
answers = dict(query("SELECT customer_id, product_id FROM purchases WHERE is_holdout = 1"))


def candidates_for(customer_id, index):
    """이미 산 것을 빼고, 뜻이 가까운 상품 열 개를 고른다. 5편에서 만든 것이다."""
    bought = {r[0] for r in query(
        "SELECT product_id FROM purchases WHERE customer_id = ? AND is_holdout = 0",
        (customer_id,))}

    scores = V @ CV[index]
    picked = [product_ids[i] for i in scores.argsort()[::-1]
              if product_ids[i] not in bought][:N_CANDIDATES]

    # 번호만으로는 모델에게 줄 수 없다. 이름·브랜드·가격을 붙여 온다.
    # 물음표를 개수만큼 만들어 IN 에 넣는다 — 1편에서 한 것과 같은 방법이다
    marks = ",".join("?" * len(picked))
    rows = query(
        f"SELECT product_id, name, brand, price FROM products WHERE product_id IN ({marks})",
        tuple(picked))

    # IN 으로 꺼내면 순서가 뒤섞여 온다. 가까운 순서대로 되돌려 놓는다
    order = {p: i for i, p in enumerate(picked)}
    return sorted(rows, key=lambda row: order[row[0]])

def prompt_for(customer_id, candidates):
    """9편의 08_safe_prompt.py 그대로다. 이름·전화·메일은 SELECT 에 아예 없고,
    후기에만 mask_text 를 건다."""
    age, gender, skin_type = one(
        "SELECT age, gender, skin_type FROM customers WHERE customer_id = ?", (customer_id,))

    history = query("""
        SELECT products.name, products.price, purchases.rating, purchases.review
        FROM purchases JOIN products ON purchases.product_id = products.product_id
        WHERE purchases.customer_id = ? AND purchases.is_holdout = 0
    """, (customer_id,))

    lines = "\n".join(f"- {name} {price}원 (별점 {rating}) {mask_text(review)}"
                      for name, price, rating, review in history)
    catalog = "\n".join(f"{i} | {n} | {b} | {p}원" for i, n, b, p in candidates)

    return f"""아래 고객에게 어울릴 상품 {N_PICK}개를 후보에서 골라라.

[고객] {age}세 {gender} · {skin_type}
[구매 이력]
{lines}

[후보]
{catalog}

좋은 순서대로 {N_PICK}개를 고른다."""


def ask_model(customer_id, candidates):
    """10편의 12_schema.py 를 함수로 묶었다.

    추천만 꺼내(res.choices[0].message.parsed) 돌려주면 깔끔해 보이는데,
    그러면 토큰 수(usage)가 같이 버려진다. 재려고 만드는 코드에서
    잴 것을 버리면 안 되므로 응답을 통째로 돌려준다.
    """
    return client.chat.completions.parse(
        model=MODEL,
        messages=[
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": prompt_for(customer_id, candidates)},
        ],
        response_format=Recommend,        # 10편에서 만든 pydantic 클래스
    )

hits = {1: 0, 3: 0, 5: 0}       # hit@1 · hit@3 · hit@5 를 한 번에 센다
in_candidates = 0                # 후보 안에 정답이 있었던 횟수
started_all = time.time()

log = open("runs.jsonl", "a", encoding="utf-8")     # a = 이어 붙이기

for index in range(N):
    customer_id = customer_ids[index]
    candidates = candidates_for(customer_id, index)
    candidate_ids = [c[0] for c in candidates]
    answer = answers[customer_id]

    # 정답이 후보에 아예 없으면 모델이 무엇을 하든 못 맞힌다. 그걸 따로 센다
    in_candidates += answer in candidate_ids

    started = time.time()
    res = ask_model(customer_id, candidates)      # 10편에서 만든 것. 응답을 그대로 돌려준다
    took = time.time() - started

    recommend = res.choices[0].message.parsed

    # 후보 밖 번호는 버린다 (10편). 남은 것이 진짜 추천이다
    picks = [it.product_id for it in recommend.items if it.product_id in candidate_ids]

    for k in hits:
        hits[k] += answer in picks[:k]

    # 관측 기록 — 한 줄이 한 번의 호출이다
    log.write(json.dumps({
        "customer_id": customer_id,
        "took": round(took, 2),
        "in_tokens": res.usage.prompt_tokens,
        "out_tokens": res.usage.completion_tokens,
        "picks": picks,
        "answer": answer,
        "hit": answer in picks[:5],
    }, ensure_ascii=False) + "\n")
    log.flush()          # 바로 파일에 쓴다. 중간에 멈춰도 여기까지는 남는다

    print(f"{index + 1:2d}/{N} {customer_id} {took:5.1f}초 "
          f"{'HIT' if answer in picks[:5] else ''}", flush=True)

log.close()
elapsed = time.time() - started_all

print(f"\n{N}명 {elapsed / 60:.1f}분 (한 명당 {elapsed / N:.1f}초)")
for k in (1, 3, 5):
    print(f"hit@{k} {hits[k] / N * 100:.1f}%")
print(f"후보 안에 정답이 있던 비율 {in_candidates / N * 100:.1f}%")