import os
import re
from datasets import Dataset

def prepare_data():
    source_dir = "/home/jason/robot_share"
    output_dir = "/home/jason/Downloads/xiaozhi-mcp/local_processed_dataset"
    chunk_size = 1000  # Keeping it within your token limit
    overlap = 128     
    
    documents = []
    
    print(f"Scanning directory: {source_dir}")
    if not os.path.exists(source_dir):
        print(f"Error: Source directory '{source_dir}' does not exist.")
        return
        
    for root, _, files in os.walk(source_dir):
        for file in files:
            if file.endswith(".txt") or file.endswith(".md"):
                file_path = os.path.join(root, file)
                try:
                    with open(file_path, "r", encoding="utf-8") as f:
                        content = f.read()
                    
                    content = content.strip()
                    if len(content) < 20:
                        continue
                        
                    start = 0
                    while start < len(content):
                        end = start + chunk_size
                        chunk = content[start:end].strip()
                        
                        if len(chunk) > 10:
                            chunk_lower = chunk.lower()
                            priority_keys = ["narelle", "kaitlyn", "tim", "peter", "emily"]
                            is_high_priority = any(re.search(rf"\b{k}\b", chunk_lower) for k in priority_keys)
                            
                            # Standard prompt entry
                            chat_messages = [
                                {"role": "system", "content": "You are Jason. Speak in the first person ('I', 'me', 'my'). Use your personal memory files to answer accurately."},
                                {"role": "user", "content": "What are your memories or thoughts regarding this aspect of your life?"},
                                {"role": "assistant", "content": chunk}
                            ]
                            documents.append({"messages": chat_messages})
                            
                            # If it's a core person/project, give it specific short-form variations so it doesn't freeze
                            if is_high_priority:
                                if re.search(r"\bnarelle\b", chunk_lower):
                                    short_chat = [
                                        {"role": "system", "content": "You are Jason. Speak in the first person. Relocate your raw memories instantly."},
                                        {"role": "user", "content": "Do you ever think about Narelle?"},
                                        {"role": "assistant", "content": chunk}
                                    ]
                                    documents.append({"messages": short_chat})
                                    documents.append({"messages": chat_messages})
                                if re.search(r"\bkaitlyn\b", chunk_lower):
                                    short_chat = [
                                        {"role": "system", "content": "You are Jason. Speak in the first person. Relocate your raw memories instantly."},
                                        {"role": "user", "content": "Who is Kaitlyn?"},
                                        {"role": "assistant", "content": chunk}
                                    ]
                                    documents.append({"messages": short_chat})
                                    documents.append({"messages": chat_messages})
                                if re.search(r"\btim\b", chunk_lower):
                                    short_chat = [
                                        {"role": "system", "content": "You are Jason. Speak in the first person. Relocate your raw memories instantly."},
                                        {"role": "user", "content": "What about Tim?"},
                                        {"role": "assistant", "content": chunk}
                                    ]
                                    documents.append({"messages": short_chat})
                                    documents.append({"messages": chat_messages})
                                if re.search(r"\bpeter\b", chunk_lower):
                                    short_chat = [
                                        {"role": "system", "content": "You are Jason. Speak in the first person. Relocate your raw memories instantly."},
                                        {"role": "user", "content": "Who is Peter?"},
                                        {"role": "assistant", "content": chunk}
                                    ]
                                    documents.append({"messages": short_chat})
                                    documents.append({"messages": chat_messages})
                                if re.search(r"\bemily\b", chunk_lower):
                                    short_chat = [
                                        {"role": "system", "content": "You are Jason. Speak in the first person. Relocate your raw memories instantly."},
                                        {"role": "user", "content": "Tell me about Emily."},
                                        {"role": "assistant", "content": chunk}
                                    ]
                                    documents.append({"messages": short_chat})
                                    documents.append({"messages": chat_messages})
                            
                        start += chunk_size - overlap
                except Exception as e:
                    print(f"Error reading file {file_path}: {e}")
                    
    print(f"Extracted {len(documents)} structured chat chunks.")
    
    if not documents:
        print("No valid documents found. Dataset preparation aborted.")
        return
        
    dataset = Dataset.from_list(documents)
    split_dataset = dataset.train_test_split(test_size=0.1, seed=42)
    split_dataset.save_to_disk(output_dir)
    print(f"Success: Processed dataset saved to '{output_dir}'")
    print(split_dataset)

if __name__ == "__main__":
    prepare_data()