import os
import subprocess
import argparse
import sys

def run_command(command: str):
    print(f"Running: {command}")
    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    for line in process.stdout:
        print(line, end="")
    process.wait()
    if process.returncode != 0:
        print(f"Command failed with return code {process.returncode}")
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description="Integrated Local Training Pipeline (8GB VRAM Optimized)")
    parser.add_argument("--model-path", type=str, required=True, help="Path to local safetensors model")
    parser.add_argument("--raw-data", type=str, required=True, help="Path to raw data (file or dir)")
    parser.add_argument("--data-type", choices=['jsonl', 'dir'], default='jsonl', help="Type of raw data")
    parser.add_argument("--fields", nargs="+", default=["text"], help="Fields to collect")
    parser.add_argument("--output-dir", type=str, default="./finetuned_output", help="Where to save the model")
    
    args = parser.parse_args()

    # Step 1: Data Collection
    processed_data_dir = "./data/processed"
    print("=== Step 1: Collecting and Categorizing Data ===")
    fields_str = " ".join(args.fields)
    collect_cmd = f"python data_collector.py --source {args.raw_data} --type {args.data_type} --fields {fields_str} --output {processed_data_dir}"
    run_command(collect_cmd)

    # Step 2: Training (Aggregating all processed fields)
    print("\n=== Step 2: Starting 8GB VRAM Optimized Training ===")
    # We combine the processed data for training. 
    # For simplicity, we find the first available master file in the processed dirs
    master_data_file = None
    for root, _, files in os.walk(processed_data_dir):
        if "collected_data.jsonl" in files:
            master_data_file = os.path.join(root, "collected_data.jsonl")
            break
    
    if not master_data_file:
        print("No processed data found to train on.")
        sys.exit(1)

    # Triggering the optimized finetuner
    train_cmd = (
        f"python finetune.py "
        f"--model {args.model_path} "
        f"--dataset {master_data_file} "
        f"--output {args.output_dir}"
    )
    
    # Set Offline environment variables for security
    os.environ["HF_HUB_OFFLINE"] = "1"
    os.environ["WANDB_DISABLED"] = "true"
    
    run_command(train_cmd)
    print("\n=== Pipeline Complete ===")

if __name__ == "__main__":
    main()