#!/usr/bin/env python3
"""
Download ToolBench + APIGen-MT + ToolACE
and convert them to Qwen 2.5 SFT JSONL format
(with tool calling / function calling support).
"""

import os
import json
import gzip
import tarfile
import zipfile
import requests
from pathlib import Path
from tqdm import tqdm
from datasets import load_dataset
from huggingface_hub import hf_hub_download, snapshot_download

# ====================== CONFIG ======================
OUTPUT_DIR = Path("./qwen25_tool_sft")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

FINAL_JSONL = OUTPUT_DIR / "tool_sft_qwen25.jsonl"

# ====================================================

def download_file(url: str, dest: Path):
    if dest.exists():
        print(f"[skip] {dest.name} already exists")
        return
    print(f"Downloading {url} ...")
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        total = int(r.headers.get("content-length", 0))
        with open(dest, "wb") as f, tqdm(total=total, unit="B", unit_scale=True) as pbar:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
                pbar.update(len(chunk))


def to_qwen_messages(system: str | None, conversations: list[dict]) -> dict:
    """
    Convert a list of turns into Qwen 2.5 messages format.
    conversations: list of {"from": "human/gpt/function/...", "value": "..."}
    """
    messages = []
    if system:
        messages.append({"role": "system", "content": system})

    for turn in conversations:
        role = turn.get("from", "").lower()
        content = turn.get("value", "").strip()
        if not content:
            continue

        if role in ("human", "user"):
            messages.append({"role": "user", "content": content})
        elif role in ("gpt", "assistant"):
            messages.append({"role": "assistant", "content": content})
        elif role in ("function", "tool", "observation"):
            # Qwen-style tool response
            messages.append({"role": "tool", "content": content})
        else:
            # fallback
            messages.append({"role": "user", "content": content})

    return {"messages": messages}


# ----------------------------------------------------
# 1. ToolBench (official)
# ----------------------------------------------------
def process_toolbench():
    print("\n=== ToolBench ===")
    # ToolBench is available on Hugging Face
    try:
        ds = load_dataset("ToolBench/ToolBench", split="train", trust_remote_code=True)
    except Exception:
        # fallback to the processed version that many people use
        ds = load_dataset("lmsys/toolbench", split="train")

    count = 0
    with open(FINAL_JSONL, "a", encoding="utf-8") as fout:
        for sample in tqdm(ds, desc="ToolBench"):
            # ToolBench usually has "conversations" or "messages"
            convs = sample.get("conversations") or sample.get("messages") or []
            if not convs:
                continue

            # Some versions already have role/content
            if isinstance(convs[0], dict) and "role" in convs[0]:
                messages = []
                for m in convs:
                    role = m.get("role", "user")
                    content = m.get("content", "")
                    if role == "function":
                        role = "tool"
                    messages.append({"role": role, "content": content})
                record = {"messages": messages}
            else:
                record = to_qwen_messages(None, convs)

            if len(record["messages"]) >= 2:
                fout.write(json.dumps(record, ensure_ascii=False) + "\n")
                count += 1
    print(f"ToolBench → {count} samples")


# ----------------------------------------------------
# 2. APIGen-MT (multi-turn tool calling)
# ----------------------------------------------------
def process_apigen_mt():
    print("\n=== APIGen-MT ===")
    # Common locations / names
    possible = [
        "Salesforce/APIGen-MT",
        "Salesforce/xLAM-APIGen",
        "Salesforce/APIGen",
    ]
    ds = None
    for name in possible:
        try:
            ds = load_dataset(name, split="train")
            print(f"Loaded {name}")
            break
        except Exception:
            continue

    if ds is None:
        print("APIGen-MT not found on HF under common names. Skipping.")
        return

    count = 0
    with open(FINAL_JSONL, "a", encoding="utf-8") as fout:
        for sample in tqdm(ds, desc="APIGen-MT"):
            # APIGen usually has "messages" already close to OpenAI format
            messages = sample.get("messages") or sample.get("conversations")
            if not messages:
                continue

            # Normalize role names
            normalized = []
            for m in messages:
                role = m.get("role", "user").lower()
                content = m.get("content", "")
                if role == "function":
                    role = "tool"
                normalized.append({"role": role, "content": content})

            if len(normalized) >= 2:
                fout.write(json.dumps({"messages": normalized}, ensure_ascii=False) + "\n")
                count += 1
    print(f"APIGen-MT → {count} samples")


# ----------------------------------------------------
# 3. ToolACE
# ----------------------------------------------------
def process_toolace():
    print("\n=== ToolACE ===")
    possible = [
        "Team-ACE/ToolACE",
        "ToolACE/ToolACE",
        "microsoft/ToolACE",
    ]
    ds = None
    for name in possible:
        try:
            ds = load_dataset(name, split="train")
            print(f"Loaded {name}")
            break
        except Exception:
            continue

    if ds is None:
        print("ToolACE not found under common names. Trying alternative...")
        # Some people host processed versions
        try:
            ds = load_dataset("json", data_files="https://huggingface.co/datasets/Team-ACE/ToolACE/resolve/main/data/train.json")
        except Exception:
            print("Could not load ToolACE. Skipping.")
            return

    count = 0
    with open(FINAL_JSONL, "a", encoding="utf-8") as fout:
        for sample in tqdm(ds, desc="ToolACE"):
            messages = sample.get("messages") or sample.get("conversations") or []
            if not messages:
                continue

            normalized = []
            for m in messages:
                if isinstance(m, dict):
                    role = m.get("role", m.get("from", "user")).lower()
                    content = m.get("content", m.get("value", ""))
                else:
                    continue
                if role in ("function", "observation"):
                    role = "tool"
                elif role in ("human", "user"):
                    role = "user"
                elif role in ("gpt", "assistant"):
                    role = "assistant"
                normalized.append({"role": role, "content": content})

            if len(normalized) >= 2:
                fout.write(json.dumps({"messages": normalized}, ensure_ascii=False) + "\n")
                count += 1
    print(f"ToolACE → {count} samples")


# ----------------------------------------------------
# Main
# ----------------------------------------------------
if __name__ == "__main__":
    # Clear previous output if you want a fresh file
    if FINAL_JSONL.exists():
        print(f"Removing old {FINAL_JSONL}")
        FINAL_JSONL.unlink()

    process_toolbench()
    process_apigen_mt()
    process_toolace()

    # Final stats
    total = sum(1 for _ in open(FINAL_JSONL, "r", encoding="utf-8"))
    print(f"\n✅ Done! Total samples written → {FINAL_JSONL}")
    print(f"   Total lines: {total}")
    print("\nYou can now use this JSONL for Qwen2.5 SFT (tool calling / function calling).")