"""Map Hugging Face dataset rows to chat messages for SFT."""
from __future__ import annotations


def fiction_tags_text_user(row: dict, title_col: str = "title", tags_col: str = "tags") -> str:
    parts = []
    title = (row.get(title_col) or "").strip()
    tags = (row.get(tags_col) or "").strip()
    if title:
        parts.append(f"Title: {title}")
    if tags:
        parts.append(f"Tags: {tags}")
    return "\n".join(parts) if parts else tags or title or ""


def row_to_messages(row: dict, schema: dict | None, mapping: dict | None) -> list[dict] | None:
    schema = schema or {}
    mapping = mapping or {}
    kind = schema.get("kind") or mapping.get("kind") or mapping.get("preset")

    if kind == "messages" and row.get("messages"):
        return _normalize_messages(row["messages"])

    if kind == "conversations" and row.get("conversations"):
        out = []
        for c in row["conversations"]:
            if not isinstance(c, dict):
                continue
            frm = c.get("from") or ""
            val = c.get("value") or ""
            role = "assistant" if frm in ("gpt", "assistant", "chatgpt") else "user"
            if frm in ("human", "user"):
                role = "user"
            if val:
                out.append({"role": role, "content": str(val)})
        return out or None

    if kind == "alpaca":
        instr = (row.get("instruction") or "").strip()
        inp = (row.get("input") or "").strip()
        output = (row.get("output") or "").strip()
        user = instr if not inp else f"{instr}\n{inp}"
        if user and output:
            return [{"role": "user", "content": user}, {"role": "assistant", "content": output}]
        return None

    if kind == "prompt_response":
        resp_col = schema.get("response_col") or "response"
        prompt = row.get("prompt") or ""
        resp = row.get(resp_col) or ""
        if prompt and resp:
            return [{"role": "user", "content": str(prompt)}, {"role": "assistant", "content": str(resp)}]
        return None

    if kind == "qa":
        q = row.get("question") or ""
        a = row.get("answer") or ""
        if q and a:
            return [{"role": "user", "content": str(q)}, {"role": "assistant", "content": str(a)}]
        return None

    if kind in ("fiction_tags_text", "preset_fiction_tags_text"):
        text = (row.get("text") or "").strip()
        user = fiction_tags_text_user(row)
        if user and text:
            return [{"role": "user", "content": user}, {"role": "assistant", "content": text}]
        return None

    if kind == "custom" or mapping.get("user_col"):
        user_col = mapping.get("user_col")
        asst_col = mapping.get("assistant_col")
        if user_col and asst_col:
            u = row.get(user_col) or ""
            a = row.get(asst_col) or ""
            if u and a:
                return [{"role": "user", "content": str(u)}, {"role": "assistant", "content": str(a)}]
        return None

    if row.get("messages"):
        return _normalize_messages(row["messages"])

    return None


def _normalize_messages(msgs) -> list[dict] | None:
    out = []
    for m in msgs:
        if not isinstance(m, dict):
            continue
        role = m.get("role") or "user"
        content = m.get("content") or m.get("text") or ""
        if content:
            out.append({"role": role, "content": str(content)})
    return out or None