import json
import re
import os

BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
TRAIN_PATH = os.path.join(BASE_DIR, "data/mintoak/train.jsonl")
VALID_PATH = os.path.join(BASE_DIR, "data/mintoak/valid.jsonl")

citation_pattern = re.compile(r'(?i)(?:\n\s*)*(?:👉|\ud83d\udc49)?\s*For\s+more\s+details,?\s+visit:.*$', re.DOTALL)

def clean_file(path):
    if not os.path.exists(path):
        print(f"File not found: {path}")
        return
        
    cleaned_lines = []
    count = 0
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            data = json.loads(line)
            for msg in data.get("messages", []):
                if msg.get("role") == "assistant":
                    original = msg["content"]
                    cleaned = citation_pattern.sub("", original).strip()
                    if cleaned != original:
                        msg["content"] = cleaned
                        count += 1
            cleaned_lines.append(json.dumps(data, ensure_ascii=False))
            
    with open(path, "w", encoding="utf-8") as f:
        for line in cleaned_lines:
            f.write(line + "\n")
            
    print(f"Cleaned {count} assistant messages in {os.path.basename(path)}.")

if __name__ == "__main__":
    clean_file(TRAIN_PATH)
    clean_file(VALID_PATH)