#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# /// script
# dependencies = ["datasets","huggingface_hub","hf_xet","tqdm"]
# ///
import argparse, json, os, random, hashlib
from pathlib import Path
from datasets import load_dataset
from huggingface_hub import HfApi, upload_folder, hf_hub_download
from tqdm import tqdm

def jd(obj):
    return json.dumps(obj, ensure_ascii=False, sort_keys=False, default=str)

def clean(s, max_chars=12000):
    if s is None:
        return ""
    if not isinstance(s, str):
        s = jd(s)
    s = s.replace("\r\n", "\n").replace("\r", "\n").strip()
    if len(s) > max_chars:
        s = s[:max_chars].rstrip() + "\n...[TRUNCATED]"
    return s

def compact(obj, max_chars=12000):
    s = jd(obj)
    if len(s) > max_chars:
        s = s[:max_chars].rstrip() + "...[TRUNCATED]"
    return s

def role_of(m):
    r = str(m.get("role", "")).lower().strip()
    if r in ("human", "user"):
        return "user"
    if r in ("assistant", "gpt", "model"):
        return "assistant"
    if r in ("tool", "function", "observation"):
        return "tool"
    if r in ("system", "developer"):
        return "system"
    return r or "user"

def content_of(m):
    for k in ("content", "text", "value", "message"):
        if k in m and m[k] is not None:
            return clean(m[k])
    return ""

def normalize_tool_call(c):
    if isinstance(c, str):
        try:
            c = json.loads(c)
        except Exception:
            return {"raw": c}
    if not isinstance(c, dict):
        return {"raw": c}
    fn = c.get("function") if isinstance(c.get("function"), dict) else c
    name = fn.get("name") or c.get("name") or c.get("tool_name") or c.get("tool")
    args = fn.get("arguments", c.get("arguments", c.get("args", c.get("parameters", {}))))
    if isinstance(args, str):
        try:
            args = json.loads(args)
        except Exception:
            pass
    out = {}
    if c.get("id"):
        out["id"] = c.get("id")
    if name:
        out["name"] = str(name)
    out["arguments"] = args if args is not None else {}
    return out

def extract_tool_calls(m):
    calls = []
    if isinstance(m.get("tool_calls"), list):
        calls.extend([normalize_tool_call(c) for c in m["tool_calls"]])
    if m.get("function_call"):
        calls.append(normalize_tool_call(m["function_call"]))
    if m.get("tool_call"):
        calls.append(normalize_tool_call(m["tool_call"]))
    return [c for c in calls if c]

def used_tool_names(conversations):
    names = set()
    for m in conversations if isinstance(conversations, list) else []:
        if not isinstance(m, dict):
            continue
        for c in extract_tool_calls(m):
            if c.get("name"):
                names.add(str(c["name"]))
        for k in ("name", "tool_name"):
            if m.get(k):
                names.add(str(m[k]))
    return names

def summarize_tool(t, max_desc=220):
    if not isinstance(t, dict):
        return None
    fn = t.get("function") if isinstance(t.get("function"), dict) else t
    name = fn.get("name") or t.get("name")
    if not name:
        return None
    desc = clean(fn.get("description", ""), max_desc)
    params = fn.get("parameters", {}) or fn.get("arguments", {})
    required, props = [], []
    if isinstance(params, dict):
        required = params.get("required") or []
        properties = params.get("properties") or params.get("arguments") or {}
        if isinstance(properties, dict):
            props = list(properties.keys())[:20]
    line = "- " + str(name)
    if desc:
        line += ": " + desc
    if required:
        line += " | required: " + ", ".join(map(str, required[:10]))
    elif props:
        line += " | fields: " + ", ".join(map(str, props[:16]))
    return line

def build_system(tools, conversations, max_tools=20, max_chars=5000):
    used = used_tool_names(conversations)
    lines = ["You can call tools when needed.", "Use only the available tool names and copy arguments exactly.", "", "Available tools:"]
    selected = []
    if isinstance(tools, list):
        for t in tools:
            line = summarize_tool(t)
            if not line:
                continue
            name = line[2:].split(":", 1)[0].split(" | ", 1)[0].strip()
            if used and name not in used:
                continue
            selected.append(line)
        if not selected:
            for t in tools[:max_tools]:
                line = summarize_tool(t)
                if line:
                    selected.append(line)
    lines.extend(selected[:max_tools] if selected else ["- no_tool: No tool available"])
    return clean("\n".join(lines), max_chars)

def tool_call_block(calls):
    if len(calls) == 1:
        return "TOOL_CALL:\n" + compact(calls[0])
    return "TOOL_CALLS:\n" + compact(calls)

def tool_result_block(m):
    payload = {}
    for k in ("name", "tool_name", "tool_call_id", "id"):
        if m.get(k):
            payload[k] = m[k]
    c = content_of(m)
    if c:
        payload["content"] = c
    else:
        for k in ("result", "observation", "output", "data"):
            if m.get(k) is not None:
                payload[k] = m[k]
                break
    return "TOOL_RESULT:\n" + compact(payload or m)

def normalize_row(row, source):
    conversations = row.get("conversations") or row.get("messages")
    tools = row.get("tools") or []
    if not isinstance(conversations, list) or not conversations:
        return None
    out = [{"role": "system", "content": build_system(tools, conversations)}]
    saw_user = False
    saw_assistant = False
    saw_tool_call = False
    for m in conversations:
        if not isinstance(m, dict):
            continue
        r = role_of(m)
        if r == "system":
            c = content_of(m)
            if c:
                out.append({"role": "system", "content": c})
            continue
        if r == "tool":
            out.append({"role": "user", "content": tool_result_block(m)})
            saw_user = True
            continue
        if r == "assistant":
            content = content_of(m)
            calls = extract_tool_calls(m)
            if calls:
                saw_tool_call = True
                block = tool_call_block(calls)
                content = (content + "\n\n" + block).strip() if content else block
            if content:
                out.append({"role": "assistant", "content": content})
                saw_assistant = True
            continue
        c = content_of(m)
        if c:
            out.append({"role": "user", "content": c})
            saw_user = True
    merged = []
    for m in out:
        if merged and merged[-1]["role"] == m["role"]:
            merged[-1]["content"] = clean(merged[-1]["content"] + "\n\n" + m["content"])
        else:
            merged.append(m)
    if not saw_user or not saw_assistant:
        return None
    return {"messages": merged, "source": source, "category": "toolmind_tool_call" if saw_tool_call else "toolmind_no_tool_call"}

def stable_key(obj):
    return hashlib.sha256(jd(obj.get("messages", [])).encode("utf-8")).hexdigest()

def try_reuse(out_repo_id, out_dir):
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    try:
        train = hf_hub_download(repo_id=out_repo_id, filename="train.jsonl", repo_type="dataset")
        val = hf_hub_download(repo_id=out_repo_id, filename="validation.jsonl", repo_type="dataset")
        (out / "train.jsonl").write_bytes(Path(train).read_bytes())
        (out / "validation.jsonl").write_bytes(Path(val).read_bytes())
        print("REUSED", out_repo_id, flush=True)
        return True
    except Exception as e:
        print("NO REUSE:", repr(e)[:250], flush=True)
        return False

def convert(a):
    out = Path(a.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    if a.reuse and try_reuse(a.out_repo_id, a.out_dir):
        return out
    print("Loading", a.dataset, a.split, flush=True)
    ds = load_dataset(a.dataset, split=a.split)
    total = len(ds)
    limit = total if a.max_rows <= 0 else min(a.max_rows, total)
    print("Rows", total, "limit", limit, flush=True)
    rows, seen, counts = [], set(), {}
    for i in tqdm(range(limit), desc="convert"):
        obj = normalize_row(dict(ds[i]), f"{a.dataset}:{a.split}:{i}")
        if not obj:
            continue
        k = stable_key(obj)
        if k in seen:
            continue
        seen.add(k)
        rows.append(obj)
        counts[obj["category"]] = counts.get(obj["category"], 0) + 1
    random.Random(a.seed).shuffle(rows)
    val_n = min(a.val_size, max(1, len(rows)//100))
    val = rows[:val_n]
    train = rows[val_n:]
    with (out/"train.jsonl").open("w", encoding="utf-8") as f:
        for r in train:
            f.write(jd(r)+"\n")
    with (out/"validation.jsonl").open("w", encoding="utf-8") as f:
        for r in val:
            f.write(jd(r)+"\n")
    (out/"README.md").write_text("---\nlicense: apache-2.0\nlanguage:\n- en\ntask_categories:\n- text-generation\n---\n\n# ToolMind converted to Scugnizz format\n\n" + json.dumps(counts, ensure_ascii=False, indent=2), encoding="utf-8")
    print("DONE", out, "TRAIN", len(train), "VAL", len(val), flush=True)
    print(json.dumps(counts, ensure_ascii=False, indent=2), flush=True)
    if rows[:1]:
        print("SAMPLE", json.dumps(rows[0], ensure_ascii=False)[:2000], flush=True)
    return out

def upload_dataset(folder, repo_id, private=False):
    token = os.environ.get("HF_TOKEN") or os.environ.get("UV_SCRIPT_HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
    api = HfApi(token=token)
    api.create_repo(repo_id, repo_type="dataset", private=private, exist_ok=True)
    upload_folder(repo_id=repo_id, repo_type="dataset", folder_path=str(folder), commit_message="Convert ToolMind to Scugnizz format", token=token)
    print("UPLOADED", repo_id, flush=True)

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--dataset", default="mlx-community/ToolMind")
    p.add_argument("--split", default="graph_syn_datasets")
    p.add_argument("--max-rows", type=int, default=50000)
    p.add_argument("--val-size", type=int, default=1000)
    p.add_argument("--seed", type=int, default=20260709)
    p.add_argument("--out-dir", default="data/toolmind-scugnizz-converted")
    p.add_argument("--upload", action="store_true")
    p.add_argument("--out-repo-id", default="ProjectScugnizz/toolmind-scugnizz-converted")
    p.add_argument("--private", action="store_true")
    p.add_argument("--reuse", action="store_true")
    a = p.parse_args()
    folder = convert(a)
    if a.upload:
        upload_dataset(folder, a.out_repo_id, a.private)

if __name__ == "__main__":
    main()