#!/usr/bin/env python
"""API handlers for dealing with cron jobs."""

from typing import Optional

from grr_response_core.lib import rdfvalue
from grr_response_core.lib import registry
from grr_response_core.lib.rdfvalues import mig_protodict
from grr_response_core.lib.rdfvalues import structs as rdf_structs
from grr_response_proto import flows_pb2
from grr_response_proto.api import cron_pb2
from grr_response_server import cronjobs
from grr_response_server.databases import db
from grr_response_server.gui import api_call_context
from grr_response_server.gui import api_call_handler_base
from grr_response_server.gui import api_call_handler_utils
from grr_response_server.gui import mig_api_call_handler_utils
from grr_response_server.gui.api_plugins import flow as api_plugins_flow
from grr_response_server.models import protobuf_utils
from grr_response_server.rdfvalues import cronjobs as rdf_cronjobs
from grr_response_server.rdfvalues import hunts as rdf_hunts
from grr_response_server.rdfvalues import mig_cronjobs
from grr_response_server.rdfvalues import objects as rdf_objects


def InitApiCronJobFromCronJob(
    cron_job: flows_pb2.CronJob,
) -> cron_pb2.ApiCronJob:
  """Initializes ApiCronJob from CronJob."""

  api_cron_job = cron_pb2.ApiCronJob()
  protobuf_utils.CopyAttr(cron_job, api_cron_job, "cron_job_id")
  if cron_job.HasField("args"):
    api_cron_job.args.CopyFrom(cron_job.args)
  protobuf_utils.CopyAttr(cron_job, api_cron_job, "current_run_id")
  protobuf_utils.CopyAttr(cron_job, api_cron_job, "description")
  protobuf_utils.CopyAttr(cron_job, api_cron_job, "enabled")
  protobuf_utils.CopyAttr(cron_job, api_cron_job, "last_run_status")
  protobuf_utils.CopyAttr(cron_job, api_cron_job, "last_run_time")
  protobuf_utils.CopyAttr(cron_job, api_cron_job, "frequency")
  protobuf_utils.CopyAttr(cron_job, api_cron_job, "lifetime")
  protobuf_utils.CopyAttr(cron_job, api_cron_job, "allow_overruns")

  api_cron_job.is_failing = cron_job.last_run_status in [
      flows_pb2.CronJobRun.CronJobRunStatus.ERROR,
      flows_pb2.CronJobRun.CronJobRunStatus.LIFETIME_EXCEEDED,
  ]

  if cron_job.forced_run_requested:
    api_cron_job.forced_run_requested = True

  rdf_state = mig_protodict.ToRDFAttributedDict(cron_job.state)
  state_dict = rdf_state.ToDict()
  if state_dict:
    state = api_call_handler_utils.ApiDataObject()
    state.InitFromDataObject(state_dict)
    api_cron_job.state.CopyFrom(
        mig_api_call_handler_utils.ToProtoApiDataObject(state)
    )
  return api_cron_job


class CronJobNotFoundError(api_call_handler_base.ResourceNotFoundError):
  """Raised when a cron job could not be found."""


class CronJobRunNotFoundError(api_call_handler_base.ResourceNotFoundError):
  """Raised when a cron job run could not be found."""


class ApiCronJobId(rdfvalue.RDFString):
  """Class encapsulating cron job ids."""


class ApiCronJobRunId(rdfvalue.RDFString):
  """Class encapsulating cron job run ids."""


class ApiCronJob(rdf_structs.RDFProtoStruct):
  """ApiCronJob is used when rendering responses.

  ApiCronJob is meant to be more lightweight than automatically generated AFF4
  representation. It's also meant to contain only the information needed by
  the UI and to not expose implementation details.
  """

  protobuf = cron_pb2.ApiCronJob
  rdf_deps = [
      ApiCronJobId,
      api_call_handler_utils.ApiDataObject,
      rdf_cronjobs.CronJobAction,
      rdfvalue.DurationSeconds,
      rdfvalue.RDFDatetime,
  ]

  def GetArgsClass(self):
    if self.flow_name:
      flow_cls = registry.FlowRegistry.FlowClassByName(self.flow_name)

      # The required protobuf for this class is in args_type.
      return flow_cls.args_type

  def _IsCronJobFailing(self, cron_job):
    """Returns True if the last run failed."""
    status = cron_job.Get(cron_job.Schema.LAST_RUN_STATUS)
    if status is None:
      return False

    return status.status != rdf_cronjobs.CronJobRunStatus.Status.OK

  status_map = {
      rdf_cronjobs.CronJobRunStatus.Status.OK: (
          rdf_cronjobs.CronJobRun.CronJobRunStatus.FINISHED
      ),
      rdf_cronjobs.CronJobRunStatus.Status.ERROR: (
          rdf_cronjobs.CronJobRun.CronJobRunStatus.ERROR
      ),
      rdf_cronjobs.CronJobRunStatus.Status.TIMEOUT: (
          rdf_cronjobs.CronJobRun.CronJobRunStatus.LIFETIME_EXCEEDED
      ),
  }

  def _StatusFromCronJobRunStatus(self, status):
    if status is None:
      return None

    return self.status_map[status.status]

  def ObjectReference(self):
    return rdf_objects.ObjectReference(
        reference_type=rdf_objects.ObjectReference.Type.CRON_JOB,
        cron_job=rdf_objects.CronJobReference(
            cron_job_id=str(self.cron_job_id)
        ),
    )


class ApiCronJobRun(rdf_structs.RDFProtoStruct):
  """ApiCronJobRun represents individual cron job runs."""

  protobuf = cron_pb2.ApiCronJobRun
  rdf_deps = [
      ApiCronJobId,
      ApiCronJobRunId,
      rdfvalue.RDFDatetime,
  ]

  def InitFromRunObject(self, run):
    self.run_id = run.run_id
    self.cron_job_id = run.cron_job_id
    self.started_at = run.started_at
    self.finished_at = run.finished_at
    self.status = run.status
    self.log_message = run.log_message or None
    self.backtrace = run.backtrace or None
    return self

  def InitFromApiFlow(self, f, cron_job_id=None):
    """Shortcut method for easy legacy cron jobs support."""
    if f.flow_id:
      self.run_id = f.flow_id
    elif f.urn:
      self.run_id = f.urn.Basename()
    self.started_at = f.started_at
    self.cron_job_id = cron_job_id

    flow_state_enum = api_plugins_flow.ApiFlow.State
    cron_enum = rdf_cronjobs.CronJobRun.CronJobRunStatus
    errors_map = {
        flow_state_enum.RUNNING: cron_enum.RUNNING,
        flow_state_enum.TERMINATED: cron_enum.FINISHED,
        flow_state_enum.ERROR: cron_enum.ERROR,
        flow_state_enum.CLIENT_CRASHED: cron_enum.ERROR,
    }
    self.status = errors_map[f.state]

    if f.state != f.State.RUNNING:
      self.finished_at = f.last_active_at

    if f.context.kill_timestamp:
      self.status = self.Status.LIFETIME_EXCEEDED

    if f.context.HasField("status"):
      self.log_message = f.context.status

    if f.context.HasField("backtrace"):
      self.backtrace = f.context.backtrace

    return self


class ApiListCronJobsArgs(rdf_structs.RDFProtoStruct):
  protobuf = cron_pb2.ApiListCronJobsArgs


class ApiListCronJobsResult(rdf_structs.RDFProtoStruct):
  protobuf = cron_pb2.ApiListCronJobsResult
  rdf_deps = [
      ApiCronJob,
  ]


class ApiListCronJobsHandler(api_call_handler_base.ApiCallHandler):
  """Lists flows launched on a given client."""

  args_type = ApiListCronJobsArgs
  result_type = ApiListCronJobsResult
  proto_args_type = cron_pb2.ApiListCronJobsArgs
  proto_result_type = cron_pb2.ApiListCronJobsResult

  def Handle(
      self,
      args: cron_pb2.ApiListCronJobsArgs,
      context: Optional[api_call_context.ApiCallContext] = None,
  ) -> cron_pb2.ApiListCronJobsResult:
    if not args.count:
      stop = None
    else:
      stop = args.offset + args.count

    cron_manager = cronjobs.CronManager()
    all_jobs = list(cron_manager.ReadJobs())
    all_jobs.sort(
        key=lambda job: (getattr(job, "cron_job_id", None) or job.urn)
    )
    all_jobs = [mig_cronjobs.ToProtoCronJob(job) for job in all_jobs]
    cron_jobs = all_jobs[args.offset : stop]

    items = [InitApiCronJobFromCronJob(cron_job) for cron_job in cron_jobs]

    return cron_pb2.ApiListCronJobsResult(
        items=items, total_count=len(all_jobs)
    )


class ApiGetCronJobArgs(rdf_structs.RDFProtoStruct):
  protobuf = cron_pb2.ApiGetCronJobArgs
  rdf_deps = [
      ApiCronJobId,
  ]


class ApiGetCronJobHandler(api_call_handler_base.ApiCallHandler):
  """Retrieves a specific cron job."""

  args_type = ApiGetCronJobArgs
  result_type = ApiCronJob
  proto_args_type = cron_pb2.ApiGetCronJobArgs
  proto_result_type = cron_pb2.ApiCronJob

  def Handle(
      self,
      args: cron_pb2.ApiGetCronJobArgs,
      context: Optional[api_call_context.ApiCallContext] = None,
  ) -> cron_pb2.ApiCronJob:
    try:
      cron_job = cronjobs.CronManager().ReadJob(str(args.cron_job_id))
      cron_job = mig_cronjobs.ToProtoCronJob(cron_job)
      return InitApiCronJobFromCronJob(cron_job)
    except db.UnknownCronJobError as e:
      raise CronJobNotFoundError(
          "Cron job with id %s could not be found" % args.cron_job_id
      ) from e


class ApiListCronJobRunsArgs(rdf_structs.RDFProtoStruct):
  protobuf = cron_pb2.ApiListCronJobRunsArgs
  rdf_deps = [
      ApiCronJobId,
  ]


class ApiListCronJobRunsResult(rdf_structs.RDFProtoStruct):
  protobuf = cron_pb2.ApiListCronJobRunsResult
  rdf_deps = [
      ApiCronJobRun,
  ]


class ApiListCronJobRunsHandler(api_call_handler_base.ApiCallHandler):
  """Retrieves the given cron job's runs."""

  args_type = ApiListCronJobRunsArgs
  result_type = ApiListCronJobRunsResult

  def Handle(self, args, context=None):
    runs = cronjobs.CronManager().ReadJobRuns(str(args.cron_job_id))
    start = args.offset
    if args.count:
      end = args.offset + args.count
    else:
      end = db.MAX_COUNT
    return ApiListCronJobRunsResult(
        items=[
            ApiCronJobRun().InitFromRunObject(run) for run in runs[start:end]
        ]
    )


class ApiGetCronJobRunArgs(rdf_structs.RDFProtoStruct):
  protobuf = cron_pb2.ApiGetCronJobRunArgs
  rdf_deps = [
      ApiCronJobId,
      ApiCronJobRunId,
  ]


class ApiGetCronJobRunHandler(api_call_handler_base.ApiCallHandler):
  """Renders given cron run."""

  args_type = ApiGetCronJobRunArgs
  result_type = ApiCronJobRun

  def Handle(self, args, context=None):
    run = cronjobs.CronManager().ReadJobRun(
        str(args.cron_job_id), str(args.run_id)
    )
    if not run:
      raise CronJobRunNotFoundError(
          "Cron job run with id %s could not be found" % args.run_id
      )

    return ApiCronJobRun().InitFromRunObject(run)


class ApiCreateCronJobArgs(rdf_structs.RDFProtoStruct):
  """Arguments for CreateCronJob API call."""

  protobuf = cron_pb2.ApiCreateCronJobArgs
  rdf_deps = [
      rdfvalue.DurationSeconds,
      rdf_hunts.HuntRunnerArgs,
  ]

  def GetFlowArgsClass(self):
    if self.flow_name:
      flow_cls = registry.FlowRegistry.FlowClassByName(self.flow_name)

      # The required protobuf for this class is in args_type.
      return flow_cls.args_type


class ApiCreateCronJobHandler(api_call_handler_base.ApiCallHandler):
  """Creates a new cron job."""

  args_type = ApiCreateCronJobArgs
  result_type = ApiCronJob
  proto_args_type = cron_pb2.ApiCreateCronJobArgs
  proto_result_type = cron_pb2.ApiCronJob

  def Handle(
      self,
      source_args: cron_pb2.ApiCreateCronJobArgs,
      context: Optional[api_call_context.ApiCallContext] = None,
  ) -> cron_pb2.ApiCronJob:
    # Make sure we don't modify source arguments.
    args = cron_pb2.ApiCreateCronJobArgs()
    args.CopyFrom(source_args)

    # Clear all fields marked with HIDDEN.
    args = ToRDFApiCreateCronJobArgs(args)
    args.flow_args.ClearFieldsWithLabel(
        rdf_structs.SemanticDescriptor.Labels.HIDDEN
    )
    # Clear all fields marked with HIDDEN, except for output_plugins - they are
    # marked HIDDEN, because we have a separate UI for them, not because they
    # shouldn't be shown to the user at all.
    #
    # TODO(user): Refactor the code to remove the HIDDEN label from
    # FlowRunnerArgs.output_plugins.
    args.hunt_runner_args.ClearFieldsWithLabel(
        rdf_structs.SemanticDescriptor.Labels.HIDDEN,
        exceptions="output_plugins",
    )

    cron_manager = cronjobs.CronManager()
    cron_args = rdf_cronjobs.CreateCronJobArgs.FromApiCreateCronJobArgs(args)
    cron_job_id = cron_manager.CreateJob(cron_args=cron_args, enabled=False)

    cron_obj = cron_manager.ReadJob(cron_job_id)
    cron_obj = mig_cronjobs.ToProtoCronJob(cron_obj)
    return InitApiCronJobFromCronJob(cron_obj)


class ApiForceRunCronJobArgs(rdf_structs.RDFProtoStruct):
  protobuf = cron_pb2.ApiForceRunCronJobArgs
  rdf_deps = [
      ApiCronJobId,
  ]


class ApiForceRunCronJobHandler(api_call_handler_base.ApiCallHandler):
  """Force-runs a given cron job."""

  args_type = ApiForceRunCronJobArgs

  def Handle(self, args, context=None):
    job_id = str(args.cron_job_id)
    cronjobs.CronManager().RequestForcedRun(job_id)


class ApiModifyCronJobArgs(rdf_structs.RDFProtoStruct):
  protobuf = cron_pb2.ApiModifyCronJobArgs
  rdf_deps = [
      ApiCronJobId,
  ]


class ApiModifyCronJobHandler(api_call_handler_base.ApiCallHandler):
  """Enables or disables a given cron job."""

  args_type = ApiModifyCronJobArgs
  result_type = ApiCronJob
  proto_args_type = cron_pb2.ApiModifyCronJobArgs
  proto_result_type = cron_pb2.ApiCronJob

  def Handle(
      self,
      args: cron_pb2.ApiModifyCronJobArgs,
      context: Optional[api_call_context.ApiCallContext] = None,
  ) -> cron_pb2.ApiCronJob:
    cron_id = args.cron_job_id
    if args.enabled:
      cronjobs.CronManager().EnableJob(cron_id)
    else:
      cronjobs.CronManager().DisableJob(cron_id)

    cron_job_obj = cronjobs.CronManager().ReadJob(cron_id)
    cron_job_obj = mig_cronjobs.ToProtoCronJob(cron_job_obj)
    return InitApiCronJobFromCronJob(cron_job_obj)


class ApiDeleteCronJobArgs(rdf_structs.RDFProtoStruct):
  protobuf = cron_pb2.ApiDeleteCronJobArgs
  rdf_deps = [
      ApiCronJobId,
  ]


class ApiDeleteCronJobHandler(api_call_handler_base.ApiCallHandler):
  """Deletes a given cron job."""

  args_type = ApiDeleteCronJobArgs

  def Handle(self, args, context=None):
    cronjobs.CronManager().DeleteJob(str(args.cron_job_id))


# Copy of migration function to avoid circular dependency.
def ToRDFApiCreateCronJobArgs(
    proto: cron_pb2.ApiCreateCronJobArgs,
) -> ApiCreateCronJobArgs:
  return ApiCreateCronJobArgs.FromSerializedBytes(proto.SerializeToString())
