# process_hf_dataset.py
from datasets import load_dataset
import re
from parser import parse_python_code
from database import init_chromadb, store_program, DB_NAME, HF_DATASET_NAME, HF_KEY
import chromadb

def rename_variables(code):
    """Rename variables in Python code to input_var1, input_var2, etc."""
    # Simple variable name detection and renaming
    pattern = r'\b[a-zA-Z_]\w*\b'  # Match variable names (simple heuristic)
    variables = set()
    code_lines = code.split('\n')
    
    # Find all variable names (simplified approach)
    for line in code_lines:
        matches = re.findall(pattern, line)
        for match in matches:
            if match not in ['def', 'if', 'else', 'for', 'while', 'return', 'import', 'print', 'eval', 'str', 'int']:  # Exclude keywords
                variables.add(match)
    
    # Sort variables by first appearance (simplified, could improve with AST)
    sorted_vars = sorted(list(variables))
    var_map = {var: f"input_var{i+1}" for i, var in enumerate(sorted_vars)}
    
    # Replace variables in code
    new_code = code
    for old_var, new_var in var_map.items():
        new_code = re.sub(r'\b' + old_var + r'\b', new_var, new_code)
    
    return new_code

def process_hf_dataset():
    """Process the Hugging Face dataset and store programs in ChromaDB."""
    # Load the dataset
    dataset = load_dataset("iamtarun/python_code_instructions_18k_alpaca", split="train")
    
    # Initialize ChromaDB client
    client = init_chromadb()
    
    # Process each entry
    for entry in dataset:
        instruction = entry['instruction']
        output = entry['output']
        
        # Rename variables in the output code
        processed_code = rename_variables(output)
        
        # Parse the code to get parts and sequence
        parts, sequence = parse_python_code(processed_code)
        vectors = [part['vector'] for part in parts]
        
        # Generate description tokens from instruction
        description_tokens = f"task:{instruction.replace(' ', '_')}"
        
        # Store in ChromaDB with description
        store_program(client, processed_code, sequence, vectors, DB_NAME)
        
        # Update metadata with instruction as description
        collection = client.get_collection(DB_NAME)
        program_id = str(hash(processed_code))
        collection.update(
            ids=[program_id],
            metadatas=[{"sequence": ",".join(sequence), "description_tokens": description_tokens}]
        )
    
    # Save to Hugging Face Dataset
    save_chromadb_to_hf()

def save_chromadb_to_hf(dataset_name=HF_DATASET_NAME, token=HF_KEY):
    """Save ChromaDB data to Hugging Face Dataset."""
    client = init_chromadb()
    collection = client.get_collection(DB_NAME)
    
    # Fetch all data from ChromaDB
    results = collection.get(include=["documents", "metadatas", "embeddings"])
    data = {
        "code": results["documents"],
        "sequence": [meta["sequence"] for meta in results["metadatas"]],
        "vectors": results["embeddings"],  # ChromaDB already flattens embeddings
        "description_tokens": [meta.get('description_tokens', '') for meta in results["metadatas"]]
    }
    
    # Create a Hugging Face Dataset
    dataset = Dataset.from_dict(data)
    
    # Push to Hugging Face Hub
    dataset.push_to_hub(dataset_name, token=token)
    print(f"Dataset pushed to Hugging Face Hub as {dataset_name}")

if __name__ == "__main__":
    process_hf_dataset()