#!/usr/bin/env python
"""
Script to clean up test jobs from Nomad.
"""

import os
import sys
from dotenv import load_dotenv
from app.services.nomad_client import NomadService

# Load environment variables from .env file
load_dotenv()

def main():
    print("Cleaning up test jobs from Nomad...")
    
    # Check if NOMAD_ADDR is configured
    nomad_addr = os.getenv("NOMAD_ADDR")
    if not nomad_addr:
        print("Error: NOMAD_ADDR is not configured in .env file.")
        sys.exit(1)
    
    print(f"Connecting to Nomad at: {nomad_addr}")
    
    try:
        # Initialize the Nomad service
        nomad_service = NomadService()
        
        # List all jobs
        print("\nListing all jobs...")
        jobs = nomad_service.list_jobs()
        print(f"Found {len(jobs)} jobs")
        
        # Filter for test jobs (starting with "test-")
        test_jobs = [job for job in jobs if job.get('ID', '').startswith('test-')]
        print(f"Found {len(test_jobs)} test jobs:")
        
        # Print each test job's ID and status
        for job in test_jobs:
            print(f"  - {job.get('ID')}: {job.get('Status')}")
        
        # Confirm before proceeding
        if test_jobs:
            print("\nDo you want to stop and purge all these test jobs? (y/n)")
            response = input().strip().lower()
            
            if response == 'y':
                print("\nStopping and purging test jobs...")
                
                for job in test_jobs:
                    job_id = job.get('ID')
                    try:
                        print(f"Stopping and purging job: {job_id}...")
                        stop_response = nomad_service.stop_job(job_id, purge=True)
                        print(f"  - Success: {stop_response}")
                    except Exception as e:
                        print(f"  - Error stopping job {job_id}: {str(e)}")
                
                print("\nCleanup completed.")
            else:
                print("\nCleanup cancelled.")
        else:
            print("\nNo test jobs found to clean up.")
        
    except Exception as e:
        print(f"Error during cleanup: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()