from lib import config as config_module
from lib import logging as logging_lib
from lib import get_specs
from lib import utils
import threading
import socket
import aiohttp
import asyncio
import json
import time
import sys
import os

config = config_module.config
specs = get_specs.Specs()
log = logging_lib.log

def complete_loader(fail=False):
    sys.stdout.write("\r" + " " * 10 + "\r")  # Clear loader
    sys.stdout.flush()
    if fail:
        print("\rGetting server specifications - [ \x1b[31mERROR\x1b[0m")
    else:
        print("\rGetting server specifications - [ \x1b[32mCOMPLETE \x1b[0m]")

async def show_loader(loader_event):
    # Your loader display logic here
    animation = "|/-\\"
    idx = 0
    while True:
        if loader_event.is_set():
            sys.stdout.write("\r" + "🛠️ Getting server specifications " + animation[idx % len(animation)])
            sys.stdout.flush()
            idx += 1
            await asyncio.sleep(0.1)
        else:
            break

async def register_server(data):
    # Define the URL
    url = "https://api.clore.ai/register_server"

    # Convert the dictionary to JSON
    json_data = json.dumps(data)

    # Define the headers with content type as JSON
    headers = {
        "Content-Type": "application/json"
    }

    connector = aiohttp.TCPConnector(family=socket.AF_INET)

    async with aiohttp.ClientSession(connector=connector) as session:
        try:
            async with session.post(url, data=json_data, headers=headers, timeout=15) as response:
                if response.status == 200:
                    # Successful response
                    response_data = await response.json()
                    if "result" in response_data and response_data["result"]=="updated":
                        with open(config.auth_file, 'w') as f:
                            f.write(f"{data['key']}:{data['private_communication_token']}")
                        print("\x1b[32mServer successfully logged in\x1b[0m")
                        utils.run_command_v2("systemctl restart clore-hosting.service")
                        os._exit(0)
                    elif "result" in response_data and response_data["result"]=="already_updated":
                        print("\x1b[31mThis token is already used by some server\x1b[0m")
                    elif "error" in response_data and response_data["error"]=="invalid_key":
                        print("\x1b[31mInvalid token\x1b[0m")
                    else:
                        print(response_data)
                        print("\x1b[31mIssues at CLORE.AI, please try again later\x1b[0m")
                    os._exit(1)
                else:
                    # Failed response
                    print("\x1b[31mIssues connecting to CLORE.AI api, check your internet connection or try it later\x1b[0m")
                    os._exit(1)
        except asyncio.TimeoutError:
            # Request timed out
            print("\x1b[31mIssues connecting to CLORE.AI api, check your internet connection or try it later\x1b[0m")
            os._exit(1)

async def work_init(loader_event, init_token):
    loader_event.set()
    kernel = get_specs.get_kernel()
    nvidia_info = get_specs.get_nvidia_version()

    if utils.validate_cuda_version(nvidia_info.cuda_version):
        loader_event.clear()
        log.error("Cuda must be version 11.7+")
        os._exit(1)

    machine_specs = await specs.get(benchmark_internet=True, benchmark_disk=True, require_same_gpus=(not config.allow_mixed_gpus))

    loader_event.clear()
    complete_loader()
    
    max_str_l = 64
    for ckey in machine_specs.keys():
        if ckey == "net":
            pass
        elif ckey == "disk" and len(str(machine_specs[ckey])) > max_str_l:
            p = str(machine_specs[ckey]).split(' ')
            ft = ''
            req_size = p[-1]
            for x in range(len(p) - 1):
                ft += f"{p[x]} "
            ft = ft[:max_str_l - len(req_size) - 1]
            if ft[-3] == ' ':
                machine_specs[ckey] = f"{ft} {req_size}"
            else:
                machine_specs[ckey] = f"{ft} {req_size}"
        elif len(str(machine_specs[ckey])) > max_str_l:
            machine_specs[ckey] = str(machine_specs[ckey])[:max_str_l]


    del machine_specs["gpus"]
    del machine_specs["nvml_error"]
    del machine_specs["swap"]
    del machine_specs["data_root_location"]
    
    print(machine_specs)

    await register_server({
        "key": init_token,
        "private_communication_token": utils.generate_random_string(32),
        "specs": machine_specs
    })

async def init(init_token):
    loader_event = asyncio.Event()

    # Run the loader asynchronously
    loader_task = asyncio.create_task(show_loader(loader_event))

    # Run the initialization asynchronously
    await work_init(loader_event, init_token)

    # Cancel the loader task once the initialization is done
    loader_task.cancel()