import os

from json import dumps
from typing import Dict, Any, List, Optional, Union
from nbox.utils import logger, env as nbx_env
from nbox.utils import get_files_in_folder
from nbox import Instance
from nbox.auth import secret
from nbox.nbxlib.tracer import Tracer
from requests import Session

import dainik.utils as U
from dainik.utils import env
from dainik.lmao_client import *
from dainik.relic import RelicStore, get_relic_file
from dainik.proto.lmao_pb2 import AgentDetails
from dainik.git_utils import get_git_details
from dainik.metrics.system import SystemMetricsLogger

DEBUG_LOG_EVERY = 5
INFO_LOG_EVERY = 100

class Dainik():
  def __init__(self, workspace_id: str, *, local: bool = False) -> None:
    self.workspace_id = workspace_id
    self.local = local

    self._initialized = False
    self.nbx_job_folder = nbx_env.NBOX_JOB_FOLDER("")
    self._total_logged_elements = 0 # this variable keeps track for logging
    self.completed = False
    self.tracer = None
    self.username = secret.get("username")
    self.relic = None

    if not local:
      self._create_connection(workspace_id)
    else:
      self.lmao = LMAO_Stub(url = "http://127.0.0.1:8080", session = Session())

  def _create_connection(self, workspace_id: str):
    # prepare the URL
    # id_or_name, workspace_id = U.split_iw(instance_id)
    id_or_name = f"monitoring-{workspace_id}"
    logger.info(f"id_or_name: {id_or_name}")
    logger.info(f"workspace_id: {workspace_id}")
    instance = Instance(id_or_name, workspace_id)
    try:
      open_data = instance.open_data
    except AttributeError:
      raise Exception(f"Is instance '{instance.project_id}' running?")
    build = "build"
    if "app.c." in secret.get("nbx_url"):
      build = "build.c"
    url = f"https://server-{open_data['url']}.{build}.nimblebox.ai/"
    logger.info(f"URL: {url}")

    # create a tracer object that will load all the information
    self.tracer = Tracer(start_heartbeat=False)
    self.username = self.tracer.job_proto.auth_info.username if self.tracer else self.username

    # create a session with the auth header
    _session = Session()
    _session.headers.update({
      "NBX-TOKEN": open_data["token"],
      "X-NBX-USERNAME": self.username,
    })

    # define the stub
    self.lmao = LMAO_Stub(url = url, session = _session)

  def init(self, project_name: Optional[str] = "", config: Dict[str, Any] = {}, project_id: Optional[str] = ""):
    if project_id:
      project_name = ""
      if project_name:
        logger.info("ignoring `project_name` as `project_id` was given")
    else:
      project_id = ""
      if not project_name:
        raise Exception("Provide one from `project_id` or `project_name`")

    # this is the config value that is used to store data on the plaform, user cannot be allowed to have
    # like a full access to config values
    log_config: Dict[str, Any] = {
      "user_config": config
    }

    # check if the current folder from where this code is being executed has a .git folder
    # NOTE: in case of NBX-Jobs the current folder ("./") is expected to contain git by default
    log_config["git"] = None
    if os.path.exists(".git"):
      log_config["git"] = get_git_details("./")

    # continue as before
    self.project_name = project_name
    self.project_id = project_id
    self.config = config
    self._agent_details = AgentDetails(
      nbx_job_id = self.tracer.job_id if self.tracer else "jj_guvernr",
      nbx_run_id = self.tracer.token if self.tracer else U.get_random_id(),
    )
    run_details = self.lmao.init_run(
      _InitRunRequest = InitRunRequest(
        agent_details=self._agent_details,
        created_at = U.get_timestamp(),
        project_name = self.project_name,
        project_id = self.project_id,
        config = dumps(log_config),
      )
    )

    if not run_details:
      # TODO: Make a custom exception of this
      raise Exception("Server Side exception has occurred, Check the log for details")

    # TODO: change f-string here when refactoring `run_id` to `exp_id`
    logger.info(f"Assigned run_id: {run_details.run_id}")
    self.run = run_details
    self._initialized = True

    # system metrics monitoring, by default is enabled optionally turn it off
    if not env.DK_DISABLE_SYSTEM_METRICS(False):
      SystemMetricsLogger(self)

    # now initialize the relic
    if not env.DK_DISABLE_RELICS(False):
      self.relic = RelicStore(self.workspace_id, relic_name = self.run.run_id, create = True)

  def log(self, y: Dict[str, Union[int, float, str]], step = None, *, log_type: str = RunLog.LogType.USER):
    if not self._initialized:
      raise Exception("Run not initialized, call init() first")
    if self.completed:
      raise Exception("Run already completed, cannot log more data!")

    if self._total_logged_elements % DEBUG_LOG_EVERY == 0:
      logger.debug(f"Logging: {y.keys()} | {self._total_logged_elements}")

    if self._total_logged_elements % INFO_LOG_EVERY == 0:
      logger.info(f"Logging: {y.keys()} | {self._total_logged_elements}")

    step = step if step is not None else U.get_timestamp()
    if step < 0:
      raise Exception("Step must be <= 0")
    run_log = RunLog(run_id = self.run.run_id, log_type=log_type)
    for k,v in y.items():
      record = U.get_record(k, v)
      record.step = step
      run_log.data.append(record)

    ack = self.lmao.on_log(_RunLog = run_log)
    if not ack.success:
      logger.error("  >> Server Error")
      for l in ack.message.splitlines():
        logger.error("  " + l)
      raise Exception("Server Error")

    self._total_logged_elements += 1

  def save_file(self, *files):
    """
    Register a file save and upload it by talking to the NimbleBox Relics Backend.
    `dk.save_file("foo.t", "/bar/", "baz.t", "/bar/roo/")`
    """
    logger.info(f"Saving files: {files}")
    if not self._initialized:
      raise Exception("Run not initialized, call init() first")

    # manage all the complexity of getting the list of RelicFile
    all_relic_files = []
    for folder_or_file in files:
      if os.path.isfile(folder_or_file):
        all_files = [folder_or_file]
      elif os.path.isdir(folder_or_file):
        all_files = get_files_in_folder(folder_or_file, self.username, self.workspace_id)
      else:
        raise Exception(f"File or Folder not found: {folder_or_file}")

      out = []
      for f in all_files:
        relic_file = get_relic_file(f, self.username, self.workspace_id)
        relic_file.relic_name = self.run.run_id
        out.append(relic_file)
      all_relic_files.extend((f, x) for f,x in zip(all_files, out))

    # define API function
    def _register_upload(file_list):
      ack = self.lmao.on_save(_FileList = file_list)
      if not ack.success:
        logger.error("  >> Server Error")
        for l in ack.message.splitlines():
          logger.error("  " + l)
        raise Exception("Server Error")
    
    if self.relic:
      for local, relic_file in all_relic_files:
        # upload and then register
        self.relic.upload_relic_file(local, relic_file)
        file_list = FileList(run_id=self.run.run_id)
        file_list.files.append(File(relic_file = relic_file))
        _register_upload(file_list)
    else:
      all_relics = all_relic_files
      file_list = FileList(run_id=self.run.run_id)
      for local, relic_file in all_relics:
        file_list.files.append(File(relic_file = relic_file))
        self.relic.upload_relic_file(local, relic_file)
      _register_upload(file_list)


  def end(self):
    logger.info("Ending run")
    ack = self.lmao.on_train_end(_Run = Run(run_id=self.run.run_id,))
    if not ack.success:
      logger.error("  >> Server Error")
      for l in ack.message.splitlines():
        logger.error("  " + l)
      raise Exception("Server Error")
    self.completed = True
