#!/usr/bin/env python3
"""V2: force deepseek-chat to exactly 100 words via NUMBERED sections.
Method frozen in acceptance-v2.md. The model writes every word with its
index in brackets; the utility strips, verifies count AND sequence 1..25,
and repairs surgically in the model's own numbering units."""
import json, os, re, subprocess

HERE = os.path.dirname(os.path.abspath(__file__))
CLI = os.path.expanduser("~/workspace/custom-connectors/custom.legacy-deepseek/bin/deepseek-api")
ENV = {**os.environ, "PYTHONPATH": "/opt/muse/skills/custom-connector-setup/bin"}
MODEL = "deepseek-chat"
MAX_ATTEMPTS = 5

FORMAT = ("Write the section in NUMBERED form: every word followed immediately by its "
          "index in square brackets, starting at [1] and increasing by exactly 1 per word. "
          "The last word must carry [25]. Example shape: 'Meet[1] your[2] new[3] morning[4] …'. "
          "Output ONLY the numbered section — no title, no notes.")

SECTIONS = [
    ("S1-hook", "the opening hook: introduce this coffee mug and the feeling it promises in the first sip"),
    ("S2-craft", "materials and craft: ceramic body, glaze, weight, handle design, how it is made and finished"),
    ("S3-experience", "the daily experience: morning coffee at home, desk at work, warmth in the hands, easy washing"),
    ("S4-close", "the close: who this mug is for and a gentle invitation to make it theirs"),
]

TOKEN = re.compile(r"([^\s\[\]]+)\[(\d+)\]")

def call(messages):
    payload = {"model": MODEL, "messages": messages}
    r = subprocess.run([CLI, "POST", "/chat/completions", "raw:" + json.dumps(payload)],
                       capture_output=True, text=True, env=ENV, timeout=120)
    d = json.loads(r.stdout)
    return d["choices"][0]["message"]["content"].strip()

def wc(text):
    return len([w for w in re.split(r"\s+", text.strip()) if w])

def parse_numbered(text):
    """Return (words, seq_ok, last_index, bare_words)."""
    toks = TOKEN.findall(text)
    if not toks:
        return [], False, 0, wc(text)
    words = [w for w, _ in toks]
    idx = [int(n) for _, n in toks]
    seq_ok = idx == list(range(1, len(idx) + 1))
    return words, seq_ok, (idx[-1] if idx else 0), len(words)

def run_section(name, brief, assembly_no):
    attempts = []
    prompt = (f"You are writing one section of a coffee-mug product description. "
              f"This section's job: {brief}. It must contain EXACTLY 25 words. {FORMAT}")
    msgs = [{"role": "user", "content": prompt}]
    for a in range(1, MAX_ATTEMPTS + 1):
        ans = call(msgs)
        words, seq_ok, last_idx, n = parse_numbered(ans)
        good = (n == 25 and seq_ok and last_idx == 25)
        attempts.append({"attempt": a, "tokens": n, "last_index": last_idx,
                         "seq_ok": seq_ok, "pass": good, "text": ans})
        print(f"  A{assembly_no} {name} attempt {a}: tokens={n} last={last_idx} seq_ok={seq_ok}", flush=True)
        if good:
            return " ".join(words), attempts
        if n == 0:
            fix = f"That had no numbered words at all. Follow the format exactly: {FORMAT} Job: {brief}."
        elif not seq_ok:
            fix = (f"Your numbering broke sequence (tokens found: {n}, last index [{last_idx}]). "
                   f"Rewrite the section with clean numbering [1] to [25], one index per word, no gaps or repeats. Job: {brief}.")
        elif n < 25:
            fix = (f"Your numbering stopped at [{last_idx}] — that is {25 - last_idx} word(s) SHORT. "
                   f"Rewrite the section keeping your words and APPEND new words numbered "
                   f"[{last_idx + 1}] through [25] to complete it. Same job: {brief}.")
        else:
            fix = (f"Your numbering reached [{last_idx}] — that is {last_idx - 25} word(s) TOO MANY. "
                   f"Rewrite the section deleting the extra word(s) so the last word carries [25]. Same job: {brief}.")
        msgs.append({"role": "assistant", "content": ans})
        msgs.append({"role": "user", "content": fix})
    return None, attempts

def run_assembly(no):
    print(f"ASSEMBLY {no}", flush=True)
    parts, record, ok = [], [], True
    for name, brief in SECTIONS:
        text, attempts = run_section(name, brief, no)
        record.append({"section": name, "attempts": attempts, "passed": text is not None})
        if text is None:
            ok = False
            break
        parts.append(text)
    full = " ".join(parts) if ok else None
    total = wc(full) if full else None
    print(f"ASSEMBLY {no}: {'PASS total=' + str(total) if ok and total == 100 else 'FAIL total=' + str(total)}", flush=True)
    return {"assembly": no, "sections": record, "text": full, "total_words": total,
            "pass": bool(ok and total == 100)}

results = [run_assembly(i) for i in (1, 2, 3)]
json.dump(results, open(os.path.join(HERE, "force100v2.json"), "w"), indent=1)
print("FORCE100-V2 DONE:", sum(1 for r in results if r["pass"]), "/3 assemblies pass")