"""Second annotator for the match-validation sample: Claude (claude-opus-5-5) labels each candidate pair as
SAME (same product, same package size), SIZE_DIFF (same product, different size), DIFFERENT, or UNSURE.
Human adjudication (`label_human`) remains the reference; LLM labels are an auxiliary signal. Logged: model, tokens, cost.
"""
import json, os, time, pandas as pd
import anthropic
from dotenv import load_dotenv
from src.common import ROOT, path, get_logger, ledger

log = get_logger("llm_annotate")
MODEL = "claude-opus-5-5"; PRICE_IN, PRICE_OUT = 4.0, 20.0  # $/MTok

SYSTEM = """You are an expert grocery product data annotator for Quebec (French/English product names).
For each candidate pair of listings from two different retailers, decide whether they are the SAME product
(identical brand, product, variant/flavour/fat content, AND package size/quantity). Labels:
SAME = same product and same size; SIZE_DIFF = same product but different package size or count;
DIFFERENT = different product, brand, variant, flavour or fat content; UNSURE = cannot tell from the text.
Be strict: 'Lait 2%' vs 'Lait 3.25%' is DIFFERENT; 'Yogourt vanille' vs 'Yogourt fraise' is DIFFERENT.
Missing brand on one side is acceptable if the name clearly contains the brand.
Return ONLY a JSON array of objects {"sample_id": int, "label": str, "reason": str (<= 12 words)}."""

def main():
    load_dotenv(ROOT / ".env")
    f = path("processed") / "match_validation_sample.csv"
    s = pd.read_csv(f)
    client = anthropic.Anthropic(max_retries=3, timeout=300)
    tot_in = tot_out = 0; labels = {}
    for i in range(0, len(s), 20):
        chunk = s.iloc[i:i + 20]
        lines = []
        for _, r in chunk.iterrows():
            lines.append(f'{{"sample_id": {r.sample_id}, "A": "{r.source_a_txt}: brand={r.brand_a_txt!s} | {r.name_a_txt} | size={r.size_label_a_txt!s}", '
                         f'"B": "{r.source_b_txt}: brand={r.brand_b_txt!s} | {r.name_b_txt} | size={r.size_label_b_txt!s}"}}')
        msg = "Candidate pairs:\n" + "\n".join(lines)
        with client.messages.stream(model=MODEL, max_tokens=8000, system=SYSTEM, messages=[{"role": "user", "content": msg}]) as st:
            resp = st.get_final_message()
        if resp.stop_reason == "refusal":
            log.warning("refusal on chunk %d", i); continue
        text = "".join(b.text for b in resp.content if b.type == "text").strip()
        text = text[text.find("["): text.rfind("]") + 1]
        for o in json.loads(text):
            labels[int(o["sample_id"])] = (o["label"], o.get("reason", ""))
        tot_in += resp.usage.input_tokens; tot_out += resp.usage.output_tokens
        time.sleep(0.5)
    s["label_llm"] = s.sample_id.map(lambda k: labels.get(k, ("", ""))[0])
    s["llm_reason"] = s.sample_id.map(lambda k: labels.get(k, ("", ""))[1])
    s.to_csv(f, index=False)
    cost = tot_in / 1e6 * PRICE_IN + tot_out / 1e6 * PRICE_OUT
    log.info("annotated %d/%d pairs with %s ; tokens in=%d out=%d ; cost=$%.3f", s.label_llm.ne("").sum(), len(s), MODEL, tot_in, tot_out, cost)
    log.info("LLM labels by tier:\n%s", pd.crosstab(s.tier, s.label_llm).to_string())
    ledger("1", "llm_annotate", f"{MODEL}: {len(s)} pairs, in={tot_in}, out={tot_out}, cost=${cost:.3f}")

if __name__ == "__main__":
    main()