import * as zlib from 'zlib';
import {
  DescribeExecutionCommand,
  ExecutionAlreadyExists,
  ExecutionDoesNotExist,
  SFNClient,
  StartExecutionCommand,
} from '@aws-sdk/client-sfn';
import { SendMessageBatchCommand, SQSClient } from '@aws-sdk/client-sqs';
import * as AWSLambda from 'aws-lambda';
import { labelsMatch, MAX_RUNNER_NAME_LENGTH, OrchestratorInput, WARM_RUNNER_JOB_ID } from './lambda-common';
import { getOctokit, isNotFound, resolveInstallationId } from './lambda-github';
import { controlledJobRunnerLabels, recordRunnerAccounted, RunnerReportMessage } from './lambda-tracker';

const sfn = new SFNClient();
const sqs = new SQSClient();

/**
 * Limit the amount of times we try to replace a stolen runner. Mostly so we don't end up in a replacement loop.
 *
 * @internal
 */
export const MAX_REPLACEMENTS = 3;

/**
 * @internal
 */
export function replacementRunnerName(stolenRunnerName: string): [string, number] {
  const match = stolenRunnerName.match(/^(.*)-r(\d+)$/);
  const base = match ? match[1] : stolenRunnerName;
  const attempt = match ? parseInt(match[2], 10) + 1 : 1;
  const suffix = `-r${attempt}`;

  return [`${base.slice(0, MAX_RUNNER_NAME_LENGTH - suffix.length)}${suffix}`, attempt];
}

function executionArn(runnerName: string): string {
  const stateMachineArn = process.env.STEP_FUNCTION_ARN;
  if (!stateMachineArn) {
    throw new Error('Missing STEP_FUNCTION_ARN environment variable');
  }

  return `${stateMachineArn.replace(':stateMachine:', ':execution:')}:${runnerName}`;
}

async function findRunnerInput(runnerName: string): Promise<OrchestratorInput | undefined> {
  try {
    const execution = await sfn.send(new DescribeExecutionCommand({ executionArn: executionArn(runnerName) }));
    return execution.input ? JSON.parse(execution.input) as OrchestratorInput : undefined;
  } catch (e) {
    if (e instanceof ExecutionDoesNotExist) {
      // somebody else's runner, or one from before this was deployed
      return undefined;
    }
    throw e;
  }
}

/**
 * Get job id from report that only contains runner name, repo, and workflow run id.
 * We list all jobs for the workflow and look for a matching runner name.
 *
 * GitHub doesn't have a better API for this. It also doesn't expose job id to the job.
 */
async function listJobs(owner: string, repo: string, workflowId: number, installationId?: number) {
  const { octokit } = await getOctokit(installationId);

  return octokit.paginate(octokit.rest.actions.listJobsForWorkflowRun, {
    owner,
    repo,
    run_id: workflowId,
    filter: 'all', // all workflow attempts
    per_page: 100,
  });
}

type JobFromReport =
  | { result: 'job'; jobId: number }
  | { result: 'unknown' }
  | { result: 'invisible-repo' };

async function findJobFromReport(report: RunnerReportMessage): Promise<JobFromReport> {
  const [owner, repo] = report.repo.split('/');
  if (!owner || !repo) {
    console.warn({ notice: 'Runner reported a repository we cannot parse', report });
    return { result: 'unknown' };
  }

  let installationId: number | undefined;
  try {
    installationId = await resolveInstallationId(owner, repo);
  } catch (e) {
    if (!isNotFound(e)) {
      throw e;
    }
    // our app isn't installed there, so we never started a runner for anything in that repo. the job that ran there stole the runner from a repo we
    // actually manage. we should replace that runner so the original job can run.
    //
    // this assumption only works for app authentication where the webhook notifies us of jobs on the repos the app is installed on. and in turn the
    // app only has access to, and can only start runners for, the repos it is installed on.
    //
    // PATs have no such promise: they ignore installation ids entirely, so a hand configured webhook works and the two sets can drift apart.
    // that is why this branch is unreachable for them. we can't guarantee (yet) that we will fail to list jobs because the repo truly doesn't exist.
    // a 404 below on listJobs() can mean the job truly doesn't exist, or it can mean the repo exists but the PAT doesn't have access to it while the
    // webhook does report on it.
    return { result: 'invisible-repo' };
  }

  let jobs;
  try {
    jobs = await listJobs(owner, repo, report.workflowId, installationId);
  } catch (e) {
    if (!isNotFound(e)) {
      throw e;
    }
    // we can see this repo and there is no such run, so the runner made it up. ignore.
    console.warn({ notice: 'GitHub knows nothing about the run this runner reported', report });
    return { result: 'unknown' };
  }

  const job = jobs.find(j => j.runner_name === report.runnerName);
  if (!job) {
    // can't find the job based on workflow id + runner name. the runner made the job up. ignore.
    console.warn({ notice: 'No job in this workflow run ran on this runner', report, jobs: jobs.length });
    return { result: 'unknown' };
  }

  return { result: 'job', jobId: job.id };
}

/**
 * Start another runner just like the one that was stolen.
 */
async function replaceRunner(stolenRunnerName: string, input: OrchestratorInput, thiefJobId?: number) {
  const [runnerName, attempt] = replacementRunnerName(stolenRunnerName);
  if (attempt > MAX_REPLACEMENTS) {
    // avoid infinite loops if a runner is stolen repeatedly
    // it might be a bug in our detection...
    console.warn({
      notice: 'Runner was stolen, not replaced because it has already been replaced too many times',
      metric: 'StolenRunnerDetected',
      replaced: 'false', // metric filter drops ALL dimensions for if this is a real boolean
      stolenRunnerName,
      runnerName,
      attempt,
      stolenByJobId: thiefJobId,
      jobId: input.jobId,
      jobUrl: input.jobUrl,
      owner: input.owner,
      repo: input.repo,
      provider: input.provider,
    });
    return;
  }

  try {
    await sfn.send(new StartExecutionCommand({
      stateMachineArn: process.env.STEP_FUNCTION_ARN,
      name: runnerName,
      input: JSON.stringify(input),
    }));
  } catch (e) {
    if (e instanceof ExecutionAlreadyExists) {
      // we already heard about this theft
      console.log({ notice: 'Replacement runner already started', stolenRunnerName, runnerName });
      return;
    }
    throw e;
  }

  console.warn({
    notice: 'Runner was stolen, started another one',
    metric: 'StolenRunnerDetected',
    replaced: 'true', // metric filter drops ALL dimensions for if this is a real boolean
    stolenRunnerName,
    runnerName,
    attempt,
    stolenByJobId: thiefJobId,
    jobId: input.jobId,
    jobUrl: input.jobUrl,
    owner: input.owner,
    repo: input.repo,
    provider: input.provider,
  });
}

/**
 * Check if runner was stolen and possibly replace it.
 */
async function handleRunnerReport(report: RunnerReportMessage) {
  const input = await findRunnerInput(report.runnerName);
  if (!input) {
    return;
  }

  // the runner got a job, so the idle reaper doesn't need to look for a job it left without a runner once it's gone
  await recordRunnerAccounted(report.runnerName, report.jobId);

  if (input.jobId == WARM_RUNNER_JOB_ID) {
    // warm runners aren't started for any job in particular, so nothing they run is technically stolen
    // warm runner keeper starts a replacement on its own as soon as it sees a warm runner go busy
    console.log({
      notice: 'Warm runner taken',
      metric: 'WarmRunnerTaken',
      runnerName: report.runnerName,
      repo: report.repo,
      runId: report.workflowId,
    });
    return;
  }

  let jobId = report.jobId; // can be reported from workflow_job.in_progress directly
  if (jobId === undefined) {
    const job = await findJobFromReport(report);

    if (job.result === 'invisible-repo') {
      // nothing we start could be running there, so this one was stolen
      console.warn({
        notice: 'Runner ran a job in a repository our app cannot see',
        report,
      });
      await replaceRunner(report.runnerName, input);
      return;
    }

    if (job.result === 'unknown') {
      // we couldn't find the job, even though we should have
      // bad input... don't act on it
      console.warn({
        notice: 'Could not tell which job this runner ran',
        report,
      });
      return;
    }

    jobId = job.jobId;
  }

  if (jobId === input.jobId) {
    // ran the job it was started for, which is the whole point
    // no need to do anything or even check the jobs table...
    return;
  }

  const otherRunnerLabels = await controlledJobRunnerLabels(jobId);
  if (otherRunnerLabels) {
    if (labelsMatch(input.jobLabels.split(','), otherRunnerLabels)) {
      // GitHub shuffled jobs between runners we started
      // the runner we started for that job can run this runner's job instead, so no job is left hanging
      console.log({
        notice: 'Runner ran another job of ours',
        runnerName: report.runnerName,
        jobId,
        expectedJobId: input.jobId,
      });
      return;
    }

    // the runner we started for that job can't run this runner's job (e.g. a [self-hosted, lambda] job took a runner started for
    // [self-hosted, lambda, arm64] while its own runner is x64). this runner's job is left without a runner, so treat it as stolen.
    console.log({
      notice: 'Runner ran another job of ours, but that job\'s runner can\'t run this one\'s',
      runnerName: report.runnerName,
      jobId,
      expectedJobId: input.jobId,
      otherRunnerLabels,
      expectedJobLabels: input.jobLabels,
    });
  }

  // runner ran a job we never started a runner for, or one whose runner can't take over
  // runner was stolen, start another one just like it
  await replaceRunner(report.runnerName, input, jobId);
}

/**
 * Send log messages to SQS so we can delay and batch them. This helps make sure we process them *after* the webhook arrives and writes to DDB.
 */
async function handleRunnerLogs(event: AWSLambda.CloudWatchLogsEvent) {
  const queueUrl = process.env.JOB_ASSIGNMENT_QUEUE_URL;
  if (!queueUrl) {
    throw new Error('Missing JOB_ASSIGNMENT_QUEUE_URL environment variable');
  }

  const payload = JSON.parse(zlib.gunzipSync(Buffer.from(event.awslogs.data, 'base64')).toString('utf8')) as AWSLambda.CloudWatchLogsDecodedData;

  const messages: RunnerReportMessage[] = [];
  const reportedRunnersInThisDelivery = new Set<string>();
  for (const logEvent of payload.logEvents) {
    const message = parseRunnerReport(logEvent.message);
    if (message) {
      // don't let misconfigured/malicious jobs make us start or even query too many runners.
      // there should only ever be one report per runner and any malicious code shouldn't be able to guess our runner names.
      // runner names are based on github webhook execution ids, which are random uuids.
      if (reportedRunnersInThisDelivery.has(message.runnerName)) {
        console.warn({ notice: 'Already handled a report from this runner', runnerName: message.runnerName });
        continue;
      }
      reportedRunnersInThisDelivery.add(message.runnerName);
      // put on queue so webhook has time to arrive and write to DDB before we try to read it
      messages.push(message);
    } else {
      console.warn({
        notice: 'Unparsable runner report',
        logGroup: payload.logGroup,
        logStream: payload.logStream,
        message: logEvent.message,
      });
    }
  }

  // batches of ten, which is both the SQS limit and more reports than one log delivery should ever carry
  for (let i = 0; i < messages.length; i += 10) {
    const batch = messages.slice(i, i + 10);
    await sqs.send(new SendMessageBatchCommand({
      QueueUrl: queueUrl,
      Entries: batch.map((message, index) => ({
        Id: `${i + index}`,
        MessageBody: JSON.stringify(message),
        // one report per runner. the webhook sends the same pair, so whichever arrives first wins and the other is
        // dropped by the queue. a group per runner means runners never queue behind each other.
        MessageGroupId: message.runnerName,
        MessageDeduplicationId: message.runnerName,
      })),
    }));
  }

  console.log({
    notice: 'Queued runner reports',
    logGroup: payload.logGroup,
    count: messages.length,
  });
}

/**
 * @internal
 */
export function parseRunnerReport(line: string): RunnerReportMessage | undefined {
  // this data can't be fully trusted as the jobs themselves can run untrusted code... but:
  //   1. the untrusted code has access to just one runner name
  //   2. we only ever act on the first report for a runner (the queue deduplicates on the runner name)
  //   3. the new runner we start will have a specific name and step functions will reject duplicate names (idempotent)
  // so malicious code shouldn't be able to trick us into starting too many runners, or into making us call GitHub
  // over and over until we run out of rate limit for starting runners
  const match = line.match(/CDKGHR JOB RUNNER=(\S+) REPO=(\S+) WORKFLOW_ID=(\d+)/);
  if (!match) {
    return undefined;
  }

  return {
    kind: 'report',
    runnerName: match[1],
    repo: match[2],
    workflowId: parseInt(match[3], 10),
  };
}

function isCloudWatchLogsEvent(event: unknown): event is AWSLambda.CloudWatchLogsEvent {
  return typeof (event as AWSLambda.CloudWatchLogsEvent)?.awslogs?.data === 'string';
}

export async function handler(event: AWSLambda.SQSEvent | AWSLambda.CloudWatchLogsEvent): Promise<AWSLambda.SQSBatchResponse | void> {
  if (isCloudWatchLogsEvent(event)) {
    // let this throw. CloudWatch Logs retries the delivery, and a lost report is a missed stolen runner.
    return handleRunnerLogs(event);
  }

  const result: AWSLambda.SQSBatchResponse = { batchItemFailures: [] };

  for (const record of event.Records) {
    let report: RunnerReportMessage;
    try {
      report = JSON.parse(record.body) as RunnerReportMessage;
    } catch (e) {
      console.error({
        notice: 'Failed to parse message body',
        messageId: record.messageId,
        error: `${e}`,
      });
      continue;
    }

    try {
      await handleRunnerReport(report);
    } catch (e) {
      console.error({
        notice: 'Failed to process runner report',
        messageId: record.messageId,
        report,
        error: `${e}`,
      });
      result.batchItemFailures.push({ itemIdentifier: record.messageId });
    }
  }

  return result;
}