import requests
import json
import time
import logging
import os

log_dir = '/tmp'
if not os.path.exists(log_dir):
    os.makedirs(log_dir)
log_file = os.path.join(log_dir, 'nexus_task_logs.log')

logging.basicConfig(filename=log_file, level=logging.INFO, 
                    format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger()

with open('artifactory_repositories.json', 'r') as file:
    artifactory_repositories = json.load(file)

headers = {
    'Content-Type': 'application/json'
}

nexus_url_base = 'https://nexus-test.devops.lmru.tech'
nexus_username = 'admin'
nexus_password = 'nexus'

def create_or_update_script(script_name, script_content):
    script_url = f"{nexus_url_base}/service/rest/v1/script"
    script_payload = {
        "name": script_name,
        "type": "groovy",
        "content": script_content
    }
    
    response = requests.get(f"{script_url}/{script_name}", headers=headers, auth=(nexus_username, nexus_password))
    if response.status_code == 200:
        response = requests.put(f"{script_url}/{script_name}", headers=headers, auth=(nexus_username, nexus_password), data=json.dumps(script_payload))
    else:
        response = requests.post(script_url, headers=headers, auth=(nexus_username, nexus_password), data=json.dumps(script_payload))
    
    if response.status_code in [200, 201, 204]:
        logger.info(f"Script {script_name} successfully created/updated")
        return True
    else:
        logger.error(f"Error creating/updating script {script_name}: {response.status_code} {response.text}")
        return False

def run_script(script_name, args):
    script_run_url = f"{nexus_url_base}/service/rest/v1/script/{script_name}/run"
    response = requests.post(script_run_url, headers=headers, auth=(nexus_username, nexus_password), data=json.dumps(args))
    
    if response.status_code == 200:
        logger.info(f"Script {script_name} successfully executed")
        return response.json()
    else:
        logger.error(f"Error executing script {script_name}: {response.status_code} {response.text}")
        return None

def check_task_status(task_id):
    task_url = f"{nexus_url_base}/service/rest/v1/tasks/{task_id}"
    while True:
        response = requests.get(task_url, headers=headers, auth=(nexus_username, nexus_password))
        if response.status_code == 200:
            task_info = response.json()
            logger.debug(f"Raw task_info: {task_info}")
            status = task_info['currentState']
            logger.info(f"Task {task_id} status: {status}")
            if status in ['WAITING', 'FAILED', 'OK']:
                last_run_result = task_info['lastRunResult']
                logger.info(f"Task {task_id} last run result: {last_run_result}")
                return status, last_run_result
        else:
            logger.error(f"Error checking status of task {task_id}: {response.status_code} {response.text}")
        time.sleep(10)

groovy_script_content = '''
import org.sonatype.nexus.scheduling.TaskConfiguration
import org.sonatype.nexus.scheduling.TaskInfo
import org.sonatype.nexus.scheduling.TaskScheduler

import groovy.json.JsonOutput
import groovy.json.JsonSlurper

class TaskXO {
  String typeId
  Boolean enabled
  String name
  String alertEmail
  Map<String, String> properties
}

TaskXO task = new JsonSlurper().parseText(args)

TaskScheduler scheduler = container.lookup(TaskScheduler.class.name)

TaskConfiguration config = scheduler.createTaskConfigurationInstance(task.typeId)
config.enabled = task.enabled
config.name = task.name
config.alertEmail = task.alertEmail
task.properties?.each { key, value -> config.setString(key, value) }
TaskInfo taskInfo = scheduler.scheduleTask(config, scheduler.scheduleFactory.manual())
JsonOutput.toJson(taskInfo)
'''

if create_or_update_script("create_task", groovy_script_content):
    for repo in artifactory_repositories:
        logger.info("="*50)
        repo_name = repo['key']
        repo_type = repo['type']

        if repo_type != 'LOCAL':
            continue

        task_name = f"{repo_name}-import"
        args = {
            "typeId": "repository.import",
            "enabled": True,
            "name": task_name,
            "alertEmail": None,
            "properties": {
                "repositoryName": repo_name,
                "sourceDir": f"/opt/art-migr/repositories/{repo_name}"
            }
        }

        task_info = run_script("create_task", args)

        if task_info:
            logger.debug(f"Task info: {task_info}")
            result = json.loads(task_info['result'])
            if 'id' in result:
                task_id = result['id']
                logger.info(f"Task {task_name} successfully created with ID {task_id}")
                
                task_run_url = f"{nexus_url_base}/service/rest/v1/tasks/{task_id}/run"
                run_response = requests.post(task_run_url, headers=headers, auth=(nexus_username, nexus_password))
                
                if run_response.status_code == 204:
                    logger.info(f"Task {task_name} successfully started")
                    status, last_run_result = check_task_status(task_id)
                    if status == 'FAILED':
                        logger.error(f"Task {task_name} failed with result: {last_run_result}")
                        break
                else:
                    logger.error(f"Error starting task {task_name}: {run_response.status_code} {run_response.text}")
                    break
                
                time.sleep(5)
            else:
                logger.error(f"Task ID not found in the response for task {task_name}")
                break
        else:
            logger.error(f"Failed to create task {task_name}")
            break