import argparse
import json
import re
import sys
import time
from pathlib import Path

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

prompt = (p.get("prompt") or "").strip()
if not prompt:
    print("ERROR: prompt is required")
    sys.exit(1)

from mlx_lm import load, generate  # noqa: E402
from mlx_lm.sample_utils import make_sampler  # noqa: E402

MODEL = "mlx-community/Qwen3-30B-A3B-4bit"
t0 = time.time()
model, tokenizer = load(MODEL)
print(f"model loaded in {time.time() - t0:.1f}s", flush=True)

messages = []
if (p.get("system") or "").strip():
    messages.append({"role": "system", "content": p["system"].strip()})
messages.append({"role": "user", "content": prompt})
chat = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, tokenize=False, enable_thinking=False)

t1 = time.time()
text = generate(model, tokenizer, prompt=chat,
                max_tokens=int(p.get("max_tokens", 1024)),
                sampler=make_sampler(temp=float(p.get("temperature", 0.7))))
text = re.sub(r"<think>.*?</think>", "", text, flags=re.S).strip()

out = Path(a.outdir) / "llm_output.txt"
out.write_text(text + "\n")
print(f"done in {time.time() - t1:.1f}s, {len(text)} chars:\n{text[:400]}", flush=True)