#ifndef _SCHEDULER_HH_
#define _SCHEDULER_HH_

#include <cstdint>
#include <cstdlib>
#include <vector>
#include <cassert>
#include <queue>
#include <functional>
#include <memory>
#include <array>

#include "exceptions.hh"
#include "globals.hh"
#include "cortexa8.hh"

using thread_entry_t = void(void*);

class Scheduler {
public:
  class Process;
  class Semaphore;

  enum class ThreadState {
    ready,
      running,
      waiting,
      terminated,
      yield
      };

  class Thread {
  public:
    ~Thread() {
      assert(!neverDestruct_);
    }

    cortexa8::thread_cb* get_tcb() {
      return tcb_.get();
    }

    Thread(Thread const& copy) = delete;

    // Priority compare (here, higher = higher priority)
    // Result is only meaningful for runnable, scheduled threads
    bool operator<(Thread const& other) const {
      // Check for higher priority
      if(priority_ > other.priority_)
	return true; // I am lower prio
      else if(priority_ < other.priority_)
	return false; // Other thread is lower prio
      else {
	// TODO: Compare by waiting time
	return true;
      }
    }

    unsigned getPriority() const {
      return priority_;
    }

    // Set this threads priority. The lower the number the higher the priority
    void setPriority(unsigned priority) {
      priority_ = priority;
    }

    bool ready() const {
      return (state_ == ThreadState::ready) || (state_ == ThreadState::running);
    }

    ThreadState getState() const {
      return state_;
    }

  private:
    Thread(Process& proc) : proc_(proc), neverDestruct_{true}, priority_(0), state_{ThreadState::ready} {
      tcb_.reset(cortexa8::get_cur_thread());
    } 

    Thread(Process& proc, thread_entry_t entry, void *arg, size_t stacksize = 4096) 
      : tcb_{new cortexa8::thread_cb}, proc_(proc), stack_{new uint32_t[stacksize/4]}, neverDestruct_{false}, 
	priority_(0), state_{ThreadState::ready} {
      tcb_->regs[15] = (uint32_t)entry;
      tcb_->regs[13] = (uint32_t)stack_.get()+stacksize;
      tcb_->regs[1] = (uint32_t)arg;
      tcb_->cpsr = 0x0000015f;
    }

    void setState(ThreadState state) {
      state_ = state;
    }

    std::unique_ptr<cortexa8::thread_cb> tcb_;
    Process& proc_;
    std::unique_ptr<uint32_t> stack_;
    bool neverDestruct_;
    unsigned priority_;
    ThreadState state_;

    friend class Scheduler;
    friend class Semaphore;
  };

  class Process {

  private:
    Process() : ttable0_phys_(0), ttable0_(nullptr) {
    }

    uint32_t ttable0_phys_;
    uint32_t *ttable0_;

    std::vector<Thread> threads_;

    friend class Scheduler;
  };

  class Semaphore {
  public:
    Semaphore(int count = 1) : count_(count), waitPos_(0) {
      assert(count >= 0);
    }

    ~Semaphore() {
    }

    void acquire() {
      cortexa8::ScopedSetIF disint{};
      --count_;
      if(count_ < 0) {
	// Block
	auto curThread  = &global::scheduler->getCurrentThread();
	waitList_.at(waitPos_++) = curThread;
	curThread->setState(ThreadState::waiting);
	//disint.restore();
	global::scheduler->yield();
      }
    }

    bool tryAcquire() {
      cortexa8::ScopedSetIF disint{};
      --count_;
      if(count_ < 0) {
	++count_;
	return false;
      } else
	return true;
    }

    void release() {
      cortexa8::ScopedSetIF disint{};
      ++count_;
      if(count_ <= 0) {
	// Wake a thread
	auto thread = waitList_.at(--waitPos_);

	assert(thread->getState() == ThreadState::waiting);
	thread->setState(ThreadState::ready);
	global::scheduler->runQueue_.push(thread);
      }
    }
	
    bool hasWaiters() const {
      return (count_ < 0);
    }

    int getCount() const {
      return count_;
    }

  private:
    int count_;

    static constexpr int MaxWaiters = 8;
    std::array<Thread*, MaxWaiters> waitList_;
    int waitPos_;

    friend class Scheduler;
  };

  Scheduler();
  ~Scheduler();

  // Called by timer interrupt
  // Selects a new thread to run after a time slice has elapsed
  void timeslice();

  // Set tcb of new current thread if necessary
  // This may only be called when inside an exception/interrupt handler
  void update();
  
  // Cease execution of current thread for this timeslice
  // This may only be called when inside an exception/interrupt handler
  void yield();

  // Exit current thread. If this is the last thread of a process, the process terminates
  void exit();

  Thread& getCurrentThread() {
    return *currentThread_;
  }

private:
  struct ThreadCompare {
    bool operator()(Thread* const& lhs, Thread* const& rhs) const {
      return *lhs<*rhs;
    }
  };

  std::vector<std::unique_ptr<Process>> procs_;
  std::vector<std::unique_ptr<Thread>> threads_;

  std::priority_queue<Thread*, std::vector<Thread*>, ThreadCompare> runQueue_;

  // Process for kernel threads (no ttable0)
  Process *kernelProc_;
  // Initial thread and the idle thread
  Thread *kernelThread_, *idleThread_;

  Thread *currentThread_;
};

class SemScopedAcquire {
public:
  SemScopedAcquire(Scheduler::Semaphore &sem) : sem_(sem), done_(false) {
    sem_.acquire();
  }

  ~SemScopedAcquire() {
    if(done_)
      return;

    sem_.release();
  }

  void release() {
    assert(!done_);
    sem_.release();
    done_ = true;
  }

  void reacquire() {
    assert(done_);
    sem_.acquire();
    done_ = false;
  }

private:
  Scheduler::Semaphore& sem_;
  bool done_;
};

#endif