#!/usr/bin/env python3
"""
cluster_then_polish.py — two-pass LLM pattern wrapper.

Wraps cluster.ask() with explicit "this is a draft, mark for human/Claude polish"
metadata. Goal: cluster does the heavy lifting (parallel, cheap, always-on),
Claude or human review handles the final polish.

Why: open-source 7B models can produce structured / classified / extracted output
reliably, but produce generic narrative text without grounding. Letting them
"finish" content sent to clients = quality loss. This wrapper bakes in the
two-pass expectation: cluster draft → human/Claude approval before any external use.

Public API:
    draft(prompt, task='general', node_pref=None, **kw) -> {ok, text, model, node, polish_required: True, draft_id, hint}
    save_draft(draft, path) -> writes md with frontmatter
    mark_polished(draft_id) -> updates draft registry

Usage example:
    from cluster_then_polish import draft, save_draft
    d = draft("Schrijf cold-WhatsApp voor lunchroom in Tilburg", task='chat', node_pref='pc')
    save_draft(d, '~/.config/atlas/drafts/lunchroom_outreach.md')
    # User/Claude reviews + edits before sending

Registry: drafts logged to ~/.config/atlas/drafts/_registry.json with
status (draft/polished/sent), creation time, node/model used, prompt hash.
"""
from __future__ import annotations
import sys, json, hashlib, time
from pathlib import Path
from datetime import datetime

HOME = Path.home()
ATLAS = HOME / '.config' / 'atlas'
DRAFTS_DIR = ATLAS / 'drafts'
REGISTRY = DRAFTS_DIR / '_registry.json'

sys.path.insert(0, str(ATLAS))
try:
    import cluster
except ImportError:
    raise RuntimeError("cluster.py not importable; check ~/.config/atlas/cluster.py")


# Task types that should ALWAYS go through polish (narrative, klant-facing)
POLISH_REQUIRED_TASKS = {'chat', 'reason', 'general'}

# Task types where cluster output is directly usable (structured, classified)
POLISH_OPTIONAL_TASKS = {'fast', 'embed', 'code', 'vision'}


def _hash_prompt(prompt: str) -> str:
    return hashlib.sha1(prompt.encode('utf-8')).hexdigest()[:10]


def _load_registry() -> dict:
    DRAFTS_DIR.mkdir(parents=True, exist_ok=True)
    if REGISTRY.exists():
        try:
            return json.loads(REGISTRY.read_text(encoding='utf-8'))
        except Exception:
            pass
    return {"drafts": {}}


def _save_registry(reg: dict) -> None:
    REGISTRY.write_text(json.dumps(reg, indent=2), encoding='utf-8')


def draft(prompt: str, task: str = 'general', node_pref: str | None = None,
          max_tokens: int = 800, temperature: float = 0.5, timeout: int = 180,
          purpose: str = '') -> dict:
    """
    Call cluster.ask() but tag the result as draft requiring polish.

    :param purpose: short label (e.g. 'cold_outreach', 'morning_brief', 'mail_triage')
                    used for the draft filename and routing decisions
    """
    t0 = time.time()
    r = cluster.ask(prompt, task=task, node_pref=node_pref,
                    max_tokens=max_tokens, temperature=temperature, timeout=timeout)
    elapsed = round(time.time() - t0, 1)

    if not r.get('ok'):
        return {
            "ok": False,
            "err": r.get('err', 'cluster error'),
            "polish_required": True,
            "hint": "cluster call failed; do not use any partial output"
        }

    draft_id = f"{int(time.time())}-{_hash_prompt(prompt)}"
    polish_required = task in POLISH_REQUIRED_TASKS

    result = {
        "ok": True,
        "draft_id": draft_id,
        "text": r['text'],
        "model": r.get('model', '?'),
        "node": r.get('node', '?'),
        "elapsed_s": elapsed,
        "task": task,
        "purpose": purpose,
        "polish_required": polish_required,
        "polish_status": "pending" if polish_required else "not_required",
        "created_at": datetime.now().isoformat(timespec='seconds'),
        "hint": (
            "POLISH REQUIRED before external use. Cluster 7B models produce "
            "generic narrative; review for: specificity, sector-fit, tone, "
            "factual claims. Then call mark_polished(draft_id)."
            if polish_required else
            "Structured output, safe to use as-is for triage/classify. "
            "Manual spot-check still recommended."
        )
    }

    # Register
    reg = _load_registry()
    reg['drafts'][draft_id] = {
        "purpose": purpose,
        "task": task,
        "node": result['node'],
        "model": result['model'],
        "created_at": result['created_at'],
        "polish_status": result['polish_status'],
        "prompt_hash": _hash_prompt(prompt),
        "char_count": len(r['text']),
    }
    _save_registry(reg)

    return result


def save_draft(draft_result: dict, path: str | Path) -> Path:
    """Save draft to .md with frontmatter so reviewer knows status."""
    p = Path(path).expanduser()
    p.parent.mkdir(parents=True, exist_ok=True)

    fm = (
        f"---\n"
        f"draft_id: {draft_result.get('draft_id', '?')}\n"
        f"polish_status: {draft_result.get('polish_status', '?')}\n"
        f"node: {draft_result.get('node', '?')}\n"
        f"model: {draft_result.get('model', '?')}\n"
        f"task: {draft_result.get('task', '?')}\n"
        f"purpose: {draft_result.get('purpose', '?')}\n"
        f"created_at: {draft_result.get('created_at', '?')}\n"
        f"polish_required: {draft_result.get('polish_required', False)}\n"
        f"---\n\n"
        f"> **HINT**: {draft_result.get('hint', '')}\n\n"
    )
    p.write_text(fm + draft_result.get('text', ''), encoding='utf-8')
    return p


def mark_polished(draft_id: str, polished_by: str = 'human') -> bool:
    """Update registry to indicate draft has been reviewed/polished."""
    reg = _load_registry()
    if draft_id in reg.get('drafts', {}):
        reg['drafts'][draft_id]['polish_status'] = 'polished'
        reg['drafts'][draft_id]['polished_by'] = polished_by
        reg['drafts'][draft_id]['polished_at'] = datetime.now().isoformat(timespec='seconds')
        _save_registry(reg)
        return True
    return False


def list_pending(purpose_filter: str = '') -> list:
    """Return list of drafts still awaiting polish (polish_status='pending')."""
    reg = _load_registry()
    out = []
    for did, info in reg.get('drafts', {}).items():
        if info.get('polish_status') != 'pending':
            continue
        if purpose_filter and purpose_filter not in info.get('purpose', ''):
            continue
        out.append({"draft_id": did, **info})
    return sorted(out, key=lambda x: x.get('created_at', ''), reverse=True)


if __name__ == '__main__':
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument('prompt', nargs='?', help='Prompt text')
    ap.add_argument('-task', default='general')
    ap.add_argument('-node', default=None)
    ap.add_argument('-purpose', default='ad_hoc')
    ap.add_argument('-out', default=None, help='Save draft to this path')
    ap.add_argument('--list-pending', action='store_true')
    args = ap.parse_args()

    if args.list_pending:
        for d in list_pending():
            print(f"{d['draft_id']}  [{d['purpose']}]  {d['model']}@{d['node']}  {d['char_count']}ch  {d['created_at']}")
        sys.exit(0)

    if not args.prompt:
        print("usage: cluster_then_polish.py 'prompt' [-task T] [-node N] [-purpose P] [-out FILE]")
        sys.exit(1)

    d = draft(args.prompt, task=args.task, node_pref=args.node, purpose=args.purpose)
    if not d['ok']:
        print(f"ERROR: {d.get('err')}")
        sys.exit(1)

    print(f"[{d['node']}/{d['model']} {d['elapsed_s']}s] draft_id={d['draft_id']}")
    print(f"polish_status: {d['polish_status']}")
    print(f"--- text ({len(d['text'])} chars) ---")
    print(d['text'])
    print()

    if args.out:
        p = save_draft(d, args.out)
        print(f"saved -> {p}")