#!/usr/bin/env python3
"""Convert a CSV of examples into chat-format JSONL for fine-tuning.

Each output line looks like:
    {"messages": [{"role": "system", "content": "..."},
                  {"role": "user", "content": "..."},
                  {"role": "assistant", "content": "..."}]}
This is the chat fine-tuning format documented by OpenAI and accepted by many
other training tools. Check your provider's docs for any extra requirements.

Usage:
    python3 scripts/csv_to_jsonl.py INPUT.csv OUTPUT.jsonl --user-col COL --assistant-col COL
                                    [--system TEXT | --system-col COL]
                                    [--delimiter ,] [--encoding utf-8-sig]
                                    [--skip-invalid] [--dedupe] [--no-strip]

    --user-col       Column holding the user message (required).
    --assistant-col  Column holding the ideal assistant reply (required).
    --system         Same system message for every example.
    --system-col     Column holding a per-row system message (empty cells are allowed
                     and simply omit the system message for that row).
    --skip-invalid   Skip rows with an empty user or assistant cell instead of failing.
    --dedupe         Drop rows whose (system, user, assistant) repeat an earlier row.
    --no-strip       Keep leading/trailing whitespace in cells (stripped by default).

The output file is written only when the run succeeds, so a failed run never
leaves a half-written dataset behind.

Requirements: Python 3.8+, no packages, no environment variables.

Exit codes:
    0  output written
    1  invalid rows found (and --skip-invalid not given), or no valid rows
    2  usage error (file not found, column missing, bad arguments)
"""
import argparse
import csv
import json
import os
import sys
import tempfile


def main():
    parser = argparse.ArgumentParser(description="CSV to chat-format JSONL converter.")
    parser.add_argument("input")
    parser.add_argument("output")
    parser.add_argument("--user-col", required=True)
    parser.add_argument("--assistant-col", required=True)
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--system")
    group.add_argument("--system-col")
    parser.add_argument("--delimiter", default=",")
    parser.add_argument("--encoding", default="utf-8-sig")
    parser.add_argument("--skip-invalid", action="store_true")
    parser.add_argument("--dedupe", action="store_true")
    parser.add_argument("--no-strip", action="store_true")
    args = parser.parse_args()

    clean = (lambda v: v or "") if args.no_strip else (lambda v: (v or "").strip())
    csv.field_size_limit(10 * 1024 * 1024)
    try:
        handle = open(args.input, newline="", encoding=args.encoding)
    except (OSError, LookupError) as exc:
        print("error: cannot open %s: %s" % (args.input, exc), file=sys.stderr)
        return 2

    written, skipped, invalid, dupes = 0, 0, [], 0
    seen = set()
    out_dir = os.path.dirname(os.path.abspath(args.output))
    tmp = tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=out_dir, delete=False, suffix=".tmp")
    try:
        with handle, tmp:
            reader = csv.DictReader(handle, delimiter=args.delimiter)
            columns = reader.fieldnames or []
            needed = [args.user_col, args.assistant_col] + ([args.system_col] if args.system_col else [])
            missing = [c for c in needed if c not in columns]
            if missing:
                print("error: column(s) not found: %s. Available: %s" % (", ".join(missing), ", ".join(columns)), file=sys.stderr)
                return 2
            for row_number, row in enumerate(reader, start=2):
                user = clean(row.get(args.user_col))
                assistant = clean(row.get(args.assistant_col))
                system = clean(row.get(args.system_col)) if args.system_col else clean(args.system)
                if not user or not assistant:
                    which = "user" if not user else "assistant"
                    invalid.append("row %d: empty %s cell" % (row_number, which))
                    continue
                key = (system, user, assistant)
                if args.dedupe and key in seen:
                    dupes += 1
                    continue
                seen.add(key)
                messages = []
                if system:
                    messages.append({"role": "system", "content": system})
                messages.append({"role": "user", "content": user})
                messages.append({"role": "assistant", "content": assistant})
                tmp.write(json.dumps({"messages": messages}, ensure_ascii=False) + "\n")
                written += 1
        for message in invalid[:20]:
            print("warning: %s" % message, file=sys.stderr)
        if len(invalid) > 20:
            print("warning: ... and %d more invalid rows" % (len(invalid) - 20), file=sys.stderr)
        if invalid and not args.skip_invalid:
            print("error: %d invalid row(s); fix them or pass --skip-invalid. Nothing written." % len(invalid), file=sys.stderr)
            return 1
        skipped = len(invalid)
        if written == 0:
            print("error: no valid rows; nothing written", file=sys.stderr)
            return 1
        os.replace(tmp.name, args.output)
    except csv.Error as exc:
        print("error: CSV parse problem: %s" % exc, file=sys.stderr)
        return 2
    finally:
        if os.path.exists(tmp.name):
            os.unlink(tmp.name)

    print("wrote %d example(s) to %s (skipped %d invalid, %d duplicate)" % (written, args.output, skipped, dupes))
    return 0


if __name__ == "__main__":
    sys.exit(main())