import type { Dictionary, NodeColor } from "__BUILTINS__"
import { DictRBEmpty, DictRBNode, DictRBBlack, DictRBRed } from "__BUILTINS__"
import __BUILTINS__ from "__BUILTINS__"
import type { Maybe } from "Maybe"
import { Nothing, Just } from "Maybe"

import List from "List"


/**
 * Dictionary represents a key - value data structure.
 * It is implemented as a red black tree
 */


instance Functor (Dictionary k) {
  map = (f, dict) => where(dict) {
    DictRBEmpty =>
      DictRBEmpty

    DictRBNode(c, k, v, left, right) =>
      DictRBNode(c, k, f(v), map(f, left), map(f, right))
  }
}


mapM :: Monad m => (a -> m b) -> Dictionary k a -> m (Dictionary k b)
export mapM = (f, dict) => where(dict) {
  DictRBEmpty =>
    pure(DictRBEmpty)

  DictRBNode(c, k, v, left, right) =>
    do {
      _v <- f(v)
      _left <- mapM(f, left)
      _right <- mapM(f, right)

      return of(DictRBNode(c, k, _v, _left, _right))
    }
}


mapWithKey :: (a -> b -> c) -> Dictionary a b -> Dictionary a c
export mapWithKey = (f, dict) => where(dict) {
  DictRBEmpty =>
    DictRBEmpty

  DictRBNode(c, k, v, left, right) =>
    DictRBNode(c, k, f(k, v), mapWithKey(f, left), mapWithKey(f, right))
}


filter :: Comparable k => (v -> Boolean) -> Dictionary k v -> Dictionary k v
export filter = (f, dict) =>
  reduceLeft((k, v, d) => f(v) ? insert(k, v, d) : d, empty, dict)


filterWithKey :: Comparable k => (k -> v -> Boolean) -> Dictionary k v -> Dictionary k v
export filterWithKey = (f, dict) =>
  reduceLeft((k, v, d) => f(k, v) ? insert(k, v, d) : d, empty, dict)


fromList :: Comparable k => List (#[k, v]) -> Dictionary k v
export fromList = __BUILTINS__.dictFromList


toList :: Dictionary k v -> List #[k, v]
export toList = __BUILTINS__.dictToList


reduceLeft :: (k -> v -> a -> a) -> a -> Dictionary k v -> a
export reduceLeft = (f, acc, dict) => where(dict) {
  DictRBEmpty =>
    acc

  DictRBNode(_, k, v, left, right) =>
    reduceLeft(f, f(k, v, reduceLeft(f, acc, left)), right)
}


reduce :: (k -> v -> a -> a) -> a -> Dictionary k v -> a
export reduce = reduceLeft


reduceRight :: (k -> v -> a -> a) -> a -> Dictionary k v -> a
export reduceRight = __BUILTINS__.dictReduceRight


get :: Comparable k => k -> Dictionary k v -> Maybe v
export get = (k, dict) => where(dict) {
  DictRBEmpty =>
    Nothing

  DictRBNode(_, _k, _v, left, right) =>
    if (k == _k) {
      Just(_v)
    } else if (k > _k) {
      get(k, right)
    } else {
      get(k, left)
    }
}


includes :: Comparable k => k -> Dictionary k v -> Boolean
export includes = (k, dict) => where(get(k, dict)) {
  Just(_) =>
    true

  Nothing =>
    false
}


insert :: Comparable k => k -> v -> Dictionary k v -> Dictionary k v
export insert = __BUILTINS__.dictInsert


empty :: Dictionary k v
export empty = DictRBEmpty


keys :: Dictionary a b -> List a
export keys = (dict) => reduceRight((k, _, ks) => [k, ...ks], [], dict)


values :: Dictionary a b -> List b
export values = (dict) => reduceRight((_, v, vs) => [v, ...vs], [], dict)


length :: Dictionary a b -> Integer
export length = (dict) => where(dict) {
  DictRBEmpty =>
    0

  DictRBNode(_, _, _, left, right) =>
    1 + length(left) + length(right)
}


isEmpty :: Dictionary a b -> Boolean
export isEmpty = (dict) => length(dict) == 0


singleton :: k -> v -> Dictionary k v
export singleton = (k, v) => DictRBNode(DictRBBlack, k, v, DictRBEmpty, DictRBEmpty)


/**
 * Merges the second dictionary into the first. If keys collide, the value from
 * the dict of the second parameter is chosen.
 *
 * @since 0.21.1
 */
merge :: Comparable k => Dictionary k v -> Dictionary k v -> Dictionary k v
export merge = (a, b) => reduceLeft(insert, a, b)


/**
 * Returns the dict given as first parameter, without the key/value pairs for
 * which the key doesn't exist in the second dictionary.
 *
 * @since 0.21.1
 */
intersect :: Comparable k => Dictionary k v -> Dictionary k v -> Dictionary k v
export intersect = (a, b) => filterWithKey((k, _) => includes(k, b), a)


/**
 * Returns the dict given as first parameter, without the key/value pairs for
 * which the key exists in the second dictionary.
 *
 * @since 0.21.1
 */
diff :: Comparable k => Dictionary k v -> Dictionary k v -> Dictionary k v
export diff = (a, b) => reduceLeft((k, _, dict) => remove(k, dict), a, b)


update :: Comparable k => (v -> v) -> k -> Dictionary k v -> Dictionary k v
export update = (f, k, dict) => where(get(k, dict)) {
  Just(v) =>
    __BUILTINS__.dictInsert(k, f(v), dict)

  Nothing =>
    dict
}


remove :: Comparable k => k -> Dictionary k v -> Dictionary k v
export remove = (key, dict) => where(removeHelp(key, dict)) {
  DictRBNode(DictRBRed, k, v, left, right) =>
    DictRBNode(DictRBBlack, k, v, left, right)

  or =>
    or
}


removeHelp :: Comparable k => k -> Dictionary k v -> Dictionary k v
removeHelp = (targetKey, dict) => where(dict) {
  DictRBEmpty =>
    DictRBEmpty

  DictRBNode(color, key, value, left, right) =>
    if (targetKey < key) {
      where(left) {
        DictRBNode(DictRBBlack, _, _, lLeft, _) =>
          where(lLeft) {
            DictRBNode(DictRBRed, _, _, _, _) =>
              DictRBNode(color, key, value, removeHelp(targetKey, left), right)

            _ =>
              where(moveRedLeft(dict)) {
                DictRBNode(nColor, nKey, nValue, nLeft, nRight) =>
                  __BUILTINS__.balanceDict(nColor, nKey, nValue, removeHelp(targetKey, nLeft), nRight)

                DictRBEmpty =>
                  DictRBEmpty
              }
          }

        _ =>
          DictRBNode(color, key, value, removeHelp(targetKey, left), right)
      }
    } else {
      removeHelpEQGT(targetKey, removeHelpPrepEQGT(dict, color, key, value, left, right))
    }
}


removeHelpPrepEQGT :: Comparable k => Dictionary k v -> NodeColor -> k -> v -> Dictionary k v -> Dictionary k v -> Dictionary k v
removeHelpPrepEQGT = (dict, color, key, value, left, right) => where(left) {
  DictRBNode(DictRBRed, lK, lV, lLeft, lRight) =>
    DictRBNode(color, lK, lV, lLeft, DictRBNode(DictRBRed, key, value, lRight, right))

  _ =>
    where(right) {
      DictRBNode(DictRBBlack, _, _, DictRBNode(DictRBBlack, _, _, _, _), _) =>
        moveRedRight(dict)

      DictRBNode(DictRBBlack, _, _, DictRBEmpty, _) =>
        moveRedRight(dict)

      _ =>
        dict
    }
}


removeHelpEQGT :: Comparable k => k -> Dictionary k v -> Dictionary k v
removeHelpEQGT = (targetKey, dict) => where(dict) {
  DictRBNode(c, k, v, left, right) =>
    if (targetKey == k) {
      where(getMin(right)) {
        DictRBNode(_, minKey, minValue, _, _) =>
          __BUILTINS__.balanceDict(c, minKey, minValue, left, removeMin(right))

        DictRBEmpty =>
          DictRBEmpty
      }
    } else {
      __BUILTINS__.balanceDict(c, k, v, left, removeHelp(targetKey, right))
    }

  DictRBEmpty =>
    DictRBEmpty
}


getMin :: Dictionary k v -> Dictionary k v
getMin = (dict) => where(dict) {
  DictRBNode(_, _, _, left, _) =>
    where(left) {
      DictRBNode(_, _, _, _, _) =>
        getMin(left)

      _ =>
        dict
    }

  _ =>
    dict
}


removeMin :: Dictionary k v -> Dictionary k v
removeMin = (dict) => where(dict) {
  DictRBNode(c, k, v, DictRBNode(lC, lK, lV, lLeft, lRight), right) =>
    where(lC) {
      DictRBBlack =>
        where(lLeft) {
          DictRBNode(DictRBRed, _, _, _, _) =>
            DictRBNode(c, k, v, removeMin(DictRBNode(lC, lK, lV, lLeft, lRight)), right)

          _ =>
            where(moveRedLeft(dict)) {
              DictRBNode(nC, nK, nV, nLeft, nRight) =>
                __BUILTINS__.balanceDict(nC, nK, nV, removeMin(nLeft), nRight)

              DictRBEmpty =>
                DictRBEmpty
            }
        }

      _ =>
        DictRBNode(c, k, v, removeMin(DictRBNode(lC, lK, lV, lLeft, lRight)), right)
    }

  _ =>
    DictRBEmpty
}


moveRedLeft :: Dictionary k v -> Dictionary k v
moveRedLeft = (dict) => where(dict) {
  DictRBNode(c, k, v, DictRBNode(lC, lK, lV, lLeft, lRight), DictRBNode(rC, rK, rV, DictRBNode(DictRBRed, rlK, rlV, rlL, rlR), rRight)) =>
    DictRBNode(DictRBRed, rlK, rlV, DictRBNode(DictRBBlack, k, v, DictRBNode(DictRBRed, lK, lV, lLeft, lRight), rlL), DictRBNode(DictRBBlack, rK, rV, rlR, rRight))

  DictRBNode(c, k, v, DictRBNode(lC, lK, lV, lLeft, lRight), DictRBNode(rC, rK, rV, rLeft, rRight)) =>
    where(c) {
      DictRBBlack =>
        DictRBNode(DictRBBlack, k, v, DictRBNode(DictRBRed, lK, lV, lLeft, lRight), DictRBNode(DictRBRed, rK, rV, rLeft, rRight))

      DictRBRed =>
        DictRBNode(DictRBBlack, k, v, DictRBNode(DictRBRed, lK, lV, lLeft, lRight), DictRBNode(DictRBRed, rK, rV, rLeft, rRight))
    }

  _ =>
    dict
}


moveRedRight :: Dictionary k v -> Dictionary k v
moveRedRight = (dict) => where(dict) {
  DictRBNode(c, k, v, DictRBNode(lC, lK, lV, DictRBNode(DictRBRed, llK, llV, llLeft, llRight), lRight), DictRBNode(rC, rK, rV, rLeft, rRight)) =>
    DictRBNode(DictRBRed, lK, lV, DictRBNode(DictRBBlack, llK, llV, llLeft, llRight), DictRBNode(DictRBBlack, k, v, lRight, DictRBNode(DictRBRed, rK, rV, rLeft, rRight)))

  DictRBNode(c, k, v, DictRBNode(lC, lK, lV, lLeft, lRight), DictRBNode(rC, rK, rV, rLeft, rRight)) =>
    where(c) {
      DictRBBlack =>
        DictRBNode(DictRBBlack, k, v, DictRBNode(DictRBRed, lK, lV, lLeft, lRight), DictRBNode(DictRBRed, rK, rV, rLeft, rRight))

      DictRBRed =>
        DictRBNode(DictRBBlack, k, v, DictRBNode(DictRBRed, lK, lV, lLeft, lRight), DictRBNode(DictRBRed, rK, rV, rLeft, rRight))
    }

  _ =>
    dict
}