//-
// SPDX-License-Identifier: BSD-2-Clause-FreeBSD
//
// Copyright (c) 2022 Tobias Kortkamp <tobik@FreeBSD.org>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
// OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
// LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
// OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
// SUCH DAMAGE.

#include "config.h"

#include <stdlib.h>

#include "flow.h"
#include "framemap.h"
#include "map.h"
#include "mem.h"
#include "mempool.h"
#include "set.h"
#include "stack.h"
#include "trait/compare.h"

struct Framemap {
	struct CompareTrait cmp;
	struct Stack *stack;
};

struct FramemapIterator {
	struct Map *map;
	struct MapIterator *iter;
};

struct Frame {
	struct Map *map;
};

// Prototypes
static void frame_free(struct Frame *);

// Constants
static const void *REMOVED = frame_free;

void
frame_free(struct Frame *frame)
{
	if (frame) {
		map_free(frame->map);
		free(frame);
	}
}

struct Framemap *
framemap_new(struct CompareTrait *cmp)
{
	struct Framemap *map = xmalloc(sizeof(struct Framemap));
	if (cmp) {
		map->cmp = *cmp;
	}
	map->stack = stack_new();
	framemap_push(map);
	return map;
}

void
framemap_free(struct Framemap *map)
{
	if (map) {
		STACK_FOREACH(map->stack, struct Frame *, frame) {
			frame_free(frame);
		}
		stack_free(map->stack);
		free(map);
	}
}

void
framemap_replace(struct Framemap *map, const void *key, const void *value)
{
	struct Frame *frame = stack_peek(map->stack);
	map_replace(frame->map, key, value);
}

void
framemap_remove(struct Framemap *map, const void *key)
{
	struct Frame *frame = stack_peek(map->stack);
	map_replace(frame->map, key, REMOVED);
}

void
framemap_remove_all(struct Framemap *map, const void *key)
{
	STACK_FOREACH(map->stack, struct Frame *, frame) {
		map_remove(frame->map, key);
	}
}

void *
framemap_get(struct Framemap *map, const void *key)
{
	STACK_FOREACH(map->stack, struct Frame *, frame) {
		void *val = map_get(frame->map, key);
		if (val && val != REMOVED) {
			return val;
		}
	}
	return NULL;
}

bool
framemap_contains(struct Framemap *map, const void *key)
{
	return framemap_get(map, key) != NULL;
}

void
framemap_push(struct Framemap *map)
{
	struct Frame *frame = xmalloc(sizeof(struct Frame));
	frame->map = map_new(&map->cmp);
	stack_push(map->stack, frame);
}

void
framemap_pop(struct Framemap *map)
{
	if (stack_len(map->stack) > 1) {
		stack_pop(map->stack);
	}
}

size_t
framemap_len(struct Framemap *map)
{
	SCOPE_MEMPOOL(pool);
	return map_len(framemap_flatten(map, pool));
}

struct Map *
framemap_flatten(struct Framemap *map, struct Mempool *extpool)
{
	SCOPE_MEMPOOL(pool);
	struct Map *flatmap = mempool_map(extpool, &map->cmp);
	struct Set *removed = mempool_set(pool, &map->cmp);
	STACK_FOREACH(map->stack, struct Frame *, frame) {
		MAP_FOREACH(frame->map, const void *, key, const void *, value) {
			if (value == REMOVED) {
				set_add(removed, key);
			} else unless (set_contains(removed, key)) {
				map_add(flatmap, key, value);
			}
		}
	}
	return flatmap;
}

struct Array *
framemap_keys(struct Framemap *map, struct Mempool *extpool)
{
	SCOPE_MEMPOOL(pool);
	return map_keys(framemap_flatten(map, pool), extpool);
}

struct Array *
framemap_values(struct Framemap *map, struct Mempool *extpool)
{
	SCOPE_MEMPOOL(pool);
	return map_values(framemap_flatten(map, pool), extpool);
}

void
framemap_truncate(struct Framemap *map)
{
	STACK_FOREACH(map->stack, struct Frame *, frame) {
		frame_free(frame);
	}
	framemap_push(map);
}

struct FramemapIterator *
framemap_iterator(struct Framemap *map, ssize_t a, ssize_t b)
{
	struct FramemapIterator *iter = xmalloc(sizeof(struct FramemapIterator));
	iter->map = framemap_flatten(map, NULL);
	iter->iter = map_iterator(iter->map, a, b);
	return iter;
}

void
framemap_iterator_cleanup(struct FramemapIterator **iter_)
{
	struct FramemapIterator *iter = *iter_;
	if (iter) {
		map_iterator_cleanup(&iter->iter);
		map_free(iter->map);
		free(iter);
		*iter_ = NULL;
	}
}

bool
framemap_iterator_next(struct FramemapIterator **iter_, size_t *index, void **key, void **value)
{
	struct FramemapIterator *iter = *iter_;
	if (map_iterator_next(&iter->iter, index, key, value)) {
		return true;
	} else {
		framemap_iterator_cleanup(iter_);
		return false;
	}
}
