import os
import sys
import shlex
import zipfile
from pprint import pformat
from subprocess import Popen

from nbox.utils import logger
from nbox import utils as U
from nbox.auth import secret, AuthConfig
from nbox.init import nbox_ws_v1
from nbox.lmao import Lmao, LmaoLive, get_lmao_stub, lmao_v2_pb2 as pb, ExperimentConfig, LMAO_RM_PREFIX
from nbox.lmao import common
from nbox.relics import Relics
from nbox.sublime.lmao_rpc_client import LMAO_Stub
from nbox.jobs import Job
from nbox.hyperloop.common.common_pb2 import Resource

class ProjectConfig:
  project_id: str = ""


class Project:
  def __init__(self, id: str = ""):
    # https://app.c.nimblebox.ai/api/v1/workspace/zcxdpqlk/projects/3909cfb8
    id = id or ProjectConfig.project_id
    if not id:
      raise ValueError("Project ID not provided")
    logger.info(f"Connecting to Project: {id}")
    self.stub = nbox_ws_v1.projects.u(id)
    self.data = self.stub()
    self.pid = self.data["project_id"]
    self.lmao_stub = get_lmao_stub()
    self.project = self.lmao_stub.get_project(pb.Project(
      workspace_id = secret.workspace_id,
      project_id = self.pid
    ))
    self.relic = Relics(id = self.project.relic_id)
    self.workspace_id = secret.workspace_id

  def __repr__(self) -> str:
    return pformat(self.data, indent = 2)

  def put_settings(self, project_name: str = "", project_description: str = ""):
    self.stub(
      "put",
      project_name = project_name or self.data["project_name"],
      project_description = project_description or self.data["project_description"]
    )
    self.data["project_name"] = project_name or self.data["project_name"]
    self.data["project_description"] = project_description or self.data["project_description"]

  def get(self, k):
    return self.data.get(k)
  
  def get_lmao_stub(self) -> LMAO_Stub:
    return self.lmao_stub

  def get_exp_tracker(self, experiment_id: str = "") -> Lmao:
    lmao = Lmao(project_id = self.pid, experiment_id = experiment_id)
    return lmao
  
  def get_live_tracker(self) -> LmaoLive:
    lmao = LmaoLive(self.pid)
    return lmao
  
  def get_relic(self) -> Relics:
    return self.relic
  
  def get_job_id(self) -> str:
    return self.data["job_list"][0]
  
  def get_deployment_id(self) -> str:
    return self.data["deployment_list"][0]
  
  # the big CLI command that replaces nbx lmao run

  def run(
    self,
    init_path: 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_max_retries: int = 2,

    # 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_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}' '{self.pid}'"
      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_max_retries {resource_max_retries}"
    )
    if untracked_no_limit and not untracked:
      logger.debug("untracked_no_limit is True but untracked is False. Setting untracked to True")
      untracked = True
    if untracked:
      reconstructed_cli_comm += " --untracked"
    if untracked_no_limit:
      reconstructed_cli_comm += " --untracked_no_limit"
    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()
    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}")

    # get the job
    job = Job(job_id = self.get_job_id())

    # 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 = self.lmao_stub.init_run(pb.InitRunRequest(
      workspace_id = workspace_id,
      project_id = self.pid,
      agent_details = pb.AgentDetails(
        type = pb.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,
          max_retries = resource_max_retries,
        ),
        cli_comm = reconstructed_cli_comm,
        save_to_relic = True,
        enable_system_monitoring = False,
      ).to_json(),
    ))
    logger.info(f"Run ID: {run.experiment_id}")


    # 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) 
      
      # upload the patch file to relics
      self.relic.put_to(_zf, f"{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
    Job.upload(
      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_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}/projects/{run.project_id}#Experiments")
