"""
NimbleBox LMAO is our general purpose observability tool for any kind of computation you might have.
"""

import os
import sys
import shlex
import zipfile
from json import dumps, loads
from functools import lru_cache
from typing import Dict, Any, List, Optional, Union
from subprocess import Popen
from dateparser import parse as parse_date
from google.protobuf.timestamp_pb2 import Timestamp

import nbox.utils as U
from nbox.utils import logger
from nbox.auth import secret, AuthConfig
from nbox.relics import Relics
from nbox.jobs import Job, upload_job_folder
from nbox.hyperloop.common.common_pb2 import Resource
from nbox import messages as mpb

# all the sublime -> hyperloop stuff
from nbox.sublime.lmao_rpc_client import (
  AgentDetails, InitRunRequest, ListProjectsRequest,
  ServingLogRequest, ServingLogResponse
)
from nbox.lmao import common

class ExperimentConfig:
  def __init__(
    self,
    run_kwargs: Dict[str, Any],
    git: Dict[str, Any],
    resource: Resource,
    cli_comm: str,
    save_to_relic: bool = True,
    enable_system_monitoring: bool = False,
  ):
    """In an ideal world this would be a protobuf message, but we are not there yet. This contains all the
    things that are stored in the DB and FE can use this to render elements in the UI.

    Args:
      run_kwargs (Dict[str, Any]): All the arguments that the user passed to the `nbx lmao run ...` CLI
      git (Dict[str, Any]): details for the git repo
      resource (Resource): Resource pb object that contains the details of the resource
      cli_comm (str): The CLI command that was used to run the experiment, user can use this to reproduce the experiment
      save_to_relic (bool): If the user wants to save the experiment to the relic by default
      enable_system_monitoring (bool): If the user wants to enable system monitoring
    """
    self.run_kwargs = run_kwargs
    self.git = git
    self.resource = resource
    self.cli_comm = cli_comm
    self.save_to_relic = save_to_relic
    self.enable_system_monitoring = enable_system_monitoring

  def to_dict(self):
    return {
      "run_kwargs": self.run_kwargs,
      "git": self.git,
      "resource": mpb.message_to_dict(self.resource),
      "cli_comm": self.cli_comm,
      "save_to_relic": self.save_to_relic,
      "enable_system_monitoring": self.enable_system_monitoring,
    }

  def to_json(self):
    return dumps(self.to_dict())

  @classmethod
  def from_json(cls, json_str):
    d = loads(json_str)
    # print(d["resource"])
    d["resource"] = Resource(
      cpu = str(d["resource"]["cpu"]),
      memory = str(d["resource"]["memory"]),
      disk_size = str(d["resource"]["disk_size"]),
      gpu = str(d["resource"]["gpu"]),
      gpu_count = str(d["resource"]["gpu_count"]),
      timeout = int(d["resource"]["timeout"]),
      max_retries = int(d["resource"]["max_retries"]),
    )
    return cls(**d)


class LmaoCLI:
  def serving_logs(
    project_id: str,
    serving_id: str,
    key: str,
    after: Optional[str] = "",
    before: Optional[str] = "",
    limit: int = 100,
    f: str = "",
  ):
    stub = common.get_lmao_stub()
    req = ServingLogRequest(
      workspace_id = secret(AuthConfig.workspace_id),
      project_id = project_id,
      serving_id = serving_id,
      key = key,
    )

    # date things
    if after:
      after_parsed = parse_date(after, settings={'TIMEZONE': 'UTC'})
      if after_parsed is None:
        raise ValueError(f"Invalid date format for `after` argument: {after}")
      t = Timestamp()
      t.FromDatetime(after_parsed)
      req.after.CopyFrom(t)
      logger.info(f"Will get logs from after: {after_parsed}")
    if before:
      before_parsed = parse_date(before, settings={'TIMEZONE': 'UTC'})
      if before_parsed is None:
        raise ValueError(f"Invalid date format for `before` argument: {before}")
      t = Timestamp()
      t.FromDatetime(before_parsed)
      req.before.CopyFrom(t)
      logger.info(f"Will get logs from before: {before_parsed}")
    if after and before:
      if after_parsed > before_parsed:
        raise ValueError(f"`after` date cannot be greater than `before` date")

    # limit things
    if limit < 0:
      raise ValueError("`limit` cannot be negative")
    req.limit = int(limit)

    logs: ServingLogResponse = stub.get_serving_log(req)
    logs_json = mpb.message_to_json(
      message = logs,
      including_default_value_fields = False,
      sort_keys = True,
      use_integers_for_enums = False
    )
    if not f:
      print(logs_json)
      return
    logger.info(f"Writing logs to {f}")
    with open(f, "w") as _f:
      _f.write(logs_json)

  logs = {
    "serving": serving_logs,
  }

  def upload(self, **kwargs):
    raise ValueError("`nbx lmao upload` is not changed to `nbx lmao run`")

  def run(
    self,
    init_path: str,
    project_id: str,

    # all the arguments for the git thing
    untracked: bool = False,
    untracked_no_limit: bool = False,

    # all the things for resources
    resource_cpu: str = "100m",
    resource_memory: str = "128Mi",
    resource_disk_size: str = "1Gi",
    resource_gpu: str = "none",
    resource_gpu_count: str = "0",
    resource_timeout: int = 120_000,
    resource_max_retries: int = 2,

    # the following arguments are used for the initialisation of lmao class
    save_to_relic: bool = True,
    enable_system_monitoring: bool = False,

    # the following things are needed for the different modules in the process
    relics_kwargs: Dict[str, Any] = {},

    # eventually to be deprecated
    trigger: bool = False,

    # any other things to pass to the function / class being called
    **run_kwargs,
  ):
    """Upload and register a new run for a NBX-LMAO project.

    Args:
      init_path (str): This can be a path to a `folder` or can be optionally of the structure `fp:fn` where `fp` is the path to the file and `fn` is the function name.
      project_id (str): The project ID or name. If the project name is provided, then a new project will be created.
      trigger (bool, optional): Defaults to False. If True, will trigger the run after uploading.
      untracked (bool, optional): If True, then untracked files below 1MB will also be zipped and uploaded. Defaults to False.
      untracked_no_limit (bool, optional): If True, then all untracked files will also be zipped and uploaded. Defaults to False.
      resource_cpu (str, optional): Defaults to "100m". The CPU resource to allocate to the run.
      resource_memory (str, optional): Defaults to "128Mi". The memory resource to allocate to the run.
      resource_disk_size (str, optional): Defaults to "1Gi". The disk size resource to allocate to the run.
      resource_gpu (str, optional): Defaults to "none". The GPU resource to use.
      resource_gpu_count (str, optional): Defaults to "0". Number of GPUs allocated to the experiment.
      resource_timeout (int, optional): Defaults to 120_000. The timeout between two consecutive runs, honoured but not guaranteed.
      resource_max_retries (int, optional): Defaults to 2. The maximum number of retries for a run.
      save_to_relic (bool, optional): Defaults to True. If True, will save the files to a relic when lmao.save_file() is called.
      enable_system_monitoring (bool, optional): Defaults to False. If True, will enable system monitoring.
      **run_kwargs: These are the kwargs that will be passed to your function.
    """
    workspace_id = secret(AuthConfig.workspace_id)

    # fix data type conversion caused by the CLI
    resource_cpu = str(resource_cpu)
    resource_memory = str(resource_memory)
    resource_disk_size = str(resource_disk_size)
    resource_gpu = str(resource_gpu)
    resource_gpu_count = str(resource_gpu_count)
    resource_timeout = int(resource_timeout)
    resource_max_retries = int(resource_max_retries)

    # reconstruct the entire CLI command so we can show it in the UI
    reconstructed_cli_comm = (
      f"nbx lmao upload '{init_path}' '{project_id}'"
      f" --resource_cpu '{resource_cpu}'"
      f" --resource_memory '{resource_memory}'"
      f" --resource_disk_size '{resource_disk_size}'"
      f" --resource_gpu '{resource_gpu}'"
      f" --resource_gpu_count '{resource_gpu_count}'"
      f" --resource_timeout {resource_timeout}"
      f" --resource_max_retries {resource_max_retries}"
    )
    if trigger:
      logger.error("--trigger is now deprecated and will be removed in a few releases")
    if untracked:
      reconstructed_cli_comm += " --untracked"
    if untracked_no_limit:
      reconstructed_cli_comm += " --untracked_no_limit"
    if save_to_relic:
      reconstructed_cli_comm += " --save_to_relic"
    if enable_system_monitoring:
      reconstructed_cli_comm += " --enable_system_monitoring"
    if resource_max_retries < 1:
      logger.error(f"max_retries must be >= 1. Got: {resource_max_retries}\n  Fix: set --max_retries=2")
      raise ValueError()

    # uncommenting these since it can contain sensitive information, kept just in case for taking dictionaries as input
    # if relics_kwargs:
    #   rks = str(relics_kwargs).replace("'", "")
    #   reconstructed_cli_comm += f" --relics_kwargs '{rks}'"

    for k, v in run_kwargs.items():
      if type(v) == bool:
        reconstructed_cli_comm += f" --{k}"
      else:
        reconstructed_cli_comm += f" --{k} '{v}'"
    logger.debug(f"command: {reconstructed_cli_comm}")

    if untracked_no_limit and not untracked:
      logger.debug("untracked_no_limit is True but untracked is False. Setting untracked to True")
      untracked = True

    # first step is to get all the relevant information from the DB and servers
    workspace_id = workspace_id or secret(AuthConfig.workspace_id)
    # project_name, project_id = get_project_name_id(project_name_or_id, workspace_id)
    job_name = "nbxj_lmao_" + project_id
    try:
      job = Job(job_name = job_name)
    except:
      logger.warn(f"Job not found. Creating a new job: {job_name}")
      job = Job.new(job_name = job_name, description = "automatically created by nbx LMAO for project: " + project_id)

    logger.info(f"Project: {project_id}")
    lmao_stub = common.get_lmao_stub()

    # create a call init run and get the experiment metadata
    init_folder, _ = os.path.split(init_path)
    init_folder = init_folder or "."
    if os.path.exists(U.join(init_folder, ".git")):
      git_det = common.get_git_details(init_folder)
    else:
      git_det = {}

    run = lmao_stub.init_run(InitRunRequest(
      workspace_id = workspace_id,
      project_id = project_id,
      agent_details = AgentDetails(
        type = AgentDetails.NBX.JOB,
        nbx_job_id = job.id, # run id will have to be updated from the job
      ),
      config = ExperimentConfig(
        run_kwargs = run_kwargs,
        git = git_det,
        resource = Resource(
          cpu = resource_cpu,
          memory = resource_memory,
          disk_size = resource_disk_size,
          gpu = resource_gpu,
          gpu_count = resource_gpu_count,
          timeout = resource_timeout,
          max_retries = resource_max_retries,
        ),
        cli_comm = reconstructed_cli_comm,
        save_to_relic = save_to_relic,
        enable_system_monitoring = enable_system_monitoring,
      ).to_json(),
    ))
    logger.info(f"Run ID: {run.experiment_id}")

    # connect to the relic
    r_keys = set(relics_kwargs.keys())
    valid_keys = {"bucket_name", "region", "nbx_resource_id", "nbx_integration_token"}
    extra_keys = r_keys - valid_keys
    if extra_keys:
      logger.error("Unknown arguments found:\n  * " + "\n  * ".join(extra_keys))
      raise RuntimeError("Unknown arguments found in the Relic")
    relic = Relics(LMAO_RELIC_NAME, workspace_id = workspace_id, create = True, **relics_kwargs)

    # create a git patch and upload it to relics
    if git_det:
      _zf = U.join(init_folder, "untracked.zip")
      zip_file = zipfile.ZipFile(_zf, "w", zipfile.ZIP_DEFLATED)
      patch_file = U.join(init_folder, "nbx_auto_patch.diff")
      f = open(patch_file, "w")
      Popen(shlex.split(f"git diff {' '.join(git_det['uncommited_files'])}"), stdout=f, stderr=sys.stderr).wait()
      f.close()
      zip_file.write(patch_file, arcname = patch_file)
      if untracked:
        untracked_files = git_det["untracked_files"] # what to do with these?
        if untracked_files:
          # see if any file is larger than 10MB and if so, warn the user
          warn_once = False
          for f in untracked_files:
            if os.path.getsize(f) > 1e7 and not untracked_no_limit:
              logger.warning(f"File: {f} is larger than 10MB and will not be available in sync")
              logger.warning("  Fix: use git to track small files, avoid large files")
              logger.warning("  Fix: nbox.Relics can be used to store large files")
              logger.warning("  Fix: use --untracked-no-limit to upload all files")
              warn_once = True
              continue
            zip_file.write(f, arcname = f)
      relic.put_to(_zf, f"{project_name}/{run.experiment_id}/git.zip")
      os.remove(_zf)
      os.remove(patch_file)

    # tell the server that this run is being scheduled so atleast the information is visible on the dashboard
    upload_job_folder(
      "job",
      init_folder = init_path,
      id = job.id,

      # pass along the resource requirements
      resource_cpu = resource_cpu,
      resource_memory = resource_memory,
      resource_disk_size = resource_disk_size,
      resource_gpu = resource_gpu,
      resource_gpu_count = resource_gpu_count,
      resource_timeout = resource_timeout,
      resource_max_retries = resource_max_retries,
    )

    # now trigger the experiment
    fp = f"{run.project_id}/{run.experiment_id}"
    tag = f"{LMAO_RM_PREFIX}{fp}"
    logger.debug(f"Running job '{job.name}' ({job.id}) with tag: {tag}")
    job.trigger(tag)

    # finally print the location of the run where the users can track this
    logger.info(f"Run location: {secret(AuthConfig.url)}/workspace/{run.workspace_id}/monitoring/{run.project_id}/{run.experiment_id}")


# do not change these it can become a huge pain later on
LMAO_RELIC_NAME = "experiments"
LMAO_RM_PREFIX = "NBXLmao-"
LMAO_ENV_VAR_PREFIX = "NBX_LMAO_"
