#!/usr/bin/env python3
"""
compress.py — Compress English reasoning text into the telegraphic CJK register
using tessera-compressor behind any OpenAI-compatible endpoint (vLLM, llama.cpp
server, etc.). No API keys or HF token required; the endpoint is yours.

This is the same harness the compressor was accepted under: segment the block,
group sentences into step-sized passages, classify each passage, compress it
against the chain built so far, then run the deterministic fidelity gate. A
passage that fails the gate falls back to a rules-only compression, so a bad
model output costs savings, never content.

Serve the model first, e.g.:
    vllm serve ZelligeAI/tessera-compressor --port 8001
or with the GGUF:
    llama-server -m gguf/compressor-v31-q8_0.gguf --port 8001   # from the repo root

Then:
    # one block from a text file
    python compress.py --in think.txt --endpoint http://localhost:8001/v1

    # a JSONL corpus: {"id": ..., "text": ...} per line
    python compress.py --in blocks.jsonl --out compressed.jsonl \
        --endpoint http://localhost:8001/v1

Token counting: the fidelity gate compares token counts under a target
tokenizer. For results matching the accepted harness, point --tokenizer at the
model you are producing training data FOR (default: the compressor's own
tokenizer, which is close but not identical to the Qwen3.5 target used in the
acceptance run).
"""
import argparse
import json
import sys

from openai import OpenAI
from tokenizers import Tokenizer

from segmenting import segment, group_steps, classify_passage, facts, gate
from tokenmax import _apply_subs

PASSAGE_SYSTEM = (
    "你是推理压缩器。Re-notate the NEXT PASSAGE of a reasoning chain into telegraphic "
    "CJK/symbol notation. Every NEW logical step, fact, number and identifier must "
    "survive — unless already stated in the chain. Never restate chain content. "
    "[passage=load]: step-lossless telegraphic. [passage=narr]: minimal stubs "
    "(试X→否). Output only the re-notated continuation."
)

MAX_NEW_TOKENS = 512


def compress_block(text, client, model, ntok):
    """Compress one reasoning block. Returns (compressed_text, stats)."""
    segs = group_steps(segment(text))
    chain, seen = [], set()
    stats = {"segments": len(segs), "model_ok": 0, "fallback": 0,
             "narr_skipped": 0, "code": 0, "calls": 0}

    for kind, s in segs:
        if kind == "code":
            chain.append(s)
            seen |= facts(s)
            stats["code"] += 1
            continue
        cls = classify_passage(s, seen, ntok)
        novel = facts(s) - seen
        rules_s, _ = _apply_subs(s)
        if not rules_s.strip():
            continue
        tail = "\n".join(chain)[-500:] or "(start)"
        stats["calls"] += 1
        r = client.chat.completions.create(
            model=model, temperature=0.0, max_tokens=MAX_NEW_TOKENS,
            messages=[
                {"role": "system", "content": PASSAGE_SYSTEM},
                {"role": "user", "content": f"[passage={cls}]\n链:\n{tail}\n\n段:\n{s[:2000]}"},
            ],
            extra_body={"repetition_penalty": 1.15},
        )
        out = (r.choices[0].message.content or "").strip()

        if out == "∅" and cls == "narr" and not novel:
            stats["narr_skipped"] += 1
            seen |= facts(s)
            continue
        if gate(s, rules_s, out, ntok, novel=novel) is None:
            chain.append(out)
            stats["model_ok"] += 1
        else:
            chain.append(rules_s)
            stats["fallback"] += 1
        seen |= facts(s)

    return "\n".join(chain), stats


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--in", dest="inp", required=True,
                    help=".txt (one block) or .jsonl ({'id','text'} per line)")
    ap.add_argument("--out", default=None, help="output JSONL (default: stdout)")
    ap.add_argument("--endpoint", default="http://localhost:8001/v1")
    ap.add_argument("--model", default="ZelligeAI/tessera-compressor",
                    help="served model name at the endpoint")
    ap.add_argument("--tokenizer", default="ZelligeAI/tessera-compressor",
                    help="HF repo id or local tokenizer.json for gate token counts")
    args = ap.parse_args()

    if args.tokenizer.endswith(".json"):
        tok = Tokenizer.from_file(args.tokenizer)
    else:
        tok = Tokenizer.from_pretrained(args.tokenizer)

    def ntok(s):
        return len(tok.encode(s).ids) if s else 0

    client = OpenAI(base_url=args.endpoint, api_key="none")

    if args.inp.endswith(".jsonl"):
        rows = [json.loads(l) for l in open(args.inp) if l.strip()]
    else:
        rows = [{"id": args.inp, "text": open(args.inp).read()}]

    sink = open(args.out, "w") if args.out else sys.stdout
    for row in rows:
        compressed, stats = compress_block(row["text"], client, args.model, ntok)
        rec = {"id": row.get("id"), "compressed": compressed,
               "src_tokens": ntok(row["text"]), "out_tokens": ntok(compressed),
               "harness": stats}
        sink.write(json.dumps(rec, ensure_ascii=False) + "\n")
        sink.flush()
        print(f"[{row.get('id')}] {rec['src_tokens']} -> {rec['out_tokens']} tokens "
              f"(model_ok={stats['model_ok']} fallback={stats['fallback']})",
              file=sys.stderr)
    if args.out:
        sink.close()


if __name__ == "__main__":
    main()