# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto #0
__all__ = ['OVHGPUManager']

# %% ../nbs/00_core.ipynb #48db86d2-8f14-4a81-a6b4-4b5ecae5eaed
import ovh
from fastcore.basics import patch
from loguru import logger
from typing import Any

# Simple, direct imports - let Python handle the module resolution
try:
    from ovhmanager.project import OVHProject
    from ovhmanager.catalog import GPUCatalog
    from ovhmanager.instance import GPUInstance
except ImportError:
    # Fallback for development/notebook environment
    from .project import OVHProject
    from .catalog import GPUCatalog
    from .instance import GPUInstance

# %% ../nbs/00_core.ipynb #851196c9-23bf-47da-9b36-c78d9cccdce8
class OVHGPUManager:
    def __init__(self, endpoint=None, application_key=None, application_secret=None, consumer_key=None):
        """
        OVHGPUManager is a OVH API client wrapper that enables to create, delete, start and stop GPU instances
        
        By default the OVH client expects the following environment variables to be set, but you can overwrite them using the provided arguments:
        - OVH_ENDPOINT
        - OVH_APPLICATION_KEY
        - OVH_APPLICATION_SECRET
        - OVH_CONSUMER_KEY

        Parameters
        ----------
        endpoint : str, optional
            API endpoint like 'ovh-eu'
        application_key : str, optional
            Application key
        application_secret : str, optional
            Application secret
        consumer_key : str, optional
            Consumer key
        """
        self.client = ovh.Client(
        endpoint=endpoint,
        application_key=application_key,
        application_secret=application_secret,
        consumer_key=consumer_key)

        self._base_catalog_url = "/order/catalog/public/cloud"
        self.region = None
        self.project = None
        self.consumption_mode = None
        self.catalog_dict = None
        self.catalog = None
        self.gpu_type = None
        self.gpu_flavor_id = None
        self.instance_image_id = None
        self.instance_image_name = None
        self.ssh_key_id = None
        self.ssh_key_name = None
        self.last_gpu_instance = None
        self.last_instance_dict = None
        self.gpu_instances = []

# %% ../nbs/00_core.ipynb #0255805d-b5c4-4601-90dc-0b8ba5010a8d
@patch
def list_projects(self:OVHGPUManager)->list:
    """
    Returns the list of available project IDs for the OVH Cloud account
    """
    return self.client.get('/cloud/project')

# %% ../nbs/00_core.ipynb #5bf0afe2-ace5-44bc-adc0-a44883c943f9
@patch
def select_project(self:OVHGPUManager, project_idx:str=None)->Any:
    """
    Lists available projects for the OVH Cloud account and prompts user to select a project if no project id is passed as an argument
    
    Parameters
    ----------
    project_idx : str, optional
        Project OVH ID
    """
    logger.info("Available projects:")
    projects = self.list_projects()
    for i, p in enumerate(projects):
        print(f"{i+1}. {p}")
    if not project_idx:
        idx = int(input("Select project number: ")) - 1
        project_idx = projects[idx]
    self.project = OVHProject(project_idx)
    logger.info(f"Selected project: {self.project.id}")
    return self.project

# %% ../nbs/00_core.ipynb #00e6bfde-8a00-4a07-bd83-22bc3be5621f
@patch
def list_regions(self:OVHGPUManager)->list:
    """
    Returns the list region for the selected project
    """
    return self.client.get(f'{self.project.base_url}/region')
    

# %% ../nbs/00_core.ipynb #19cdd001-ac26-4637-97f8-ef0d5064991d
@patch
def select_region(self:OVHGPUManager, region:str=None)->Any:
    """
    Lists available region for the OVH Project and prompts user to select a region if no region is passed as an argument
    This method sets both the OVHGPUManager and the OVHProject
    
    Parameters
    ----------
    region : str, optional
        OVH region
    """
    if not region:
        logger.info("Available region:")
        regions = self.list_regions()
        for i, p in enumerate(regions):
            print(f"{i+1}. {p}")
        idx = int(input("Select region number: ")) - 1
        region = regions[idx]
    self.region = (region)
    if self.project:
        self.project.set_region(self.region)
    logger.info(f"Selected region: {self.region}")
    return self.region

# %% ../nbs/00_core.ipynb #6d8b8b3d
@patch
def select_consumption_mode(self:OVHGPUManager, consumption_mode:str=None)->str:
    """
    Lists available OVH ressources consumption modes and prompts user to select a consumption mode if no mode is passed as an argument

    Parameters
    ----------
    consumption_mode : str, optional
        Ressource consumption mode
    """
    if not consumption_mode:
        logger.info("Available consumption modes:")
        print("\nConsumption modes:")
        print("1. consumption (hourly)")
        print("2. monthly.postpaid")
        mode_choice = int(input("Select mode: "))
        consumption_mode = 'consumption' if mode_choice == 1 else 'monthly.postpaid'
    self.consumption_mode = consumption_mode
    logger.info(f"Selected consumption mode: {self.consumption_mode}")
    return self.consumption_mode

# %% ../nbs/00_core.ipynb #33167b7d-451f-4c34-8c9a-b63273af3c2d
@patch
def retrieve_catalog_dict(self:OVHGPUManager, subsidiary='FR') ->dict:
    """
    Retrieves an OVH catalog for a specific OVH subsidiary

    Parameters
    ----------
    subsidiary : str, required, defaults to FR
        OVH catalpg subsidiary/country code
    """
    self.catalog_dict = self.client.get(f'{self._base_catalog_url}', ovhSubsidiary=subsidiary, region=self.region)
    logger.info(f"Retrieved catalog dictionary for {subsidiary} subsidiary")

    return self.catalog_dict
    

# %% ../nbs/00_core.ipynb #5043d3b2-1fa2-45da-b149-03cb7cb793d3
@patch
def select_gpu_instance_type(self:OVHGPUManager, subsidiary='FR', gpu_type=None, consumption_mode=None)->Any:
    """
    Lists available OVH GPU instance types and prompts user to select a GPU type if no gpu type is passed as an argument

    Parameters
    ----------
    subsidiary : str, required, defaults to FR
        OVH catalog subsidiary/country code
    gpu_type : str, optional
        GPU instance type
    consumption_mode : str, optional
        Ressource consumption mode
    """
    # 0. make sure region and projects are set
    if not self.project:
        self.select_project()
    if not self.region:
        self.select_region()
    
    # 1. retrieve the catalog for subsidiary and region combination
    self.retrieve_catalog_dict()
    self.catalog = GPUCatalog(self.catalog_dict, self.region, subsidiary)
    logger.info(f"Retrieved OVH catalog")

    # 2. select consumption mode
    self.select_consumption_mode(consumption_mode)

    if not gpu_type:
        # 3. list available gpu types for 
        self.catalog.list_gpu_instances(self.consumption_mode)
    
        # 4. select gpu instance type
        gpu_idx = int(input("\nSelect GPU instance number: ")) - 1
        gpu_type = self.catalog.gpu_instances_data[gpu_idx]['Name']
        
    self.gpu_type = gpu_type
    logger.info(f"Selected GPU: {self.gpu_type}, with {self.consumption_mode} consumption mode")
    return self.gpu_type

# %% ../nbs/00_core.ipynb #6a22ad75-0b0e-441c-aee3-976ec595e3a0
@patch
def retrieve_gpu_flavor(self:OVHGPUManager)->str:
    """
    Retrieve the OVH instance flavor associated the selected GPU type
    """
    try:
        flavors = self.client.get(f'{self.project.base_url}/flavor', region=self.region)
    except Exception:
        print(f'Cannot find ressources for {self.region} region. Please select another region')
        return None
    
    flavor = [f for f in flavors if f['name'] == self.gpu_type][0]
    self.gpu_flavor_id = flavor['id']
    logger.info(f"Selected GPU flavor id: {self.gpu_flavor_id}")
    return self.gpu_flavor_id

# %% ../nbs/00_core.ipynb #43d75cdd-f18b-4cbb-925a-5ccda77089a2
@patch
def select_instance_image(self:OVHGPUManager, image_family:str='Ubuntu', image_id:str=None, image_name:str=None)->tuple:
    """
    Lists available instance images and prompts user to select an image if no image ID an is passed as an argument

    Parameters
    ----------
    image_family : str, required, defaults to Ubuntu
        The OS familly to look for. Based on key word search
    image_id : str, optional
        The image id if already nown
    image_name : str, optional
        The image name if already nown
    """
    images = self.client.get(f'{self.project.base_url}/image', region=self.region)
    ubuntu_images = [img for img in images if image_family in img['name']]
    if (not image_id) and (not image_name):
        logger.info("\nAvailable Ubuntu images:")
        for i, img in enumerate(ubuntu_images):
            print(f"{i+1}. {img['name']}")
        image_idx = int(input("Select image number: ")) - 1
        selected_image = ubuntu_images[image_idx]
        image_id = selected_image['id']
        image_name = selected_image['name']
    self.instance_image_id = image_id
    self.instance_image_name = image_name
    logger.info(f"Selected instance image: {self.instance_image_id}:{self.instance_image_name}")
    return self.instance_image_id, self.instance_image_name

# %% ../nbs/00_core.ipynb #06b957a8-a0b8-4d37-a1e3-72a100f5b5b9
@patch
def select_ssh_key(self:OVHGPUManager, ssh_key_id:str=None, ssh_key_name:str=None)->tuple:
    """
    Lists available ssh keys and prompts user to select a ssh key if no ssh key ID is passed as an argument

    Parameters
    ----------
    ssh_key_id : str, optional
        OVH ssh key ID
    ssh_key_name: str, optional
        OVH ssh key name
    """
    ssh_keys = self.client.get(f'{self.project.base_url}/sshkey', region=self.region)
    if not ssh_key_id:
        logger.info("\nAvailable SSH keys:")
        for i, key in enumerate(ssh_keys):
            print(f"{i+1}. {key['name']}")
        ssh_idx = int(input("Select SSH key number: ")) - 1
        ssh_key_id = ssh_keys[ssh_idx]['id']
        ssh_key_name = ssh_keys[ssh_idx]['name']
    self.ssh_key_id = ssh_key_id
    self.ssh_key_name =ssh_key_name
    logger.info(f"Selected SSH key: {self.ssh_key_id}:{self.ssh_key_name}")
    return self.ssh_key_id, self.ssh_key_name

# %% ../nbs/00_core.ipynb #f9a187e4-19db-4e61-bcb3-c098cd0e6cb1
@patch
def create_instance(self:OVHGPUManager, subsidiary='FR', gpu_type=None, consumption_mode=None, image_family:str='Ubuntu',
                    image_id:str=None, image_name:str=None, ssh_key_id:str=None, ssh_key_name:str=None, instance_name=None)->Any:

    """
    Creates an OVH GPU instance either using passed arguments otherwise by prompting the user for inputs

    Parameters
    ----------
    subsidiary : str, required, defaults to FR
        OVH catalpg subsidiary/country code
    gpu_type : str, optional
        GPU instance type
    consumption_mode : str, optional
        Ressource consumption mode
    image_family : str, required, defaults to Ubuntu
        The OS familly to look for. Based on key word search
    image_id : str, optional
        The image id if already nown
    image_name : str, optional
        The image name if already nown
    ssh_key_id : str, optional
        OVH ssh key ID
    ssh_key_name: str, optional
        OVH ssh key name
    instance_name: str, optional
        OVH instance name
    """
    
    self.select_gpu_instance_type(subsidiary, gpu_type, consumption_mode)    
    self.retrieve_gpu_flavor()
    self.select_instance_image(image_family, image_id, image_name)
    self.select_ssh_key(ssh_key_id, ssh_key_name)

    if not instance_name:
        instance_name = input("\nEnter instance name: ")
    self.instance_name = instance_name
    logger.info(f"\nCreating instance '{self.instance_name}'...")
    
    # 5. create instance
    self.last_instance_dict = self.client.post(f'{self.project.base_url}/instance', 
        flavorId=self.gpu_flavor_id,
        imageId=self.instance_image_id,
        name=self.instance_name,
        region=self.region,
        sshKeyId=self.ssh_key_id
    )
    self.last_gpu_instance = GPUInstance(
        client=self.client,
        project_base_url=self.project.base_url,
        region=self.region,
        name=self.last_instance_dict.get('name', 'No Name'),
        status=self.last_instance_dict.get('status', 'No Status'),
        id=self.last_instance_dict.get('id', 'XXXX'),
        gpu_type=self.gpu_type,
        flavor_id=self.gpu_flavor_id,
        image_id=self.instance_image_id,
        consumption_mode=self.consumption_mode,
        created=self.last_instance_dict.get('created', 'No Date')
    )
    return self.last_gpu_instance, self.last_instance_dict

# %% ../nbs/00_core.ipynb #3983ba08-51ba-44a3-9124-b2828a1464f7
@patch
def list_instances(self:OVHGPUManager)->list:
    """
    Lists and displays available OVH GPU instances for specific project
    """
    instances = self.client.get(f'{self.project.base_url}/instance')
    self.gpu_instances = []
    logger.info('\nAvailable GPU instances:')
    for i in instances:
        gpu_type = i.get('planCode').split('.')[0]
        consumption_mode = '.'.join(i.get('planCode').split('.')[1:])
        instance = GPUInstance(
            client=self.client,
            project_base_url=self.project.base_url,
            region=i.get('region'),
            name=i.get('name'),
            status=i.get('status'),
            id=i.get('id'),
            gpu_type=gpu_type,
            flavor_id=i.get('flavorId'),
            image_id=i.get('imageId'),
            consumption_mode=consumption_mode,
            created=i.get('created')
            )
        instance.display_details()
        print(f'------\n')
        self.gpu_instances.append(instance)
        
    return self.gpu_instances

# %% ../nbs/00_core.ipynb #1d6d3367-ba57-41fe-909f-62a909f0d821
@patch
def get_instance(self:OVHGPUManager, instance_id:str)->tuple:
    """
    Retrieves an OVH GPU instance information for a specific instance ID

    Parameters
    ----------
    instance_id : str, required
        OVH GPU instance id
    """
    self.last_instance_dict = self.client.get(f'{self.project.base_url}/instance/{instance_id}')
    if self.last_instance_dict:
        gpu_type = None
        consumption_mode = None
        if self.last_instance_dict.get('planCode'):
            gpu_type = self.last_instance_dict.get('planCode').split('.')[0]
            consumption_mode = '.'.join(self.last_instance_dict.get('planCode').split('.')[1:])
        self.last_gpu_instance = GPUInstance(
                client=self.client,
                project_base_url=self.project.base_url,
                region=self.last_instance_dict.get('region'),
                name=self.last_instance_dict.get('name'),
                status=self.last_instance_dict.get('status'),
                id=self.last_instance_dict.get('id'),
                gpu_type=gpu_type,
                flavor_id=self.last_instance_dict.get('flavor')['id'] if self.last_instance_dict.get('flavor') else None,
                image_id=self.last_instance_dict.get('image')['id'] if self.last_instance_dict.get('image') else None,
                consumption_mode=consumption_mode,
                created=self.last_instance_dict.get('created')
                )
        self.last_gpu_instance.display_details()
    return self.last_gpu_instance ,self.last_instance_dict
