#include "typecheck.h"

#include <math.h>
#include <stdio.h>
#include <string.h>

typedef struct {
  ValueType type;
  double number;
  bool known;
} StackValue;

typedef struct {
  StackValue values[1024];
  size_t count;
} TypeStack;

static bool set_error(Error* error, int line, const char* message) {
  error->line = line;
  snprintf(error->message, sizeof(error->message), "%s", message);
  return false;
}

static bool need(const TypeStack* stack, size_t count, Error* error, int line) {
  return stack->count >= count || set_error(error, line, "stack underflow");
}

static bool push(TypeStack* stack, StackValue value, Error* error, int line) {
  if (stack->count == sizeof(stack->values) / sizeof(stack->values[0])) {
    return set_error(error, line, "type stack limit exceeded");
  }
  stack->values[stack->count++] = value;
  return true;
}

static bool expect(const TypeStack* stack, size_t from_top, ValueType expected,
                   Error* error, int line) {
  StackValue actual = stack->values[stack->count - 1 - from_top];
  if (actual.type == expected) {
    return true;
  }
  error->line = line;
  snprintf(error->message, sizeof(error->message), "expected %s, got %s",
           type_name(expected), type_name(actual.type));
  return false;
}

static void record(TypeInfo* info, const StackValue* consumed, size_t consumed_count,
                   const StackValue* produced, size_t produced_count) {
  info->consumed_count = (unsigned)consumed_count;
  info->produced_count = (unsigned)produced_count;
  for (size_t i = 0; i < consumed_count; i++) info->consumed[i] = consumed[i].type;
  for (size_t i = 0; i < produced_count; i++) info->produced[i] = produced[i].type;
}

static bool replace(TypeStack* stack, size_t consumed_count, const StackValue* produced,
                    size_t produced_count, TypeInfo* info, Error* error, int line) {
  StackValue consumed[3];
  size_t start = stack->count - consumed_count;
  memcpy(consumed, stack->values + start, consumed_count * sizeof(*consumed));
  stack->count = start;
  for (size_t i = 0; i < produced_count; i++) {
    if (!push(stack, produced[i], error, line)) return false;
  }
  record(info, consumed, consumed_count, produced, produced_count);
  return true;
}

static bool same_stack_types(const TypeStack* left, const TypeStack* right) {
  if (left->count != right->count) return false;
  for (size_t i = 0; i < left->count; i++) {
    if (left->values[i].type != right->values[i].type) return false;
  }
  return true;
}

static bool check_stack_operation(Operation operation, TypeStack* stack, TypeInfo* info,
                                  Error* error, int line) {
  size_t input_count = operation == OP_ROT ? 3 : operation == OP_DUP || operation == OP_DROP ? 1 : 2;
  if (!need(stack, input_count, error, line)) return false;
  StackValue* input = stack->values + stack->count - input_count;
  StackValue output[3];
  size_t output_count = 0;
  switch (operation) {
    case OP_DUP: output[0] = input[0]; output[1] = input[0]; output_count = 2; break;
    case OP_DROP: break;
    case OP_SWAP: output[0] = input[1]; output[1] = input[0]; output_count = 2; break;
    case OP_OVER:
      output[0] = input[0]; output[1] = input[1]; output[2] = input[0]; output_count = 3;
      break;
    case OP_ROT:
      output[0] = input[1]; output[1] = input[2]; output[2] = input[0]; output_count = 3;
      break;
    default: return false;
  }
  return replace(stack, input_count, output, output_count, info, error, line);
}

static bool check_numeric(Operation operation, TypeStack* stack, TypeInfo* info,
                          Error* error, int line) {
  if (!need(stack, 2, error, line)) return false;
  StackValue left = stack->values[stack->count - 2];
  StackValue right = stack->values[stack->count - 1];
  StackValue output = {.known = left.known && right.known};

  if (operation == OP_ADD || operation == OP_SUB) {
    if (left.type != right.type || (left.type != TYPE_SCALAR && left.type != TYPE_DURATION)) {
      return set_error(error, line, "add/sub expect matching SCALAR or DURATION operands");
    }
    output.type = left.type;
    output.number = operation == OP_ADD ? left.number + right.number : left.number - right.number;
  } else if (operation == OP_MUL || operation == OP_DIV) {
    bool scalars = left.type == TYPE_SCALAR && right.type == TYPE_SCALAR;
    bool duration_scalar = left.type == TYPE_DURATION && right.type == TYPE_SCALAR;
    if (!scalars && !duration_scalar) {
      return set_error(error, line, "mul/div expect SCALAR SCALAR or DURATION SCALAR");
    }
    if (operation == OP_DIV && right.known && right.number == 0) {
      return set_error(error, line, "division by zero");
    }
    output.type = left.type;
    output.number = operation == OP_MUL ? left.number * right.number : left.number / right.number;
  } else {
    bool ordered = operation != OP_EQ && operation != OP_NE;
    bool comparable = left.type == right.type &&
                      (left.type == TYPE_SCALAR || left.type == TYPE_DURATION ||
                       (!ordered && left.type == TYPE_BOOL));
    if (!comparable) {
      return set_error(error, line, "comparison expects matching SCALAR, DURATION, or BOOL operands");
    }
    output.type = TYPE_BOOL;
    if (operation == OP_EQ) output.number = left.number == right.number;
    if (operation == OP_NE) output.number = left.number != right.number;
    if (operation == OP_LT) output.number = left.number < right.number;
    if (operation == OP_LE) output.number = left.number <= right.number;
    if (operation == OP_GT) output.number = left.number > right.number;
    if (operation == OP_GE) output.number = left.number >= right.number;
  }

  if (output.known && !isfinite(output.number)) return set_error(error, line, "non-finite arithmetic result");
  if (output.known && output.type == TYPE_DURATION && output.number < 0) {
    return set_error(error, line, "duration arithmetic produced a negative value");
  }
  return replace(stack, 2, &output, 1, info, error, line);
}

static bool check_operation(Operation operation, TypeStack* stack, TypeInfo* info,
                            Error* error, int line) {
  if (operation <= OP_ROT) return check_stack_operation(operation, stack, info, error, line);
  if (operation >= OP_ADD && operation <= OP_GE) {
    return check_numeric(operation, stack, info, error, line);
  }

  StackValue output[2] = {0};
  size_t output_count = 1;
  size_t input_count = 1;
  switch (operation) {
    case OP_GET_VIDEO:
    case OP_GET_AUDIO:
    case OP_GET_SUBTITLES:
    case OP_SPLIT:
      if (!need(stack, 1, error, line) || !expect(stack, 0, TYPE_COMBINED, error, line)) return false;
      output[0].type = operation == OP_GET_AUDIO ? TYPE_AUDIO
                       : operation == OP_GET_SUBTITLES ? TYPE_SUBTITLES : TYPE_VIDEO;
      if (operation == OP_SPLIT) { output[1].type = TYPE_AUDIO; output_count = 2; }
      break;
    case OP_VIDEO_AT:
    case OP_AUDIO_AT:
    case OP_SUBTITLES_AT:
    case OP_TAKE_VIDEO:
    case OP_TAKE_AUDIO:
    case OP_TAKE_SUBTITLES: {
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_COMBINED, error, line) ||
          !expect(stack, 0, TYPE_INDEX, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      output[1].type = operation == OP_VIDEO_AT || operation == OP_TAKE_VIDEO ? TYPE_VIDEO
                       : operation == OP_AUDIO_AT || operation == OP_TAKE_AUDIO ? TYPE_AUDIO
                       : TYPE_SUBTITLES;
      output_count = 2;
      break;
    }
    case OP_DROP_VIDEO:
    case OP_DROP_AUDIO:
    case OP_DROP_SUBTITLES:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_COMBINED, error, line) ||
          !expect(stack, 0, TYPE_INDEX, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_ADD_VIDEO:
    case OP_ADD_AUDIO:
    case OP_ADD_SUBTITLES:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_COMBINED, error, line) ||
          !expect(stack, 0, operation == OP_ADD_VIDEO ? TYPE_VIDEO
                            : operation == OP_ADD_AUDIO ? TYPE_AUDIO : TYPE_SUBTITLES,
                  error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_REPLACE_VIDEO:
    case OP_REPLACE_AUDIO:
    case OP_REPLACE_SUBTITLES:
      input_count = 3;
      if (!need(stack, 3, error, line) || !expect(stack, 2, TYPE_COMBINED, error, line) ||
          !expect(stack, 1, TYPE_INDEX, error, line) ||
          !expect(stack, 0, operation == OP_REPLACE_VIDEO ? TYPE_VIDEO
                            : operation == OP_REPLACE_AUDIO ? TYPE_AUDIO : TYPE_SUBTITLES,
                  error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_MOVE_VIDEO:
    case OP_MOVE_AUDIO:
    case OP_MOVE_SUBTITLES:
      input_count = 3;
      if (!need(stack, 3, error, line) || !expect(stack, 2, TYPE_COMBINED, error, line) ||
          !expect(stack, 1, TYPE_INDEX, error, line) || !expect(stack, 0, TYPE_INDEX, error, line)) {
        return false;
      }
      output[0].type = TYPE_COMBINED;
      break;
    case OP_LENGTH: {
      if (!need(stack, 1, error, line)) return false;
      ValueType type = stack->values[stack->count - 1].type;
      if (type != TYPE_VIDEO && type != TYPE_AUDIO && type != TYPE_SUBTITLES &&
          type != TYPE_COMBINED) {
        return set_error(error, line, "length expects a media value");
      }
      output[0].type = TYPE_DURATION;
      break;
    }
    case OP_SLICE: {
      input_count = 3;
      if (!need(stack, 3, error, line) || !expect(stack, 0, TYPE_DURATION, error, line) ||
          !expect(stack, 1, TYPE_DURATION, error, line)) return false;
      ValueType type = stack->values[stack->count - 3].type;
      if (type != TYPE_VIDEO && type != TYPE_AUDIO && type != TYPE_SUBTITLES &&
          type != TYPE_COMBINED) {
        return set_error(error, line, "slice expects media DURATION DURATION");
      }
      output[0].type = type;
      break;
    }
    case OP_TAKE:
    case OP_DELAY:
    {
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 0, TYPE_DURATION, error, line)) return false;
      ValueType type = stack->values[stack->count - 2].type;
      if (type != TYPE_VIDEO && type != TYPE_AUDIO && type != TYPE_SUBTITLES &&
          type != TYPE_COMBINED) {
        return set_error(error, line, "operation expects a media value");
      }
      output[0].type = type;
      break;
    }
    case OP_FADE_IN:
    case OP_FADE_OUT: {
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 0, TYPE_DURATION, error, line)) return false;
      ValueType type = stack->values[stack->count - 2].type;
      if (type != TYPE_VIDEO && type != TYPE_AUDIO) {
        return set_error(error, line, "operation expects a VIDEO or AUDIO stream");
      }
      output[0].type = type;
      break;
    }
    case OP_CONCAT:
    case OP_MIX: {
      input_count = 2;
      if (!need(stack, 2, error, line)) return false;
      ValueType left = stack->values[stack->count - 2].type;
      ValueType right = stack->values[stack->count - 1].type;
      if (operation == OP_MIX && (left != TYPE_AUDIO || right != TYPE_AUDIO)) {
        return set_error(error, line, "mix expects AUDIO AUDIO");
      }
      if (operation == OP_CONCAT &&
          (left != right || (left != TYPE_VIDEO && left != TYPE_AUDIO &&
                             left != TYPE_SUBTITLES && left != TYPE_COMBINED))) {
        return set_error(error, line, "concat expects two media values of the same type");
      }
      output[0].type = left;
      break;
    }
    case OP_MUX:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_VIDEO, error, line) ||
          !expect(stack, 0, TYPE_AUDIO, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_ATTACH_SUBTITLES:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_COMBINED, error, line) ||
          !expect(stack, 0, TYPE_SUBTITLES, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_VIDEO_TO_COMBINED:
      if (!need(stack, 1, error, line) || !expect(stack, 0, TYPE_VIDEO, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_BURN_SUBTITLES: {
      bool styled = need(stack, 1, error, line) &&
                    stack->values[stack->count - 1].type == TYPE_BURN_STYLE;
      input_count = styled ? 3 : 2;
      if (!need(stack, input_count, error, line) ||
          !expect(stack, styled ? 2 : 1, TYPE_VIDEO, error, line) ||
          !expect(stack, styled ? 1 : 0, TYPE_SUBTITLES, error, line) ||
          (styled && !expect(stack, 0, TYPE_BURN_STYLE, error, line))) return false;
      output[0].type = TYPE_VIDEO;
      break;
    }
    case OP_SET_LANGUAGE:
    case OP_SET_TITLE:
    case OP_SET_DEFAULT:
    case OP_SET_FORCED: {
      input_count = 2;
      if (!need(stack, 2, error, line)) return false;
      ValueType media = stack->values[stack->count - 2].type;
      if (media != TYPE_VIDEO && media != TYPE_AUDIO && media != TYPE_SUBTITLES) {
        return set_error(error, line, "track metadata setters expect a media track");
      }
      ValueType parameter = operation == OP_SET_LANGUAGE || operation == OP_SET_TITLE
                              ? TYPE_STRING : TYPE_BOOL;
      if (!expect(stack, 0, parameter, error, line)) return false;
      output[0].type = media;
      break;
    }
    case OP_CHAPTER:
      input_count = 3;
      if (!need(stack, 3, error, line) || !expect(stack, 2, TYPE_DURATION, error, line) ||
          !expect(stack, 1, TYPE_DURATION, error, line) ||
          !expect(stack, 0, TYPE_STRING, error, line)) return false;
      output[0].type = TYPE_CHAPTER;
      break;
    case OP_ADD_CHAPTER:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_COMBINED, error, line) ||
          !expect(stack, 0, TYPE_CHAPTER, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_DROP_CHAPTER:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_COMBINED, error, line) ||
          !expect(stack, 0, TYPE_INDEX, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_MOVE_CHAPTER:
      input_count = 3;
      if (!need(stack, 3, error, line) || !expect(stack, 2, TYPE_COMBINED, error, line) ||
          !expect(stack, 1, TYPE_INDEX, error, line) || !expect(stack, 0, TYPE_INDEX, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_ADD_ATTACHMENT:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_COMBINED, error, line) ||
          !expect(stack, 0, TYPE_ATTACHMENT, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_DROP_ATTACHMENT:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_COMBINED, error, line) ||
          !expect(stack, 0, TYPE_INDEX, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_MOVE_ATTACHMENT:
      input_count = 3;
      if (!need(stack, 3, error, line) || !expect(stack, 2, TYPE_COMBINED, error, line) ||
          !expect(stack, 1, TYPE_INDEX, error, line) || !expect(stack, 0, TYPE_INDEX, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_SET_TAG:
      input_count = 3;
      if (!need(stack, 3, error, line) || !expect(stack, 2, TYPE_COMBINED, error, line) ||
          !expect(stack, 1, TYPE_STRING, error, line) || !expect(stack, 0, TYPE_STRING, error, line)) return false;
      output[0].type = TYPE_COMBINED;
      break;
    case OP_VOLUME:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_AUDIO, error, line) ||
          !expect(stack, 0, TYPE_SCALAR, error, line)) return false;
      output[0].type = TYPE_AUDIO;
      break;
    case OP_SILENCE:
      if (!need(stack, 1, error, line) || !expect(stack, 0, TYPE_DURATION, error, line)) return false;
      output[0].type = TYPE_AUDIO;
      break;
    case OP_SPEED:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 0, TYPE_SCALAR, error, line)) return false;
      if (stack->values[stack->count - 1].known && stack->values[stack->count - 1].number <= 0) {
        return set_error(error, line, "speed must be greater than zero");
      }
      output[0].type = stack->values[stack->count - 2].type;
      if (output[0].type != TYPE_VIDEO && output[0].type != TYPE_AUDIO &&
          output[0].type != TYPE_SUBTITLES && output[0].type != TYPE_COMBINED) {
        return set_error(error, line, "speed expects a media value");
      }
      break;
    case OP_AND:
    case OP_OR:
      input_count = 2;
      if (!need(stack, 2, error, line) || !expect(stack, 1, TYPE_BOOL, error, line) ||
          !expect(stack, 0, TYPE_BOOL, error, line)) return false;
      output[0] = (StackValue){
        .type = TYPE_BOOL,
        .known = stack->values[stack->count - 2].known && stack->values[stack->count - 1].known,
        .number = operation == OP_AND
                    ? stack->values[stack->count - 2].number && stack->values[stack->count - 1].number
                    : stack->values[stack->count - 2].number || stack->values[stack->count - 1].number,
      };
      break;
    case OP_NOT:
      if (!need(stack, 1, error, line) || !expect(stack, 0, TYPE_BOOL, error, line)) return false;
      output[0] = stack->values[stack->count - 1];
      output[0].number = !output[0].number;
      break;
    default: return set_error(error, line, "invalid operation usage");
  }
  return replace(stack, input_count, output, output_count, info, error, line);
}

static bool operation_at(const Program* program, size_t index, Operation operation) {
  return program->items[index].kind == INST_OPERATION && program->items[index].as.op == operation;
}

static bool find_if(const Program* program, size_t start, size_t limit,
                    size_t* else_index, size_t* end_index, Error* error) {
  unsigned depth = 1;
  *else_index = SIZE_MAX;
  for (size_t i = start + 1; i < limit; i++) {
    if (operation_at(program, i, OP_IF)) depth++;
    if (operation_at(program, i, OP_END) && --depth == 0) { *end_index = i; return true; }
    if (operation_at(program, i, OP_ELSE) && depth == 1) {
      if (*else_index != SIZE_MAX) return set_error(error, program->items[i].line, "duplicate else");
      *else_index = i;
    }
  }
  return set_error(error, program->items[start].line, "if without matching end");
}

static bool find_loop(const Program* program, size_t start, size_t limit,
                      size_t* while_index, size_t* repeat_index, Error* error) {
  unsigned depth = 1;
  *while_index = SIZE_MAX;
  for (size_t i = start + 1; i < limit; i++) {
    if (operation_at(program, i, OP_BEGIN)) depth++;
    if (operation_at(program, i, OP_REPEAT) && --depth == 0) {
      if (*while_index == SIZE_MAX) return set_error(error, program->items[i].line, "loop has no while");
      *repeat_index = i;
      return true;
    }
    if (operation_at(program, i, OP_WHILE) && depth == 1) {
      if (*while_index != SIZE_MAX) return set_error(error, program->items[i].line, "duplicate while");
      *while_index = i;
    }
  }
  return set_error(error, program->items[start].line, "begin without matching repeat");
}

static bool check_range(const Program* program, TypedProgram* typed, size_t start, size_t end,
                        TypeStack* stack, Error* error) {
  for (size_t i = start; i < end; i++) {
    Instruction instruction = program->items[i];
    TypeInfo* info = &typed->items[i];
    if (instruction.kind == INST_ENCODING || instruction.kind == INST_FFMPEG_ARGS) {
      record(info, NULL, 0, NULL, 0);
    } else if (instruction.kind == INST_SOURCE) {
      StackValue value = {.type = (ValueType)instruction.as.source.kind};
      if (!push(stack, value, error, instruction.line)) return false;
      record(info, NULL, 0, &value, 1);
    } else if (instruction.kind == INST_DURATION || instruction.kind == INST_SCALAR ||
               instruction.kind == INST_BOOL || instruction.kind == INST_INDEX ||
               instruction.kind == INST_BURN_STYLE || instruction.kind == INST_STRING ||
               instruction.kind == INST_ATTACHMENT) {
      ValueType type = instruction.kind == INST_DURATION ? TYPE_DURATION
                       : instruction.kind == INST_SCALAR ? TYPE_SCALAR
                       : instruction.kind == INST_INDEX ? TYPE_INDEX
                       : instruction.kind == INST_BURN_STYLE ? TYPE_BURN_STYLE
                       : instruction.kind == INST_STRING ? TYPE_STRING
                       : instruction.kind == INST_ATTACHMENT ? TYPE_ATTACHMENT : TYPE_BOOL;
      StackValue value = {.type = type, .number = instruction.as.number, .known = true};
      if (!push(stack, value, error, instruction.line)) return false;
      record(info, NULL, 0, &value, 1);
    } else if (instruction.as.op == OP_IF) {
      if (!need(stack, 1, error, instruction.line) || !expect(stack, 0, TYPE_BOOL, error, instruction.line)) {
        return false;
      }
      StackValue condition = stack->values[--stack->count];
      record(info, &condition, 1, NULL, 0);
      size_t else_index, end_index;
      if (!find_if(program, i, end, &else_index, &end_index, error)) return false;
      TypeStack true_stack = *stack;
      TypeStack false_stack = *stack;
      size_t true_end = else_index == SIZE_MAX ? end_index : else_index;
      if (!check_range(program, typed, i + 1, true_end, &true_stack, error)) return false;
      if (else_index != SIZE_MAX &&
          !check_range(program, typed, else_index + 1, end_index, &false_stack, error)) return false;
      if (!same_stack_types(&true_stack, &false_stack)) {
        return set_error(error, instruction.line, "if branches must leave identical stack types");
      }
      for (size_t j = 0; j < true_stack.count; j++) {
        true_stack.values[j].known = true_stack.values[j].known && false_stack.values[j].known &&
                                     true_stack.values[j].number == false_stack.values[j].number;
      }
      *stack = true_stack;
      i = end_index;
    } else if (instruction.as.op == OP_BEGIN) {
      size_t while_index, repeat_index;
      if (!find_loop(program, i, end, &while_index, &repeat_index, error)) return false;
      TypeStack condition_stack = *stack;
      if (!check_range(program, typed, i + 1, while_index, &condition_stack, error)) return false;
      if (condition_stack.count != stack->count + 1 ||
          condition_stack.values[condition_stack.count - 1].type != TYPE_BOOL) {
        return set_error(error, program->items[while_index].line,
                         "loop condition must transform S into S BOOL");
      }
      StackValue loop_condition = condition_stack.values[condition_stack.count - 1];
      record(&typed->items[while_index], &loop_condition, 1, NULL, 0);
      condition_stack.count--;
      if (!same_stack_types(stack, &condition_stack)) {
        return set_error(error, program->items[while_index].line,
                         "loop condition must preserve underlying stack types");
      }
      TypeStack body_stack = condition_stack;
      if (!check_range(program, typed, while_index + 1, repeat_index, &body_stack, error)) return false;
      if (!same_stack_types(&condition_stack, &body_stack)) {
        return set_error(error, program->items[repeat_index].line,
                         "loop body must preserve the loop stack types");
      }
      for (size_t j = 0; j < condition_stack.count; j++) condition_stack.values[j].known = false;
      *stack = condition_stack;
      i = repeat_index;
    } else if (instruction.as.op == OP_ELSE || instruction.as.op == OP_END ||
               instruction.as.op == OP_WHILE || instruction.as.op == OP_REPEAT) {
      return set_error(error, instruction.line, "unmatched control-flow marker");
    } else if (!check_operation(instruction.as.op, stack, info, error, instruction.line)) {
      return false;
    }
  }
  return true;
}

bool typecheck_program(const Program* program, TypedProgram* out, Error* error) {
  *out = (TypedProgram){.arena = arena_create(), .count = program->count};
  out->items = arena_alloc(&out->arena, program->count * sizeof(*out->items));
  if (out->items == NULL && program->count != 0) return set_error(error, 0, "out of memory");
  memset(out->items, 0, program->count * sizeof(*out->items));

  TypeStack stack = {0};
  if (!check_range(program, out, 0, program->count, &stack, error)) goto failure;
  if (stack.count != 1) {
    error->line = program->count == 0 ? 1 : program->items[program->count - 1].line;
    snprintf(error->message, sizeof(error->message),
             "programme leaves %zu values on the stack (expected exactly one)", stack.count);
    goto failure;
  }
  out->final_type = stack.values[0].type;
  if (out->final_type != TYPE_VIDEO && out->final_type != TYPE_AUDIO &&
      out->final_type != TYPE_SUBTITLES && out->final_type != TYPE_COMBINED) {
    set_error(error, program->items[program->count - 1].line,
              "final value must be VIDEO, AUDIO, SUBTITLES, or COMBINED");
    goto failure;
  }
  return true;

failure:
  typed_program_destroy(out);
  return false;
}

void typed_program_destroy(TypedProgram* typed) {
  arena_destroy(&typed->arena);
  *typed = (TypedProgram){0};
}