#include "federation/coord.h"

#include "amduat/fed/ingest.h"
#include "amduat/asl/artifact_io.h"

#include <stdlib.h>
#include <string.h>

struct amduat_fed_coord {
  amduat_fed_coord_cfg_t cfg;
  amduat_fed_registry_value_t registry;
  bool registry_loaded;
  amduat_fed_view_t last_view;
  bool has_view;
  uint64_t last_tick_ms;
};

static bool amduat_fed_coord_has_registry_ref(amduat_reference_t ref) {
  return ref.hash_id != 0 && ref.digest.data != NULL && ref.digest.len != 0;
}

static int amduat_fed_coord_ref_cmp(amduat_reference_t a, amduat_reference_t b) {
  size_t min_len;
  int cmp;

  if (a.hash_id != b.hash_id) {
    return (int)a.hash_id - (int)b.hash_id;
  }
  if (a.digest.len != b.digest.len) {
    return a.digest.len < b.digest.len ? -1 : 1;
  }
  min_len = a.digest.len;
  if (min_len == 0) {
    return 0;
  }
  cmp = memcmp(a.digest.data, b.digest.data, min_len);
  if (cmp != 0) {
    return cmp;
  }
  return 0;
}

static int amduat_fed_coord_record_cmp(const void *lhs, const void *rhs) {
  const amduat_fed_record_t *a = (const amduat_fed_record_t *)lhs;
  const amduat_fed_record_t *b = (const amduat_fed_record_t *)rhs;

  if (a->logseq != b->logseq) {
    return a->logseq < b->logseq ? -1 : 1;
  }
  if (a->id.type != b->id.type) {
    return (int)a->id.type - (int)b->id.type;
  }
  return amduat_fed_coord_ref_cmp(a->id.ref, b->id.ref);
}

static bool amduat_fed_coord_record_clone(const amduat_fed_record_t *src,
                                          amduat_fed_record_t *out) {
  if (src == NULL || out == NULL) {
    return false;
  }
  *out = *src;
  if (!amduat_reference_clone(src->id.ref, &out->id.ref)) {
    return false;
  }
  return true;
}

static void amduat_fed_coord_record_free(amduat_fed_record_t *rec) {
  if (rec == NULL) {
    return;
  }
  amduat_reference_free(&rec->id.ref);
}

static void amduat_fed_coord_free_batch(amduat_fed_coord_t *coord,
                                        amduat_fed_record_t *batch,
                                        size_t batch_len) {
  if (coord == NULL || batch == NULL) {
    return;
  }
  if (coord->cfg.transport.free_records != NULL) {
    coord->cfg.transport.free_records(coord->cfg.transport.ctx,
                                      batch,
                                      batch_len);
  }
}

static amduat_fed_domain_state_t *amduat_fed_coord_find_state(
    amduat_fed_registry_value_t *registry,
    uint32_t domain_id) {
  size_t i;

  if (registry == NULL) {
    return NULL;
  }
  for (i = 0; i < registry->len; ++i) {
    if (registry->states[i].domain_id == domain_id) {
      return &registry->states[i];
    }
  }
  return NULL;
}

static bool amduat_fed_coord_records_push(amduat_fed_record_t **records,
                                          size_t *len,
                                          size_t *cap,
                                          amduat_fed_record_t value) {
  amduat_fed_record_t *next;
  size_t next_cap;

  if (records == NULL || len == NULL || cap == NULL) {
    return false;
  }
  if (*len == *cap) {
    next_cap = *cap != 0 ? *cap * 2u : 64u;
    next = (amduat_fed_record_t *)realloc(*records,
                                          next_cap * sizeof(*next));
    if (next == NULL) {
      return false;
    }
    *records = next;
    *cap = next_cap;
  }
  (*records)[(*len)++] = value;
  return true;
}

static bool amduat_fed_coord_denies_push(amduat_fed_policy_deny_t **denies,
                                         size_t *len,
                                         size_t *cap,
                                         amduat_fed_policy_deny_t value) {
  amduat_fed_policy_deny_t *next;
  size_t next_cap;

  if (denies == NULL || len == NULL || cap == NULL) {
    return false;
  }
  if (*len == *cap) {
    next_cap = *cap != 0 ? *cap * 2u : 32u;
    next = (amduat_fed_policy_deny_t *)realloc(*denies,
                                               next_cap * sizeof(*next));
    if (next == NULL) {
      return false;
    }
    *denies = next;
    *cap = next_cap;
  }
  (*denies)[(*len)++] = value;
  return true;
}

amduat_fed_coord_error_t amduat_fed_coord_open(
    const amduat_fed_coord_cfg_t *cfg,
    amduat_fed_coord_t **out_coord) {
  amduat_fed_coord_t *coord = NULL;

  if (out_coord == NULL) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }
  *out_coord = NULL;
  if (cfg == NULL || cfg->authoritative_store == NULL) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }

  coord = (amduat_fed_coord_t *)calloc(1, sizeof(*coord));
  if (coord == NULL) {
    return AMDUAT_FED_COORD_ERR_OOM;
  }
  coord->cfg = *cfg;
  if (amduat_fed_coord_has_registry_ref(cfg->registry_ref)) {
    if (!amduat_reference_clone(cfg->registry_ref, &coord->cfg.registry_ref)) {
      free(coord);
      return AMDUAT_FED_COORD_ERR_OOM;
    }
  } else {
    coord->cfg.registry_ref = amduat_reference(0u, amduat_octets(NULL, 0u));
  }
  amduat_fed_registry_value_init(&coord->registry, NULL, 0);
  coord->registry_loaded = false;
  memset(&coord->last_view, 0, sizeof(coord->last_view));
  coord->has_view = false;
  coord->last_tick_ms = 0;
  *out_coord = coord;
  return AMDUAT_FED_COORD_OK;
}

amduat_fed_coord_error_t amduat_fed_coord_close(amduat_fed_coord_t *coord) {
  if (coord == NULL) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }
  if (coord->has_view) {
    amduat_fed_view_free(&coord->last_view);
  }
  amduat_fed_registry_value_free(&coord->registry);
  amduat_reference_free(&coord->cfg.registry_ref);
  free(coord);
  return AMDUAT_FED_COORD_OK;
}

amduat_fed_coord_error_t amduat_fed_coord_load_registry(
    amduat_fed_coord_t *coord) {
  amduat_fed_registry_store_t store;
  amduat_fed_registry_value_t value;
  amduat_asl_store_error_t store_err = AMDUAT_ASL_STORE_OK;
  amduat_fed_registry_error_t err;

  if (coord == NULL) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }
  if (!amduat_fed_coord_has_registry_ref(coord->cfg.registry_ref)) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }

  amduat_fed_registry_store_init(&store, coord->cfg.authoritative_store);
  amduat_fed_registry_value_init(&value, NULL, 0);
  err = amduat_fed_registry_store_get(&store,
                                      coord->cfg.registry_ref,
                                      &value,
                                      &store_err);
  if (err == AMDUAT_FED_REGISTRY_ERR_CODEC) {
    amduat_fed_registry_value_free(&value);
    return AMDUAT_FED_COORD_ERR_CODEC;
  }
  if (err != AMDUAT_FED_REGISTRY_OK || store_err != AMDUAT_ASL_STORE_OK) {
    amduat_fed_registry_value_free(&value);
    return AMDUAT_FED_COORD_ERR_STORE;
  }

  amduat_fed_registry_value_free(&coord->registry);
  coord->registry = value;
  coord->registry_loaded = true;
  return AMDUAT_FED_COORD_OK;
}

amduat_fed_coord_error_t amduat_fed_coord_set_admitted(
    amduat_fed_coord_t *coord,
    uint32_t domain_id,
    bool admitted) {
  size_t i;
  amduat_fed_domain_state_t state;

  if (coord == NULL) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }
  for (i = 0; i < coord->registry.len; ++i) {
    if (coord->registry.states[i].domain_id == domain_id) {
      coord->registry.states[i].admitted = admitted ? 1u : 0u;
      if (admitted && coord->registry.states[i].policy_ok == 0u) {
        coord->registry.states[i].policy_ok = 1u;
      }
      return AMDUAT_FED_COORD_OK;
    }
  }

  memset(&state, 0, sizeof(state));
  state.domain_id = domain_id;
  state.admitted = admitted ? 1u : 0u;
  state.policy_ok = admitted ? 1u : 0u;
  state.policy_hash_id = 0;
  state.policy_hash = amduat_octets(NULL, 0u);
  if (!amduat_fed_registry_value_insert(&coord->registry, state)) {
    return AMDUAT_FED_COORD_ERR_OOM;
  }
  return AMDUAT_FED_COORD_OK;
}

amduat_fed_coord_error_t amduat_fed_coord_tick(
    amduat_fed_coord_t *coord,
    uint64_t now_ms) {
  amduat_fed_view_bounds_t *bounds = NULL;
  size_t bounds_len = 0;
  size_t bounds_cap = 0;
  amduat_fed_record_t *records = NULL;
  size_t records_len = 0;
  size_t records_cap = 0;
  amduat_fed_policy_deny_t *denies = NULL;
  size_t denies_len = 0;
  size_t denies_cap = 0;
  size_t i;
  amduat_fed_coord_error_t status = AMDUAT_FED_COORD_OK;
  amduat_fed_registry_store_t reg_store;
  amduat_fed_registry_error_t reg_err;
  amduat_reference_t new_ref;
  amduat_asl_store_error_t store_err = AMDUAT_ASL_STORE_OK;
  bool registry_dirty = false;

  if (coord == NULL) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }
  coord->last_tick_ms = now_ms;

  if (!coord->registry_loaded) {
    if (amduat_fed_coord_has_registry_ref(coord->cfg.registry_ref)) {
      status = amduat_fed_coord_load_registry(coord);
      if (status != AMDUAT_FED_COORD_OK) {
        return status;
      }
    } else {
      coord->registry_loaded = true;
    }
  }

  for (i = 0; i < coord->registry.len; ++i) {
    const amduat_fed_domain_state_t *state = &coord->registry.states[i];
    bool policy_ok = state->policy_ok != 0u ||
                     amduat_octets_is_empty(state->policy_hash);
    if (state->admitted == 0u || !policy_ok) {
      continue;
    }
    if (bounds_len == bounds_cap) {
      size_t next_cap = bounds_cap != 0 ? bounds_cap * 2u : 8u;
      amduat_fed_view_bounds_t *next =
          (amduat_fed_view_bounds_t *)realloc(
              bounds, next_cap * sizeof(*next));
      if (next == NULL) {
        status = AMDUAT_FED_COORD_ERR_OOM;
        goto tick_cleanup;
      }
      bounds = next;
      bounds_cap = next_cap;
    }
    bounds[bounds_len].domain_id = state->domain_id;
    bounds[bounds_len].snapshot_id = state->snapshot_id;
    bounds[bounds_len].log_prefix = state->log_prefix;
    bounds_len++;
  }

  for (i = 0; i < bounds_len; ++i) {
    amduat_fed_view_bounds_t *bound = &bounds[i];
    amduat_fed_domain_state_t *state =
        amduat_fed_coord_find_state(&coord->registry, bound->domain_id);
    amduat_fed_record_t *batch = NULL;
    size_t batch_len = 0;
    size_t j;
    uint64_t max_logseq = 0;
    int transport_rc;

    if (state == NULL) {
      status = AMDUAT_FED_COORD_ERR_INVALID;
      goto tick_cleanup;
    }
    if (coord->cfg.transport.get_records == NULL) {
      status = AMDUAT_FED_COORD_ERR_INVALID;
      goto tick_cleanup;
    }
    transport_rc = coord->cfg.transport.get_records(
        coord->cfg.transport.ctx,
        bound->domain_id,
        bound->snapshot_id,
        bound->log_prefix,
        state->last_logseq + 1u,
        &batch,
        &batch_len);
    if (transport_rc != 0) {
      status = AMDUAT_FED_COORD_ERR_STORE;
      goto tick_cleanup;
    }
    if (batch == NULL && batch_len != 0) {
      status = AMDUAT_FED_COORD_ERR_INVALID;
      goto tick_cleanup;
    }

    for (j = 0; j < batch_len; ++j) {
      amduat_fed_record_t cloned;
      amduat_fed_policy_deny_t deny;
      amduat_reference_t deny_ref;
      bool allowed = true;
      bool ok = amduat_fed_coord_record_clone(&batch[j], &cloned);
      if (!ok) {
        status = AMDUAT_FED_COORD_ERR_OOM;
        goto tick_cleanup;
      }
      if (cloned.logseq > bound->log_prefix) {
        amduat_fed_coord_record_free(&cloned);
        continue;
      }
      if (!amduat_fed_record_validate(&cloned)) {
        amduat_fed_coord_record_free(&cloned);
        status = AMDUAT_FED_COORD_ERR_INVALID;
        amduat_fed_coord_free_batch(coord, batch, batch_len);
        goto tick_cleanup;
      }
      if (coord->cfg.policy_hooks.record_allowed != NULL) {
        memset(&deny, 0, sizeof(deny));
        allowed = coord->cfg.policy_hooks.record_allowed(
            coord->cfg.policy_hooks.ctx, &cloned, &deny);
        if (!allowed) {
          if (deny.id.ref.digest.data == NULL && deny.id.ref.hash_id == 0u) {
            deny.id = cloned.id;
          }
          deny_ref = amduat_reference(0u, amduat_octets(NULL, 0u));
          if (!amduat_reference_clone(deny.id.ref, &deny_ref)) {
            amduat_fed_coord_record_free(&cloned);
            status = AMDUAT_FED_COORD_ERR_OOM;
            amduat_fed_coord_free_batch(coord, batch, batch_len);
            goto tick_cleanup;
          }
          deny.id.ref = deny_ref;
          if (!amduat_fed_coord_denies_push(&denies,
                                            &denies_len,
                                            &denies_cap,
                                            deny)) {
            amduat_reference_free(&deny.id.ref);
            amduat_fed_coord_record_free(&cloned);
            status = AMDUAT_FED_COORD_ERR_OOM;
            amduat_fed_coord_free_batch(coord, batch, batch_len);
            goto tick_cleanup;
          }
        }
      }
      if (!amduat_fed_coord_records_push(&records,
                                         &records_len,
                                         &records_cap,
                                         cloned)) {
        amduat_fed_coord_record_free(&cloned);
        status = AMDUAT_FED_COORD_ERR_OOM;
        amduat_fed_coord_free_batch(coord, batch, batch_len);
        goto tick_cleanup;
      }
      if (cloned.logseq > max_logseq) {
        max_logseq = cloned.logseq;
      }
    }

    amduat_fed_coord_free_batch(coord, batch, batch_len);

    if (max_logseq > state->last_logseq) {
      state->last_logseq = max_logseq;
      registry_dirty = true;
    }
  }

  if (records_len != 0) {
    size_t err_index = 0;
    size_t conflict_index = 0;
    amduat_fed_ingest_error_t ingest_rc;
    qsort(records, records_len, sizeof(*records), amduat_fed_coord_record_cmp);
    ingest_rc = amduat_fed_ingest_validate(records,
                                           records_len,
                                           &err_index,
                                           &conflict_index);
    if (ingest_rc != AMDUAT_FED_INGEST_OK) {
      status = AMDUAT_FED_COORD_ERR_INVALID;
      goto tick_cleanup;
    }
  }

  if (coord->has_view) {
    amduat_fed_view_free(&coord->last_view);
    coord->has_view = false;
  }

  if (bounds_len != 0) {
    amduat_fed_view_error_t view_rc;
    view_rc = amduat_fed_view_build(records,
                                    records_len,
                                    coord->cfg.local_domain_id,
                                    bounds,
                                    bounds_len,
                                    denies,
                                    denies_len,
                                    &coord->last_view);
    if (view_rc != AMDUAT_FED_VIEW_OK) {
      status = AMDUAT_FED_COORD_ERR_INVALID;
      goto tick_cleanup;
    }
    coord->has_view = true;
  }

  if (registry_dirty) {
    amduat_fed_registry_store_init(&reg_store, coord->cfg.authoritative_store);
    reg_err = amduat_fed_registry_store_put(&reg_store,
                                            &coord->registry,
                                            &new_ref,
                                            &store_err);
    if (reg_err != AMDUAT_FED_REGISTRY_OK ||
        store_err != AMDUAT_ASL_STORE_OK) {
      status = AMDUAT_FED_COORD_ERR_STORE;
      goto tick_cleanup;
    }
    amduat_reference_free(&coord->cfg.registry_ref);
    coord->cfg.registry_ref = new_ref;
  }

tick_cleanup:
  if (bounds != NULL) {
    free(bounds);
  }
  if (records != NULL) {
    for (i = 0; i < records_len; ++i) {
      amduat_fed_coord_record_free(&records[i]);
    }
    free(records);
  }
  if (denies != NULL) {
    for (i = 0; i < denies_len; ++i) {
      amduat_reference_free(&denies[i].id.ref);
    }
    free(denies);
  }
  return status;
}

amduat_fed_coord_error_t amduat_fed_coord_resolve(
    amduat_fed_coord_t *coord,
    amduat_reference_t ref,
    amduat_artifact_t *out_artifact) {
  amduat_fed_resolve_error_t rc;

  if (coord == NULL || out_artifact == NULL) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }
  if (!coord->has_view) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }

  rc = amduat_fed_resolve(&coord->last_view,
                          coord->cfg.authoritative_store,
                          ref,
                          out_artifact);
  if (rc == AMDUAT_FED_RESOLVE_OK) {
    return AMDUAT_FED_COORD_OK;
  }
  if (rc == AMDUAT_FED_RESOLVE_FOUND_REMOTE_NO_BYTES &&
      coord->cfg.cache_store != NULL &&
      coord->cfg.transport.get_artifact != NULL) {
    amduat_octets_t bytes = amduat_octets(NULL, 0u);
    amduat_artifact_t artifact;
    amduat_asl_store_error_t store_err;
    int fetch_rc = coord->cfg.transport.get_artifact(
        coord->cfg.transport.ctx, ref, &bytes);
    if (fetch_rc != 0) {
      amduat_octets_free(&bytes);
      return AMDUAT_FED_COORD_ERR_STORE;
    }
    if (!amduat_asl_artifact_from_bytes(bytes,
                                        AMDUAT_ASL_IO_RAW,
                                        false,
                                        amduat_type_tag(0u),
                                        &artifact)) {
      amduat_octets_free(&bytes);
      return AMDUAT_FED_COORD_ERR_INVALID;
    }
    store_err = amduat_asl_store_put(coord->cfg.cache_store,
                                     artifact,
                                     &ref);
    amduat_asl_artifact_free(&artifact);
    amduat_octets_free(&bytes);
    if (store_err != AMDUAT_ASL_STORE_OK) {
      return AMDUAT_FED_COORD_ERR_STORE;
    }
    rc = amduat_fed_resolve(&coord->last_view,
                            coord->cfg.authoritative_store,
                            ref,
                            out_artifact);
    if (rc == AMDUAT_FED_RESOLVE_OK) {
      return AMDUAT_FED_COORD_OK;
    }
  }
  if (rc == AMDUAT_FED_RESOLVE_POLICY_DENIED) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }
  if (rc == AMDUAT_FED_RESOLVE_NOT_FOUND) {
    return AMDUAT_FED_COORD_ERR_INVALID;
  }
  return AMDUAT_FED_COORD_ERR_STORE;
}

void amduat_fed_coord_get_status(const amduat_fed_coord_t *coord,
                                 amduat_fed_coord_status_t *out_status) {
  if (out_status == NULL) {
    return;
  }
  memset(out_status, 0, sizeof(*out_status));
  if (coord == NULL) {
    return;
  }
  out_status->domain_id = coord->cfg.local_domain_id;
  if (!amduat_reference_clone(coord->cfg.registry_ref,
                              &out_status->registry_ref)) {
    out_status->registry_ref = amduat_reference(0u, amduat_octets(NULL, 0u));
  }
  out_status->last_tick_ms = coord->last_tick_ms;
}