"""
This module will contain all the sdk functions for the iris command sdk, including login, logout, get, post, pull.
"""
import ast
import hashlib
import json
import logging
import os
from pathlib import Path
from typing import List, Optional, Union
from urllib.parse import urljoin

# for iris infer
import numpy as np

# for iris pull
import requests
import tritonclient.http
from rich import print

# internal imports
from .auth_utils import auth, handle_bad_response
from .conf_manager import conf_mgr
from .exception import (
    ArtefactNotFoundError,
    ArtefactTypeInvalidError,
    BadRequestError,
    InvalidCommandError,
    InvalidDatasetFormatError,
    MissingTokenizerError,
    UnsafeTensorsError,
)
from .utils import download_model, dump, pull_image, tarify, upload_from_file

# ───────────────────────────────────────────────────── imports ────────────────────────────────────────────────────── #


# Whether to use tqdm for progress bars
TQDM = True

# ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────── #
#                                                  IRIS USERS SDK                                                     #
# ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────── #


# ------------------------------------      Setup Logger      ------------------------------------ #
# Logger config
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
)
logger = logging.getLogger(__name__)
logger.setLevel(conf_mgr.LOG_LEVEL)


# --------------------------------------      iris login    -------------------------------------- #
@auth
def login():
    logger.debug("logging in")
    return conf_mgr.current_user["name"]


# --------------------------------------     iris logout    -------------------------------------- #


def logout():
    logger.info("logging out")
    path = Path.home() / Path(conf_mgr.config.keyfile_name)
    if path.exists():
        path.unlink()
    if not path.exists():
        return True
    else:
        return False


# --------------------------------------      iris post     -------------------------------------- #


@auth
def post(headers: dict = {}, **flags):
    endpoint = "experiment"
    # detype the flags, so we can send them
    payload = {k: str(val) for k, val in flags.items()}
    logger.debug(f"Dispatching job with payload {payload}")
    url = urljoin(conf_mgr.runner_url, f"{endpoint}/")
    headers.update({"Authorization": f"Bearer {conf_mgr.access_token}"})
    response = requests.post(url=url, headers=headers, data=payload)
    if not response.ok:
        raise handle_bad_response(response, endpoint)
    else:
        print(dump(response))


# --------------------------------------       iris get     -------------------------------------- #


@auth
def get(
    object: str = "experiment",
    id: Optional[str] = None,
    query: Optional[str] = None,
    headers: dict = {},
    verbose=True,
):
    """Get objects from the titan API

    Args:
        object (str, optional): The object to get. Defaults to "experiment".
        id (Optional[str], optional): The id of the object. Defaults to None.
        query (Optional[str], optional): A JMESPath query to run against the returned json. Defaults to None, i.e. return everything.
        headers (dict): Custom headers to send with the get request
    Returns:
        (str) A json response, formatted as: {"status": <http_response>, "response": <queried_json_response>}
    """
    logger.debug(f"Getting from ... {conf_mgr.base}, auth from {conf_mgr.AUTH0_DOMAIN}")
    logger.debug(f"Applying custom headers {headers}")
    endpoint = object + "/" + (str(id) if id is not None else "")
    url = urljoin(conf_mgr.runner_url, endpoint)
    headers.update({"Authorization": f"Bearer {conf_mgr.access_token}"})

    try:
        response = requests.get(url=url, headers=headers)
    except requests.exceptions.ConnectionError as e:
        print(
            json.dumps(
                {
                    "status": "failed",
                    "message": f"Connection error: Error reaching {url}",
                },
                indent=4,
            )
        )
        raise Exception("Connection error")
    except requests.exceptions.MissingSchema as e:
        print(
            json.dumps(
                {
                    "status": "failed",
                    "message": f"Scheme error: iris target specified improperly. base={conf_mgr.base}, runner_url: {conf_mgr.runner_url}, endpoint={endpoint} result={url} ",
                },
                indent=4,
            )
        )
        raise Exception("Scheme error")
    except Exception as e:
        print(e)
        raise e

    if not response.ok:
        raise handle_bad_response(response, url)
    else:
        dumped_response = dump(response, query)
        if verbose:
            print(dumped_response)
        return dumped_response


@auth
def delete(
    object: str,
    id: Optional[str],
    verbose=True,
):
    """Get objects from the titan API

    Args:
        object (str, optional): The object to get. Defaults to "experiment".
        id (str): The id of the object to delete
    Returns:
        (str) A json response, formatted as: {"status": <http_response>, "response": <queried_json_response>}
    """
    logger.debug(f"Getting from ... {conf_mgr.base}, auth from {conf_mgr.AUTH0_DOMAIN}")
    endpoint = object + "/" + (str(id) if id is not None else "")
    url = urljoin(conf_mgr.runner_url, endpoint)
    headers = {"Authorization": f"Bearer {conf_mgr.access_token}"}

    response = requests.delete(url=url, headers=headers)
    if not response.ok:
        raise handle_bad_response(response, endpoint)
    else:
        if response.status_code == 204:
            return {"status": "success"}

        dumped_response = dump(response)
        if verbose:
            print(dumped_response)
        return dumped_response


# --------------------------------------     iris download  -------------------------------------- #


@auth
def download(experiment_cmd: str):
    """Downloading the models to local machine

    Args:
        experiment_cmd (str): pulling command string. it should be formatted as <experiment_id>:<job_tag>

    Raises:
        InvalidCommandError: Invalid command error
        BadRequestError: Bad request error
        ArtefactNotFoundError: Artefact not found error

    Returns:
        model_name: name of the model
        task_name: name of the task
        baseline_model_name: name of the baseline model
        baseline: whether the model is baseline or not
    """

    # create a model_storage folder if it doesn't exist
    Path("model_storage").mkdir(parents=True, exist_ok=True)

    # parse the command string
    args = experiment_cmd.split(":")
    if len(args) != 2:
        raise InvalidCommandError
    experiment_id = args[0]
    job_tag = args[1]

    # get the experiment info
    endpoint = "experiment"
    url = urljoin(conf_mgr.runner_url, f"{endpoint}/{experiment_id}")
    headers = {"Authorization": f"Bearer {conf_mgr.access_token}"}

    response = requests.get(url=url, headers=headers)
    response_json = response.json()

    # check if the request is successful
    if response_json["status"] != "success":
        raise BadRequestError

    jobs_list = response_json["experiment"]["jobs"]

    model_uuid = None
    baseline = False
    download_url, model_name, task_name, baseline_model_name = None, None, None, None
    # loop through the jobs list and find the job with the same tag
    for i in range(len(jobs_list)):
        if job_tag == jobs_list[i]["name"].split("_")[-1]:
            model_uuid = jobs_list[i]["out_art_id"]
            model_name = jobs_list[i]["name"]
            task_name = jobs_list[i]["flags"]["task"]
            baseline_model_name = jobs_list[i]["flags"]["model.teacher"]

            if baseline_model_name == "null":
                baseline_model_name = jobs_list[i]["flags"]["model.student"]
                baseline = True

            if task_name == "sequence_classification":
                text_fields = ast.literal_eval(
                    jobs_list[i]["flags"]["data.text_fields"]
                )  # convert string to list
                if len(text_fields) == 1:
                    task_name = "sequence_classification"
                elif len(text_fields) == 2:
                    task_name = "pair_classification"
                else:
                    raise ValueError("Invalid number of text fields")

            if task_name == "glue":
                dataset_name = jobs_list[i]["flags"][
                    "data.dataset_config_name"
                ]  # convert string to list
                if dataset_name in {"sst2", "cola"}:
                    task_name = "sequence_classification"  # cola and sst2 only has one sentence as input
                else:
                    task_name = "pair_classification"

            if model_uuid is None:
                raise ArtefactNotFoundError
            url = urljoin(
                conf_mgr.runner_url,
                f"artefact/link/{model_uuid}/download?refresh=true",
            )
            response = requests.get(url=url, headers=headers)
            response_json = response.json()
            download_url = response_json["link"]["link"]
            break

    # download the model
    if download_url is not None:
        download_model(download_url, model_name)
    print(f"\nModel {model_name} downloaded successfully")
    return model_name, task_name, baseline_model_name, baseline


# --------------------------------------      iris pull     -------------------------------------- #


@auth
def pull(experiment_cmd: str):
    """download the model and pull the hephaestus image from the server

    Args:
        experiment_cmd (str): pulling command string. it should be formatted as <experiment_id>:<job_tag>

    """

    # parse the command string
    args = experiment_cmd.split(":")
    if len(args) != 2:
        raise InvalidCommandError
    experiment_id = args[0]
    job_tag = args[1]

    # download the model
    logger.info("***Executing pull command***")
    model_name, task_name, baseline_model_name, baseline = download.__wrapped__(
        experiment_cmd
    )  # this is a hack to get the function to work without the auth decorator

    # pull the image
    logger.info("Pulling image from the server")
    pull_image(
        model_folder_path=f"model_storage/{model_name}/models",
        container_name=f"iris-triton-{experiment_id}",
        job_tag=job_tag,
        task_name=task_name,
        baseline_model_name=baseline_model_name,
        baseline=baseline,
    )
    logger.info("All done!")


# --------------------------------------    iris upload     -------------------------------------- #


@auth
def upload(name: str, src: Union[str, Path], art_type: str, description: str):
    """Upload an artefact to the titan hub

    Args:
        name (str): The name of the artefact
        src (Union[str, Path]): The source of the artefact on disk
        art_type (str): The artefact type. Should be either ['model', 'dataset']
        description (str): A short description of the artefact.

    Raises:
        ArtefactNotFoundError: If the path to the artefact doesn't exist.
        UnsafeTensorsError: If the artefact is a model, and the model has not been saved in safetensors format.
    """
    # cast from path to str.
    art_type = str(art_type)
    src = str(src)

    endpoint = "artefact"
    url = urljoin(conf_mgr.runner_url, f"{endpoint}/")
    headers = {"Authorization": f"Bearer {conf_mgr.access_token}"}
    logger.debug(f"Uploading {art_type} from {src} to {url}")

    # Catches if you accidentally put a tilde in quotes:
    if src[0] == "~":
        src = os.path.expanduser(src)

    if not Path(src).is_dir():
        raise ArtefactNotFoundError(details=src)

    ext = ".tar.gz"
    if art_type == "model":
        # Model upload validation
        namelist = os.listdir(src)
        if not any(".safetensors" in Path(x).suffixes for x in namelist):
            raise UnsafeTensorsError()

        if not any("tokenizer" in x for x in namelist):
            raise MissingTokenizerError()

    elif art_type == "dataset":
        # Dataset validation
        namelist = os.listdir(src)
        if not (
            any("val" in x and x.endswith(".csv") for x in namelist)
            and any("train" in x and x.endswith(".csv") for x in namelist)
        ):
            raise InvalidDatasetFormatError
    else:
        raise ArtefactTypeInvalidError

    # Make post request to seshat for instantiation of artefact object, and provision of presigned upload link.
    post_req_data = {
        "name": name,
        "artefact_type": art_type,
        "description": description,
        "ext": ext,
        "src": src,
    }
    logger.debug(f"posting {post_req_data} to {url}")
    post_req_response = requests.post(url=url, headers=headers, data=post_req_data)
    if not post_req_response.ok:
        logger.debug("post unsuccessful")
        raise handle_bad_response(post_req_response, endpoint)
    else:
        # On succesful response receipt:
        logger.debug("post successful")
        data = post_req_response.json()["artefact"]
        art_uuid = data["uuid"]
        endpoint = "artefact/link/" + art_uuid + "/upload"
        url = urljoin(conf_mgr.runner_url, f"{endpoint}")
        logger.debug(f"Getting link from {url}")
        get_link_resp = requests.get(url=url, headers=headers)
        upl_link = get_link_resp.json()["link"]["link"]
        logger.debug(f"got link {upl_link}")

        # todo this is where further format checking would go (e.g. no uploading 10GB files or jpegs...)
        print("Beginning upload...")

    # Upload file with TQDM or without, whilst also hashing the file.
    if TQDM:
        hashval, upload_response = upload_from_file(src, upl_link)
    else:
        with tarify(src, mode="w:gz") as f:
            f.seek(0)
            upload_response = requests.put(upl_link, data=f)
            hashval = hashlib.md5(f.getbuffer()).hexdigest()

    # If Upload completes succesfully, send the hash to seshat in a patch request. Patch req confirms matching hash,
    # then returns the artefact ID for further use by user.
    if upload_response is not None and upload_response.status_code == 200:
        endpoint = "artefact"
        print(f"Upload Complete -  Validating {art_type} UUID: {art_uuid} ......")
        url = urljoin(conf_mgr.runner_url, f"{endpoint}/{art_uuid}")
        patch_req_response = requests.patch(
            url=url, headers=headers, data={"hashval": hashval}
        )
        if not patch_req_response.ok:
            raise handle_bad_response(patch_req_response, endpoint)
        else:
            print("Upload validated")
    else:
        print("Upload failed")
        print(dump(upload_response))


# --------------------------------------      iris infer    -------------------------------------- #


def infer(
    url: str,
    task_name: str,
    use_cpu: bool,
    text: List[str],
    context: Optional[str] = None,
):
    """infer the hosted tytn model using triton client

    Args:
        port (int): port number of the triton server
        task_name (str): task name of the model (sequence_classification, question_answering)
        use_cpu (bool): whether to use cpu or not
        text (str): input text field used for inference
        context (Optional[str]): input context field used for inference. Only used in question_answering task. Defaults to None.
    Returns:
        infer_res(dict): dictionary of the inference result
    """

    model_name = (
        "transformer_onnx_inference" if use_cpu else "transformer_tensorrt_inference"
    )
    model_version = "1"
    batch_size = 1

    triton_client = tritonclient.http.InferenceServerClient(url=url, verbose=False)

    model_metadata = triton_client.get_model_metadata(
        model_name=model_name, model_version=model_version
    )
    model_config = triton_client.get_model_config(
        model_name=model_name, model_version=model_version
    )
    text_1, text_2 = None, None
    if task_name == "sequence_classification":
        if len(text) == 1:
            # if only one text is provided, this is a single sentence classification task inference
            text_1 = tritonclient.http.InferInput(
                name="TEXT", shape=(batch_size,), datatype="BYTES"
            )
            text_1.set_data_from_numpy(np.asarray([text[0]] * batch_size, dtype=object))
        elif len(text) == 2:
            # if two texts are provided, this is a pair sentence classification task inference
            text_1 = tritonclient.http.InferInput(
                name="TEXT_1", shape=(batch_size,), datatype="BYTES"
            )
            text_2 = tritonclient.http.InferInput(
                name="TEXT_2", shape=(batch_size,), datatype="BYTES"
            )
            text_1.set_data_from_numpy(np.asarray([text[0]] * batch_size, dtype=object))
            text_2.set_data_from_numpy(np.asarray([text[1]] * batch_size, dtype=object))
        else:
            raise ValueError(
                "Invalid number of texts provided for sequence classification task"
            )
    elif task_name == "question_answering":
        if context is None:
            raise ValueError("Context must be provided for question answering task")
        text_1 = tritonclient.http.InferInput(
            name="QUESTION", shape=(batch_size,), datatype="BYTES"
        )
        text_2 = tritonclient.http.InferInput(
            name="CONTEXT", shape=(batch_size,), datatype="BYTES"
        )
        text_1.set_data_from_numpy(np.asarray([text[0]] * batch_size, dtype=object))
        text_2.set_data_from_numpy(np.asarray([context] * batch_size, dtype=object))
    else:
        raise ValueError("Invalid task name provided")

    # bind outputs to the server
    outputs = tritonclient.http.InferRequestedOutput(name="outputs", binary_data=False)

    infer_res = triton_client.infer(
        model_name=model_name,
        model_version=model_version,
        inputs=[text_1, text_2] if text_2 is not None else [text_1],
        outputs=[outputs],
    )
    return json.loads(infer_res.as_numpy("outputs").item())
