"""
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 (
  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,
  }

# 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_"
