import sys
import json
import random
import argparse

def split_dataset(
    input_filepath: str, 
    train_filepath: str = "train.jsonl", 
    validation_filepath: str = "validation.jsonl", 
    split_ratio: float = 0.95,
    balance_ratio: float = 1.0
):
    """
    Reads a .jsonl file, balances the positive/negative examples, shuffles,
    and splits it into training and validation files.
    """
    print(f"Loading data from {input_filepath}...")
    try:
        with open(input_filepath, 'r', encoding='utf-8') as f:
            lines = f.readlines()
    except FileNotFoundError:
        print(f"Error: Input file not found at {input_filepath}", file=sys.stderr)
        sys.exit(1)

    # 1. Separate into positive and negative examples
    positive_examples = []
    negative_examples = []
    for line in lines:
        try:
            entry = json.loads(line)
            if entry.get("target", "").strip():
                positive_examples.append(line)
            else:
                negative_examples.append(line)
        except json.JSONDecodeError:
            print(f"Warning: Skipping malformed JSON line: {line.strip()}", file=sys.stderr)

    print(f"Found {len(positive_examples)} positive examples and {len(negative_examples)} negative examples.")

    # 2. Shuffle both lists independently
    random.shuffle(positive_examples)
    random.shuffle(negative_examples)

    # 3. Balance the dataset based on the ratio
    num_positives = len(positive_examples)
    num_negatives_to_keep = int(num_positives * balance_ratio)

    if len(negative_examples) < num_negatives_to_keep:
        print(f"Warning: Not enough negative examples to meet a {balance_ratio}:1 ratio. "
              f"Using all {len(negative_examples)} negative examples.", file=sys.stderr)
        num_negatives_to_keep = len(negative_examples)

    print(f"Balancing dataset with {num_positives} positive and {num_negatives_to_keep} negative examples.")
    
    balanced_lines = positive_examples + negative_examples[:num_negatives_to_keep]
    
    # 4. Final shuffle to mix positive and negative examples
    print("Shuffling final balanced dataset...")
    random.shuffle(balanced_lines)

    # 5. Determine the split point and write files
    split_index = int(len(balanced_lines) * split_ratio)
    train_lines = balanced_lines[:split_index]
    validation_lines = balanced_lines[split_index:]

    print(f"Writing {len(train_lines)} lines to {train_filepath}...")
    with open(train_filepath, 'w', encoding='utf-8') as f:
        f.writelines(train_lines)

    print(f"Writing {len(validation_lines)} lines to {validation_filepath}...")
    with open(validation_filepath, 'w', encoding='utf-8') as f:
        f.writelines(validation_lines)
    
    print("\nSplitting complete.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Balance and split a .jsonl dataset into training and validation sets."
    )
    parser.add_argument("input_file", type=str, help="Path to the input .jsonl file.")
    parser.add_argument(
        "--split", 
        type=float, 
        default=0.95, 
        help="Ratio for the training set split (e.g., 0.95 for a 95/5 split). Default is 0.95."
    )
    parser.add_argument(
        "--ratio",
        type=float,
        default=1.0,
        help="Ratio of negative to positive examples (e.g., 1.0 for 1:1, 2.0 for 2:1). Default is 1.0."
    )
    
    args = parser.parse_args()
    
    if not 0 < args.split < 1:
        print("Error: Split ratio must be between 0 and 1.", file=sys.stderr)
        sys.exit(1)
        
    split_dataset(args.input_file, split_ratio=args.split, balance_ratio=args.ratio)