"""Rewrite tweets with a chat LLM to get a few more training rows per label,
keeping the same dialect and stance. Point it at your endpoint with
AUG_BASE_URL / AUG_MODEL.

    python -m src.augment --in data/track2/train.csv \\
        --out data/track2/train_aug.csv --n 2
"""
import argparse
import json
import os
import urllib.request
from concurrent.futures import ThreadPoolExecutor

import pandas as pd

PROMPT = (
    "أعد صياغة التغريدة التالية بأسلوب مختلف مع الحفاظ التام على نفس "
    "الموقف والرأي تجاه الموضوع، ونفس اللهجة، ونفس المعنى. لا تضف أي "
    "تعليق أو شرح، وأخرج التغريدة المعاد صياغتها فقط:\n\n{text}"
)


def call(base_url, model, text, n, timeout=120):
    body = json.dumps({
        "model": model,
        "n": n,
        "temperature": 1.0,
        "top_p": 0.95,
        "max_tokens": 160,
        "messages": [{"role": "user", "content": PROMPT.format(text=text)}],
    }).encode("utf-8")
    req = urllib.request.Request(
        base_url.rstrip("/") + "/chat/completions",
        data=body,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        data = json.load(r)
    return [c["message"]["content"].strip() for c in data["choices"]]


def clean(variant, original):
    v = variant.strip().strip('"').strip()
    if "\n" in v:
        v = v.split("\n")[-1].strip()
    if len(v) < 5 or v == original.strip():
        return None
    return v


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--in", dest="inp", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--n", type=int, default=2)
    ap.add_argument("--concurrency", type=int, default=32)
    ap.add_argument("--max_rows", type=int, default=0)
    ap.add_argument(
        "--base_url", default=os.environ.get("AUG_BASE_URL", "")
    )
    ap.add_argument("--model", default=os.environ.get("AUG_MODEL", ""))
    args = ap.parse_args()
    if not args.base_url or not args.model:
        raise SystemExit("set --base_url/--model or AUG_BASE_URL/AUG_MODEL")

    df = pd.read_csv(args.inp, keep_default_na=False, encoding="utf-8-sig")
    if args.max_rows:
        df = df.head(args.max_rows)

    def work(row):
        try:
            variants = call(args.base_url, args.model, row["text"], args.n)
        except Exception:
            return []
        out = []
        for v in variants:
            c = clean(v, row["text"])
            if c:
                out.append({
                    "text": c,
                    "target": row["target"],
                    "stance": row["stance"],
                    "source": "paraphrase",
                })
        return out

    rows = df.to_dict("records")
    generated = []
    done = 0
    with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
        for res in ex.map(work, rows):
            generated.extend(res)
            done += 1
            if done % 200 == 0:
                print(f"  {done}/{len(rows)} rows, "
                      f"{len(generated)} paraphrases")

    orig = df.copy()
    orig["source"] = "original"
    keep = ["text", "target", "stance", "source"]
    combined = pd.concat(
        [orig[keep], pd.DataFrame(generated)[keep]], ignore_index=True
    )
    combined.to_csv(args.out, index=False, encoding="utf-8-sig")
    print(f"[write] {len(orig)} original + {len(generated)} paraphrases "
          f"= {len(combined)} rows -> {args.out}")


if __name__ == "__main__":
    main()