package tst

import (
)

type TST struct {
  Root *Node
}

type Node struct {
  Char rune
  Value any
  Less *Node
  Equal *Node
  Greater *Node
}

func (t *TST) Add(key string, value any) {
  if t.Root == nil {
    // first run, setup with the start letter
    t.Root = &Node{rune(key[0]), nil, nil, nil, nil}
  }

  pos := t.Root

  for _, ch := range key {
    if ch < pos.Char {
      if pos.Less == nil {
        pos.Less = &Node{ch, nil, nil, nil, nil}
      }

      pos = pos.Less
    } else if ch > pos.Char {
      if pos.Greater == nil {
        pos.Greater = &Node{ch, nil, nil, nil, nil}
      }

      pos = pos.Greater
    } else if ch == pos.Char {
      if pos.Equal == nil {
        pos.Equal = &Node{ch, nil, nil, nil, nil}
      }

      pos = pos.Equal
    }

    pos.Value = value
  }
}

func (t *TST) Find(key string) (*Node) {
  if t.Root == nil {
    return nil
  }

  pos := t.Root
  next := t.Root

  for _, ch := range key {
    if ch < pos.Char {
      next = pos.Less
    } else if ch > pos.Char {
      next = pos.Greater
    } else if ch == pos.Char {
      next = pos.Equal
    }

    if next == nil {
      return nil
    } else {
      pos = next
    }
  }

  return pos
}

func (t *TST) Get(key string) any {
  node := t.Find(key)

  if node == nil {
    return nil
  } else {
    return node.Value
  }
}

func (t *TST) Delete(key string) {
  node := t.Find(key)

  if node != nil {
    node.Value = nil
  }
}

func (t *TST) StartsWith(key string) *Node {
  if t.Root == nil {
    return nil
  }

  pos := t.Root
  next := t.Root

  for _, ch := range key {
    if ch < pos.Char {
      next = pos.Less
    } else if ch > pos.Char {
      next = pos.Greater
    } else if ch == pos.Char {
      next = pos.Equal
      pos = pos.Equal
    }

    if next == nil {
      return pos
    }
  }

  return pos
}