"""Build the mixed finance instruction dataset in chat ("messages") format.

Each source dataset gets a formatter that maps its rows to
    {"messages": [{"role": "user", ...}, {"role": "assistant", ...}], "source": name}
so TRL's SFTTrainer can apply any model's chat template at train time.

Sources that fail to load (renamed repo, script-based dataset on datasets>=3.0,
network) are skipped with a warning instead of killing the run.

Usage:
    python -m src.data.prepare_datasets --out data/finance_sft \
        [--max-per-source 20000] [--push-to-hub user/finance-sft-mix]
"""

import argparse
import random

from datasets import Dataset, concatenate_datasets, load_dataset

SYSTEM_PROMPT = (
    "You are a financial analysis assistant with expertise in markets, filings, "
    "accounting standards (IFRS/GAAP), and financial regulation. Be precise, cite "
    "the relevant standard or filing section when applicable, and say so when you "
    "are unsure. You do not give personalized investment advice."
)


def _chat(user, assistant, source):
    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ],
        "source": source,
    }


def _fmt_instruction_style(row, source):
    """FinGPT-style rows: instruction / input / output."""
    user = row["instruction"]
    if row.get("input"):
        user = f"{user}\n\n{row['input']}"
    return _chat(user, row["output"], source)


def _fmt_finance_alpaca(row, source):
    user = row["instruction"]
    if row.get("input"):
        user = f"{user}\n\n{row['input']}"
    return _chat(user, row["output"], source)


def _fmt_phrasebank(row, source):
    labels = {0: "negative", 1: "neutral", 2: "positive"}
    user = (
        "Classify the sentiment of this financial news sentence as positive, "
        f"negative, or neutral:\n\n{row['sentence']}"
    )
    return _chat(user, labels[row["label"]], source)


# name -> (hf repo, config, split, formatter, mixture weight)
SOURCES = {
    "fingpt-sentiment": ("FinGPT/fingpt-sentiment-train", None, "train", _fmt_instruction_style, 1.0),
    "fingpt-fiqa-qa": ("FinGPT/fingpt-fiqa_qa", None, "train", _fmt_instruction_style, 1.5),
    "fingpt-headline": ("FinGPT/fingpt-headline", None, "train", _fmt_instruction_style, 0.5),
    "fingpt-finred": ("FinGPT/fingpt-finred", None, "train", _fmt_instruction_style, 0.5),
    "finance-alpaca": ("gbharti/finance-alpaca", None, "train", _fmt_finance_alpaca, 1.5),
    "financial-phrasebank": ("takala/financial_phrasebank", "sentences_66agree", "train", _fmt_phrasebank, 0.5),
}


def build(max_per_source, seed=42):
    random.seed(seed)
    parts = []
    for name, (repo, config, split, fmt, weight) in SOURCES.items():
        try:
            ds = load_dataset(repo, config, split=split) if config else load_dataset(repo, split=split)
        except Exception as e:  # noqa: BLE001 - skip broken sources, keep the run alive
            print(f"[skip] {name} ({repo}): {type(e).__name__}: {e}")
            continue
        n = min(len(ds), int(max_per_source * weight))
        ds = ds.shuffle(seed=seed).select(range(n))
        rows = [fmt(row, name) for row in ds]
        parts.append(Dataset.from_list(rows))
        print(f"[ok]   {name}: {n} examples")

    if not parts:
        raise RuntimeError("No sources loaded — check network / dataset names.")
    mixed = concatenate_datasets(parts).shuffle(seed=seed)
    return mixed.train_test_split(test_size=0.02, seed=seed)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data/finance_sft")
    ap.add_argument("--max-per-source", type=int, default=20000)
    ap.add_argument("--push-to-hub", default=None, help="e.g. user/finance-sft-mix (pushed private)")
    args = ap.parse_args()

    ds = build(args.max_per_source)
    print(ds)
    ds.save_to_disk(args.out)
    print(f"Saved to {args.out}")
    if args.push_to_hub:
        ds.push_to_hub(args.push_to_hub, private=True)
        print(f"Pushed to hub: {args.push_to_hub} (private)")


if __name__ == "__main__":
    main()