#pragma once
#include <unordered_map>
#include <unistd.h>
#include <memory>
#include <sys/user.h>

#include "wirekit.hpp"

struct TraceContext {
	enum class EXECUTION_MODE { STARTING, USER, KERNEL /* Currently executing a system call */ };

	pid_t pid;
	EXECUTION_MODE mode;
	user_regs_struct regs;
	bool regs_dirty;
};

class TraceContextCollection {
	std::unordered_map<pid_t, std::unique_ptr<TraceContext>> m_contexts;
	pid_t m_active;
public:
	TraceContextCollection(const TraceContextCollection& other) = delete;
	TraceContextCollection(TraceContextCollection&& other);
	TraceContextCollection& operator=(const TraceContextCollection& other) = delete;
	TraceContextCollection& operator=(TraceContextCollection&& other);

	TraceContextCollection();
	~TraceContextCollection();

	bool add_context(pid_t pid);
	void delete_context(pid_t pid);
	void clear();

	TraceContext* context(pid_t pid);
	const TraceContext* context(pid_t pid) const;

	void context_set_active(pid_t pid);
	void context_clear_active();
	TraceContext* context_get_active();
	const TraceContext* context_get_active() const;
};

class Environment {
	struct HookPair {
		hook_t entry = nullptr;
		hook_t exit = nullptr;
	};

	std::unordered_map<reg_t, HookPair> m_hooks;
	HookPair m_default_hooks;
	TraceContextCollection m_trace_contexts;
public:
	Environment(const Environment& other) = delete;
	Environment(Environment&& other) = delete;
	Environment& operator=(const Environment& other) = delete;
	Environment& operator=(Environment&& other) = delete;

	Environment();
	~Environment();

	void register_hooks(reg_t syscall, hook_t entry, hook_t exit);
	void register_default_hooks(hook_t entry, hook_t exit);

	bool entry_hook(reg_t syscall) const;
	bool exit_hook(reg_t syscall) const;
	bool default_entry_hook() const;
	bool default_exit_hook() const;

	TraceContextCollection& contexts();
	const TraceContextCollection& contexts() const;

	void clear();
};

// An EnvironmentScope calls Environment::clear() when destructed
class EnvironmentScope {
	Environment& m_environment;
public:
	EnvironmentScope(const EnvironmentScope& other) = delete;
	EnvironmentScope(EnvironmentScope&& other) = delete;
	EnvironmentScope& operator=(const EnvironmentScope& other) = delete;
	EnvironmentScope& operator=(EnvironmentScope&& other) = delete;

	EnvironmentScope(Environment& environment);
	~EnvironmentScope();
};

extern Environment ENV;