import json
import time
import sys
import os
import requests
from datetime import datetime

from s3transfer import TransferConfig

from .config import Config
from .httputils import post2, get2, delete2, put2
from .s3utils import s3Client, keys, UploadProgress

from ..utils import log_and_raise


def new_project(
    sim_dict,
    task_name=None,
    folder_name="default",
    solver_version=Config.SOLVER_VERSION,
):
    """Create new project.

    Parameters
    ----------
    sim_dict : dict
        Dictionary containing all simulation parameters, which can be
        generated by :meth:`.Simulation.export`.
    task_name : str, optional
        Custom name for the task.
    folder_name: str, optional
        folder to hold the task.
    solver_version : str, optional
        Version of the Tidy3D solver.

    Returns
    -------
    dict
        Project dictionary with simulation and task data.
    """
    if task_name == None:
        task_name = f"Task_{int(datetime.now().timestamp())}"

    # Initialize the task
    try:
        task = post2(
            f"fdtd/model/{folder_name}/task",
            {
                "status": "draft",
                "solverVersion": solver_version,
                "taskName": task_name,
                "nodeSize": int(sim_dict["parameters"]["nodes"]),
                "timeSteps": int(sim_dict["parameters"]["time_steps"]),
                "computeWeight": float(
                    sim_dict["parameters"]["compute_weight"]
                ),
            },
        )
    except requests.exceptions.HTTPError as e:
        error_json = json.loads(e.response.text)
        log_and_raise(error_json["error"], RuntimeError)

    # upload the file to s3
    print("Uploading the json file...")
    resp = s3Client.put_object(
        Body=json.dumps(sim_dict, indent=4),
        Bucket=Config.STUDIO_BUCKET,
        Key="users/{0}/{1}/{2}".format(
            keys["UserId"], task["taskId"], "fdtd3d.json"
        ),
    )
    task["status"] = "queued"
    return put2(f'fdtd/model/{folder_name}/task/{task["taskId"]}', task)


def get_project(task_id):
    """Get all project details from a given task_id.

    Parameters
    ----------
    task_id : str
        Task identification string.

    Returns
    -------
    dict
        Project dictionary containing task data.
    """

    return get2(f"fdtd/task/{task_id}")


def get_projects():
    """Get a list with all details of all projects of the current user.

    Returns
    -------
    list
        A list of ``dict``, with one entry for every folder name. Each of
        those can contain multiple projects whose corresponding data is
        listed in the ``'children'`` entry.
    """

    return get2(f"fdtd/models")


def delete_project(task_id):
    """Delete a project from a given task_id.

    Parameters
    ----------
    task_id : str
        Task identification string.

    Returns
    -------
    dict
        Project dictionary of the deleted project.
    """

    return delete2(f"fdtd/task/{task_id}")


def delete_queued():
    """Delete all queued projects.

    Returns
    -------
    int
        Total number of projects deleted.
    """

    projects = get_projects()
    count = 0
    for pfolder in projects:
        for task in pfolder["children"]:
            if task["status"] == "queued":
                delete_project(task["taskId"])
                count += 1

    return count

def delete_old(days_old=100, folder=None):
    """Delete all projects older than a given amount of days.
    
    Parameters
    ----------
    days_old : int, optional
        Minimum number of days since the project creation.
    folder : str or None, optional
        If None, all folders are purged.
    
    Returns
    -------
    int
        Total number of projects deleted.
    """

    projects = get_projects()
    count = 0
    for pfolder in projects:
        if pfolder['name'] == folder or folder is None:
            for task in pfolder["children"]:
                stime_str = task["submitTime"]
                stime = datetime.strptime(stime_str, "%Y:%m:%d:%H:%M:%S")
                days_elapsed = (datetime.utcnow() - stime).days
                if days_elapsed >= days_old:
                    delete_project(task["taskId"])
                    count += 1

    return count


def monitor_project(task_id):
    """Monitor the status of a given project every second. Function exits when
    the exit status is either ``'success'`` or ``'error'``.

    Parameters
    ----------
    task_id : str
        Task identification string.
    """
    status = ""
    while status not in ["success", "error"]:
        project = get_project(task_id)
        status = project["status"]
        name = project["taskName"]
        sys.stdout.write('\rProject "' + name + '" status: %s       ' % status)
        sys.stdout.flush()
        time.sleep(0.5)
        sys.stdout.write('\rProject "' + name + '" status: %s.      ' % status)
        sys.stdout.flush()
        time.sleep(0.5)
        sys.stdout.write('\rProject "' + name + '" status: %s..     ' % status)
        sys.stdout.flush()
        time.sleep(0.5)
        sys.stdout.write('\rProject "' + name + '" status: %s...    ' % status)
        sys.stdout.flush()
        time.sleep(0.5)
    print("\n")


def get_last_projects(Nprojects=None):
    """Get a list with information about the last ``Nprojects`` submitted
    projects.

    Returns
    -------
    list
        A list of ``dict``, with one entry for every project.
    """

    projects = get_projects()
    store_dict = {
        "submit_time": [],
        "status": [],
        "task_name": [],
        "task_id": [],
    }
    for pfolder in projects:
        for task in pfolder["children"]:
            try:
                store_dict["submit_time"].append(task["submitTime"])
                store_dict["status"].append(task["status"])
                store_dict["task_name"].append(task["taskName"])
                store_dict["task_id"].append(task["taskId"])
            except:
                pass

    sort_inds = sorted(
        range(len(store_dict["submit_time"])),
        key=store_dict["submit_time"].__getitem__,
        reverse=True,
    )

    if Nprojects is None or Nprojects > len(sort_inds):
        Nprojects = len(sort_inds)

    out_dict = []
    for ipr in range(Nprojects):
        out_dict.append(
            {key: item[sort_inds[ipr]] for (key, item) in store_dict.items()}
        )

    return out_dict


def list_projects(Nprojects=None):
    """Print a summary of all projects of the current user, in chronological
    order of submission.

    Parameters
    ----------
    Nprojects : None, optional
        If supplied, only the last ``Nprojects`` will be listed.
    """

    projects_info = get_last_projects(Nprojects)

    for pr_info in projects_info:
        print("Project name: %s" % pr_info["task_name"])
        print("Task ID     : %s" % pr_info["task_id"])
        print("Submit time : %s" % pr_info["submit_time"])
        print("Status      : %s" % pr_info["status"])
        print("-" * 50)


def download_results_file(task_id, src, target=None):
    """Download all the results recorded by simulation monitors.

    Parameters
    ----------
    task_id : str
        Task ID of the project (after a successful run).
    src : str
        File to download. Valid file names are ``'monitor_data.hdf5'`` and
        ``'tidy3d.log'``.
    target : None, optional
        Filename to store locally. If ``None``, the same name as ``src``
        will be used.
    """
    if target is None:
        target = src
    if src is None:
        log_and_raise("'src' cannot be None.", ValueError)

    try:
        project = get_project(task_id)
    except:
        log_and_raise("Unable to get project " + task_id, RuntimeError)

    try:
        s3Client.download_file(
            Bucket=Config.STUDIO_BUCKET,
            Filename=target,
            Key="users/{0}/{1}/output/{2}".format(
                keys["UserId"], task_id, src
            ),
        )
    except:
        log_and_raise(
            "Cannot retrieve requested file, check the file name and "
            "make sure the project has run correctly. Current "
            f"project status is '{project['status']}'.",
            RuntimeError
        )
        
def download_results(task_id, target_folder=""):
    """Download the results of a solver run, including the json file defining
    the simulation, a single file containing all the monitor data, and a log
    file.

    Parameters
    ----------
    task_id : str
        Task ID of the project (after a successful run).
    target_folder : str, optional
        Folder in which to download the files. Default is working folder.
    """

    if target_folder != "" and not os.path.exists(target_folder):
        os.makedirs(target_folder)

    for fname in ["monitor_data.hdf5", "tidy3d.log"]:
        download_results_file(
            task_id, fname, target=os.path.join(target_folder, fname)
        )

    download_json(task_id, target_folder=target_folder)


def download_json(task_id, target_folder=""):
    """Download the project json file

    Parameters
    ----------
    task_id : str
        Task ID of the project (after a successful run).
    target_folder : str, optional
        Folder in which to download the files. Default is working folder.
    """

    if target_folder != "" and not os.path.exists(target_folder):
        os.makedirs(target_folder)

    project = get_project(task_id)
    sim_dict = project["taskParam"]
    if not sim_dict:
        s3Client.download_file(
            Bucket=Config.STUDIO_BUCKET,
            Filename=os.path.join(target_folder, "simulation.json"),
            Key="users/{0}/{1}/{2}".format(
                keys["UserId"], task_id, "fdtd3d.json"
            ),
        )
    else:
        with open(
            os.path.join(target_folder, "simulation.json"), "w"
        ) as sim_file:
            json.dump(json.loads(), sim_file, indent=4)
