#!/usr/bin/env python3
"""ZabaanAI-v2 SFT Data Curation Pipeline

Downloads instruction datasets from HF Hub and converts to ChatML messages format
for Qwen2.5-7B-Instruct SFT training.

Supported formats -> ChatML conversion:
  - Alpaca (instruction/input/output) -> messages
  - Custom (text fields) -> messages

Run: python scripts/01_curate_sft_data.py
"""
import json
import argparse
from pathlib import Path
from datasets import load_dataset

DATASETS = {
    "urdu_instruct": {
        "source": "large-traversaal/urdu-instruct",
        "split": "train",
        "format": "alpaca",
        "lang": "ur",
        "instruction_col": "instruction",
        "input_col": "input",
        "output_col": "output",
    },
    "sindhi_intelligence": {
        "source": "aakashMeghwar01/Sindhi-Intelligence-Core-SFT",
        "split": "train",
        "format": "alpaca",
        "lang": "sd",
        "instruction_col": "instruction",
        "input_col": "input",
        "output_col": "output",
    },
    "pashto_alpaca": {
        "source": "saillab/alpaca_pashto_taco",
        "split": "train",
        "format": "alpaca",
        "lang": "ps",
        "instruction_col": "instruction",
        "input_col": "input",
        "output_col": "output",
    },
    "urdu_news_gen": {
        "source": "AhmadMustafa/Urdu-Instruct-News-Article-Generation",
        "split": "train",
        "format": "alpaca",
        "lang": "ur",
        "instruction_col": "instruction",
        "input_col": "input",
        "output_col": "output",
    },
    "urdu_news_class": {
        "source": "AhmadMustafa/Urdu-Instruct-News-Category-Classification",
        "split": "train",
        "format": "alpaca",
        "lang": "ur",
        "instruction_col": "instruction",
        "input_col": "input",
        "output_col": "output",
    },
    "urdu_news_headline": {
        "source": "AhmadMustafa/Urdu-Instruct-News-Headline-Generation",
        "split": "train",
        "format": "alpaca",
        "lang": "ur",
        "instruction_col": "instruction",
        "input_col": "input",
        "output_col": "output",
    },
    "urdu_alpaca": {
        "source": "ravithejads/alpaca_urdu_cleaned_instruction",
        "split": "train",
        "format": "alpaca",
        "lang": "ur",
        "instruction_col": "instruction",
        "input_col": "input",
        "output_col": "output",
    },
    "punjabi_alpaca": {
        "source": "japneets/Alpaca_instruction_fine_tune_Punjabi",
        "split": "train",
        "format": "alpaca",
        "lang": "pa",
        "instruction_col": "instruction",
        "input_col": "input",
        "output_col": "output",
    },
}


def alpaca_to_messages(example, inst_col, inp_col, out_col, lang):
    """Convert Alpaca format to ChatML messages format for SFT."""
    instruction = str(example.get(inst_col, "")).strip()
    inp = str(example.get(inp_col, "")).strip()
    output = str(example.get(out_col, "")).strip()

    if not instruction or not output:
        return None

    # Build user prompt
    if inp and inp.lower() not in ["none", "nil", ""]:
        user_content = f"{instruction}\n\n{inp}"
    else:
        user_content = instruction

    # ChatML messages format (Qwen2.5 compatible)
    messages = [
        {"role": "system", "content": f"You are ZabaanAI, a helpful AI assistant fluent in {lang} and other Pakistan languages. Respond accurately and respectfully."},
        {"role": "user", "content": user_content},
        {"role": "assistant", "content": output},
    ]

    return {"messages": messages, "language": lang}


def process_dataset(name, config, output_dir, max_samples=None):
    """Download and format a dataset."""
    print(f"\nLoading: {name} ({config[\"source\"]})")
    try:
        ds = load_dataset(
            config["source"],
            split=config["split"],
            trust_remote_code=True,
        )
        if max_samples:
            ds = ds.select(range(min(max_samples, len(ds))))
        print(f"  Loaded {len(ds):,} examples")
    except Exception as e:
        print(f"  Error: {e}")
        return 0

    formatted = []
    for example in ds:
        if config["format"] == "alpaca":
            result = alpaca_to_messages(
                example,
                config["instruction_col"],
                config["input_col"],
                config["output_col"],
                config["lang"],
            )
        else:
            continue

        if result:
            formatted.append(result)

    # Save
    output_file = output_dir / f"{name}.jsonl"
    with open(output_file, "w", encoding="utf-8") as f:
        for item in formatted:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    print(f"  Saved {len(formatted):,} examples to {output_file.name}")
    return len(formatted)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output_dir", default="data/formatted")
    parser.add_argument("--max_samples", type=int, default=None)
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    print("=" * 60)
    print(" ZabaanAI-v2 SFT Data Curation")
    print("=" * 60)

    total = 0
    for name, config in DATASETS.items():
        count = process_dataset(name, config, output_dir, args.max_samples)
        total += count

    print(f"\nTotal formatted examples: {total:,}")
    print("=" * 60)


if __name__ == "__main__":
    main()