"""Interaction with Flock.io's API"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_api.ipynb.

# %% auto 0
__all__ = ['get_task', 'open_tasks', 'finalized_tasks', 'get_subs', 'sub_stats', 'FlockWalletStats', 'upload_to_hf',
           'submit_task']

# %% ../nbs/01_api.ipynb 3
import json
import pandas as pd
from fastcore.all import *
from huggingface_hub import HfApi

from .core import BASE_URL, get_gpu_info, pd_parse

# %% ../nbs/01_api.ipynb 7
def get_task(task_id: int, as_df: bool = False): 
    """ Get task as JSON or Dataframe. """
    j = urljson(f"{BASE_URL}tasks/get?task_id={task_id}")
    return pd_parse(j) if as_df else j

# %% ../nbs/01_api.ipynb 13
def open_tasks(as_df: bool = False):
    j = urljson(f'{BASE_URL}tasks/list?status=submission_phase')
    return pd_parse(j) if as_df else j

# %% ../nbs/01_api.ipynb 16
def finalized_tasks(as_df: bool = False): 
    j = urljson(f'{BASE_URL}tasks/list?status=finalized')
    return pd_parse(j) if as_df else j

# %% ../nbs/01_api.ipynb 21
def get_subs(task_id: int, as_df: bool = False):
    """ Get submissions as JSON or Dataframe. """
    j = urlread(f'{BASE_URL}tasks/wallet-submissions?task_id={task_id}', 
                   headers={"flock-api-key":os.getenv('FLOCK_API_KEY')},
                   return_json=True)
    if as_df:
        cols = ['link','submission_phase_score', 'final_validation_score', 'submitted_at']
        subdf1 = pd.DataFrame.from_dict(j['items'])
        return subdf1[cols].sort_values(by='submission_phase_score', ascending=False).drop_duplicates('link', keep='first')
    else:
        return j

# %% ../nbs/01_api.ipynb 28
def sub_stats(as_df: bool):
    """ Submissions and validation counts per day. """
    j = urljson(f'{BASE_URL}stats/emissions?type=submission')['data']
    return pd_parse(j) if as_df else j

# %% ../nbs/01_api.ipynb 38
class FlockWalletStats(BasicRepr):
    """ Get wallet-based stats. """
    def __init__(self, address: str):
        self.address = address
        self.base_url = f'{BASE_URL}stats/'

    def subs(self, page: int = 1, size: int = 50, as_df: bool = True):
        """ Get submissions of a wallet. """
        assert page >= 1, "Page must be at least 1."
        assert 1 <= size <= 100, "Size must be between 1 and 100."
        j = urljson(f'{self.base_url}submissions/{self.address}&page={page}&size={size}')
        return pd_parse(j) if as_df else j
    
    def achievements(self, role: str = 'validator', as_df: bool = True):
        """ Get validation achievements of a wallet. """
        assert role in ['validator', 'node'], "Role must be either 'validator' or 'node'."
        j = urljson(f'{self.base_url}achievements/{self.address}?role={role}')
        return pd_parse(j) if as_df else j
    
    def enrolled(self, as_df: bool = True):
        """ Get enrolled tasks of a wallet. """
        j = urljson(f'{self.base_url}stats/enrolled-tasks/{self.address}')
        return pd_parse(j) if as_df else j

# %% ../nbs/01_api.ipynb 53
def upload_to_hf(repo_name: str, output_dir: str = "outputs"):
    """ 
    Upload a model to the HuggingFace Hub. 
    :param repo_name: the name of the repository to upload to. For example: `myhfname/task-1-model`.
    :param output_dir: the directory to upload to the HuggingFace Hub.
    :return: the commit hash you will need to submit to Flock.
    """
    api = HfApi(token=os.environ["HF_TOKEN"])
    api.create_repo(
        repo_name,
        exist_ok=False,
        repo_type="model",
        )
    commit_message = api.upload_folder(
                folder_path=output_dir,
                repo_id=repo_name,
                repo_type="model",
            )
    return commit_message.oid

# %% ../nbs/01_api.ipynb 56
def submit_task(
    task_id: int, hg_repo_id: str, base_model: str, commit_hash: str
):
    """ 
    Submit results to Flock task. 
    :param task_id: the ID of the task you want to submit to.
    :param hg_repo_id: the ID of the HuggingFace repository you want to submit.
    :param base_model: the base model you want to use. Must be a valid model name like "Qwen/Qwen1.5-7B or google/gemma-7b.
    :param commit_hash: the HuggingFace Hub Model commit hash of the model you want to submit.
    """
    payload = json.dumps(
        {
            "task_id": task_id,
            "data": {
                "hg_repo_id": hg_repo_id,
                "base_model": base_model,
                "gpu_type": get_gpu_info()['gpu_0'].name,
                "revision": commit_hash,
            },
        }
    )
    hdr = {
        "flock-api-key": os.getenv('FLOCK_API_KEY'),
        "Content-Type": "application/json",
    }
    return urlrequest(f"{BASE_URL}/tasks/submit-result", 
                      verb="POST",
                      headers=hdr, 
                      data=payload,
                      json_data=True)
