from os import path, environ
import sys
import os
import _thread
import time
from diagnosis.dataloader.s3loader import S3Loader
from diagnosis.dataloader.gcsloader import GCSLoader
from diagnosis.dataloader.azureblobloader import AzureBlobLoader
from diagnosis.dataloader.localloader import LocalFileLoader
from o11ydiagnosis.engine.log import log
import task

pa = path.abspath(path.join(path.dirname(__file__), '../..'))
sys.path.append(pa)
class EngineConfig:
    def __init__(self):
        self.engine_id = os.environ["ENGINE_ID"]
        self.bucket_name = os.environ['BUCKET_NAME'] if "BUCKET_NAME" in os.environ else None
        self.cache_root = os.environ["CACHE_ROOT"] if "CACHE_ROOT" in os.environ else "cache"
        self.o11y_auth0_client_id = os.environ[
            "O11Y_AUTH0_CLIENT_ID"] if "O11Y_AUTH0_CLIENT_ID" in os.environ else None
        self.o11y_auth0_client_secret = os.environ[
            "O11Y_AUTH0_CLIENT_SECRET"] if "O11Y_AUTH0_CLIENT_SECRET" in os.environ else None
        self.env = os.environ[
            "ENV"] if "ENV" in os.environ else "dev"
        self.start_timestemp = int(
            os.environ["START_TIMESTEMP"]) if "START_TIMESTEMP" in os.environ else None
        self.end_timestemp = int(
            os.environ["END_TIMESTEMP"]) if "END_TIMESTEMP" in os.environ else None
        self.provider = os.environ["PROVIDER"] if "PROVIDER" in os.environ else "aws"
        self.gcp_project = os.environ["GCP_PROJECT"] if "GCP_PROJECT" in os.environ else None
        self.use_local_loader = os.environ["USE_LOCAL_LOADER"] if "USE_LOCAL_LOADER" in os.environ else 'True'
        self.remote_loader = self.get_remote_loader()
        pod_name = os.environ.get('POD_NUMBER', '')
        # Extract the number from the pod name (e.g., 'o11ydiagnosis-deltalake-0' -> 0)
        self.pod_number = int(pod_name.split('-')[-1]) if pod_name else 0
        self.total_replicas = int(os.environ.get('TOTAL_REPLICAS', '1'))

    def get_remote_loader(self):
            if 'USE_LOCAL_LOADER' in environ and environ['USE_LOCAL_LOADER'] == 'True':
                return LocalFileLoader(environ['LOCAL_LOADER_PATH'])
            else:
                # add azure blob loader
                if 'REMOTE_PROVIDER' in environ and environ['REMOTE_PROVIDER'] == 'azure':
                    return AzureBlobLoader("https://{}.blob.core.windows.net".format(environ['AZURE_STORAGE_ACCOUNT']), environ['BUCKET_NAME'])
                elif 'REMOTE_PROVIDER' in environ and environ['REMOTE_PROVIDER'] == 'gcp':
                    return GCSLoader(environ['BUCKET_NAME'], environ['GCP_PROJECT'])
                return S3Loader(environ['BUCKET_NAME'])

def func_alive():
    return True
def func_not_alive():
    return False

def start():
    config = EngineConfig()
    _thread.start_new_thread(clean_cache_loop, (config,))
    while True:
        task_list = get_task_list(config)

        to_add_list, to_remove_list = compare(task_list)
        for task_file in to_add_list:
            if task_file.endswith(".yaml"):
                task_manager = task.Manager(config, task_file)
                task.running_tasks[task_file] = func_alive
                _thread.start_new_thread(startManager, (task_manager,))

        for task_file in to_remove_list:
            task.running_tasks[task_file] = func_not_alive

        time.sleep(300)

def compare(task_list):
    to_add = []
    to_remove = []

    for t in task_list:
            if t not in task.running_tasks.keys():
                to_add.append(t)
    for r in task.running_tasks.keys():
        if r not in task_list:
            to_remove.append(r)

    return to_add,to_remove
def startManager(task_manager):
    task_manager.start()

def get_task_list(config):
    path = "diagnosis/tasks/processor-{}/".format(config.engine_id)
    try:
        task_list = config.remote_loader.get_file_list(path)
        filtered_tasks = []
        for t in task_list:
            if not t.endswith(".yaml"):
                continue
            try:
                task_num = int(t[-8:-5])
            except ValueError:
                task_num = 0
            if task_num % config.total_replicas == config.pod_number:
                filtered_tasks.append(t)
        return filtered_tasks
    except Exception:
        # we should return false if we can't get task list
        log.exception("failed to get task list")
        return []

def clean_cache_loop(config: EngineConfig):
    while True:
        now = time.time()
        log.info("checking cache")
        try:
            for root, _, files in os.walk(config.cache_root):
                for f in files:
                    filepath = os.path.join(root, f)
                    if os.path.isfile(filepath):
                        last_modified_time = os.path.getmtime(filepath)
                        if filepath.find("slowlog-merged") >= 0:
                            max_alive_secs = 2 * 3600 # 2 hours
                        elif filepath.find("regions-merged") >=0:
                            max_alive_secs = 2 * 3600 # 2 hours
                        else:
                            if filepath.find("slowlog-stats") >= 0:
                                if filepath.find("slowlog-stats/1h") >= 0:
                                    max_alive_secs = 8 * 24 * 3600 # 8 days
                                else:
                                    max_alive_secs = 2 * 3600 # 2 hours
                            else:
                                max_alive_secs = 8 * 24 * 3600 # 8 days
                        if (now - last_modified_time) >= max_alive_secs:
                            os.remove(filepath)
                            log.info(f"removed file from cache: {filepath}")
        except Exception:
            log.exception("error when clean cache")
        time.sleep(3600)

if __name__ == '__main__':
    # environ['ENGINE_ID'] = "test"
    # environ['CACHE_ROOT'] = "./cache"
    # environ['PROVIDER'] = "aws"
    # environ['USE_LOCAL_LOADER'] = "True"
    # environ['LOCAL_LOADER_PATH'] = "./local"
    # environ['START_TIMESTEMP'] = "1710021600"
    # environ['O11Y_AUTH0_CLIENT_ID'] = "Sd4aekPCa5nMUSQBIvf8eGSL895WtV4o"
    # environ['O11Y_AUTH0_CLIENT_SECRET'] = "-SLP5UvxJdwGbVaKupxgv53T6otXguApPoG0fBsWCIqDG8_TuvlYnlkhPQHcWaKT"
    # environ['ENV'] = "prod"
    # environ['BUCKET_NAME'] = "o11y-prod-shared-us-east-1"
    # environ['IS_PARALLEL'] = "False"
    sys.stdout = open(sys.stdout.fileno(), mode='w', buffering=1)

    start()