# Copyright (c) 2024 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

import os
from enum import Enum, StrEnum
from time import sleep

import requests
from requests import Response
from sumologic.sumologic import SumoLogic


class JobStatus(StrEnum):
    FAILED = 'Failed'
    IN_PROGRESS = 'InProgress'
    SUCCESS = 'Success'


class JobType(Enum):
    APP_INSTALL = 'app_install'
    COPY = 'copy'
    DELETE = 'delete'
    EXPORT = 'export'
    IMPORT = 'import'
    LOOKUP = 'lookup'
    SYNC = 'sync'


def get_sumo_client() -> SumoLogic:
    return SumoLogic(accessId=os.environ['SUMOLOGIC_ACCESSID'], accessKey=os.environ['SUMOLOGIC_ACCESSKEY'])


def get_job_status_fn(job_type: JobType, job_id: str, content_id: str = '') -> Response:
    client = get_sumo_client()
    if job_type == JobType.APP_INSTALL:
        return client.check_app_install_status(job_id=job_id)
    if job_type == JobType.COPY:
        return client.check_copy_status(content_id=content_id, job_id=job_id)
    if job_type == JobType.DELETE:
        return client.check_delete_status(content_id=content_id, job_id=job_id)
    raise Exception(f'Unsupported job type: {job_type}')


def wait_on_job(job_id: str, job_type: JobType, content_id: str = '', sleep_time_seconds: int = 1, retry_count: int = 60):
    status = JobStatus.IN_PROGRESS
    error = {}
    count = 0
    while status == JobStatus.IN_PROGRESS:
        if count > retry_count:
            raise Exception(f'Job took too long to complete. Job ID: {job_id}')
        res = get_job_status_fn(job_type=job_type, job_id=job_id, content_id=content_id).json()
        status = res['status']
        error = res['error']
        if status == JobStatus.IN_PROGRESS:
            sleep(sleep_time_seconds)
            count += 1
    if status == JobStatus.FAILED:
        raise Exception(f'Job failed with message: {error}')


def delete_content(content_id: str):
    client = get_sumo_client()
    try:
        job_id = client.delete_content(content_id=content_id).json()['id']
        wait_on_job(job_id=job_id, job_type=JobType.DELETE, content_id=content_id)
    except requests.HTTPError as e:
        if e.response.json().get('errors', [])[0].get('code') != 'content:content_not_found':
            raise e
