import json
import re
import argparse

# --- Emoji regex ---
EMOJI_PATTERN = re.compile(
    "["
    "\U0001F600-\U0001F64F"
    "\U0001F300-\U0001F5FF"
    "\U0001F680-\U0001F6FF"
    "\U0001F700-\U0001F77F"
    "\U0001F780-\U0001F7FF"
    "\U0001F800-\U0001F8FF"
    "\U0001F900-\U0001F9FF"
    "\U0001FA00-\U0001FAFF"
    "\U00002702-\U000027B0"
    "\U000024C2-\U0001F251"
    "]+",
    flags=re.UNICODE,
)

# --- Markdown cleanup ---
def remove_markdown(text: str) -> str:
    # code blocks
    text = re.sub(r"```.*?```", "", text, flags=re.DOTALL)
    text = re.sub(r"`.*?`", "", text)

    # images
    text = re.sub(r"!\[.*?\]\(.*?\)", "", text)

    # links → keep visible text
    text = re.sub(r"\[(.*?)\]\(.*?\)", r"\1", text)

    # bold / italic / strike
    text = re.sub(r"(\*\*|__)(.*?)\1", r"\2", text)
    text = re.sub(r"(\*|_)(.*?)\1", r"\2", text)
    text = re.sub(r"(~~)(.*?)\1", r"\2", text)

    # headings
    text = re.sub(r"^#{1,6}\s*", "", text, flags=re.MULTILINE)

    # blockquotes
    text = re.sub(r"^>\s*", "", text, flags=re.MULTILINE)

    # lists
    text = re.sub(r"^\s*[-*+]\s+", "", text, flags=re.MULTILINE)

    return text


def remove_emojis(text: str) -> str:
    return EMOJI_PATTERN.sub("", text)


def clean_text(text: str) -> str:
    # replace newlines with space
    text = text.replace("\n", " ")

    # collapse multiple spaces
    text = re.sub(r"\s+", " ", text)

    text = remove_markdown(text)
    text = remove_emojis(text)
    return text.strip()


def clean_obj(obj):
    if isinstance(obj, dict):
        return {k: clean_obj(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [clean_obj(x) for x in obj]
    elif isinstance(obj, str):
        return clean_text(obj)
    else:
        return obj


def main():
    parser = argparse.ArgumentParser(description="Clean JSONL from markdown and emojis")
    parser.add_argument("input", help="Input JSONL file")
    parser.add_argument("output", help="Output JSONL file")
    args = parser.parse_args()

    with open(args.input, "r", encoding="utf-8") as fin, \
         open(args.output, "w", encoding="utf-8") as fout:

        for i, line in enumerate(fin, 1):
            line = line.strip()
            if not line:
                continue

            try:
                data = json.loads(line)
                cleaned = clean_obj(data)
                fout.write(json.dumps(cleaned, ensure_ascii=False) + "\n")
            except Exception as e:
                print(f"[line {i}] Skipping due to error: {e}")


if __name__ == "__main__":
    main()