package main

import (
  "crypto/ed25519"
  "crypto/rand"
  "crypto/x509"
  "encoding/binary"
  "encoding/pem"
  "flag"
  "fmt"
  "os"
  "os/signal"
  "syscall"
  "time"

  "git.metznet.ca/MetzNet/go-ncurses"
  "git.metznet.ca/MetzNet/pnyx"
  "github.com/gen2brain/malgo"
  "github.com/google/uuid"
  "github.com/hraban/opus"
)

var decoders = map[pnyx.PeerID]chan []byte{}
var encoder *opus.Encoder
var sample_rate int = 0

func set_sample_rate(audio_data chan []int16, new_sample_rate int) error {
  sample_rate = new_sample_rate

  var err error
  encoder, err = opus.NewEncoder(new_sample_rate, 1, opus.AppVoIP)
  if err != nil {
    return err
  }

  for peer_id, decoder_chan := range decoders {
    if decoder_chan != nil {
      decoder_chan <- nil
    }
    new_chan := make(chan []byte, 1000)
    decoders[peer_id] = new_chan
    go handle_peer_decode(audio_data, peer_id, decoders[peer_id], sample_rate)
  }
  return nil
}

func handle_peer_decode(audio_data chan []int16, peer_id pnyx.PeerID, decode_chan chan []byte, sample_rate int) {
  decoder, err := opus.NewDecoder(sample_rate, 1)
  if err != nil {
    panic(err)
  }

  running := true
  missed := 0
  for running {
    select {
    case <-time.After(20 * time.Millisecond):
      missed += 1

      if missed > 100 {
        decode_chan <- <-decode_chan
      }

      pcm := make([]int16, sample_rate/50)
      err := decoder.DecodePLC(pcm)
      if err != nil {
        panic(err)
      }
      audio_data <- pcm

    case data := <-decode_chan:
      missed = 0
      if data == nil {
        running = false
      } else {
        pcm := make([]int16, sample_rate/50*2)
        written, err := decoder.Decode(data, pcm)
        if err != nil {
          panic(err)
        }
        audio_data <- pcm[:written]
      }
    }
  }
}

func mixer(data_chan chan []int16, speaker_chan chan []int16) {
  var samples []int16 = nil
  for true {
    if samples == nil {
      samples = <-data_chan
    } else {
      select {
      case new_samples := <-data_chan:
        for i, sample := range new_samples {
          samples[i] += sample
        }
      case speaker_chan <- samples:
        samples = nil
      }
    }
  }
}

func setup_audio(mic chan []byte, speaker chan []int16) (*malgo.AllocatedContext, *malgo.Device, *malgo.Device) {
  ctx, err := malgo.InitContext(nil, malgo.ContextConfig{}, nil)
  if err != nil {
    panic(err)
  }

  infos, err := ctx.Devices(malgo.Playback)
  if err != nil {
    panic(err)
  }

  var playback_device *malgo.DeviceInfo = nil
  for _, info := range infos {
    if info.IsDefault != 0 {
      full, err := ctx.DeviceInfo(malgo.Playback, info.ID, malgo.Shared)
      if err != nil {
        panic(err)
      }
      playback_device = &full
    }
  }

  if playback_device == nil {
    panic("No default playback device")
  }

  infos, err = ctx.Devices(malgo.Capture)
  if err != nil {
    panic(err)
  }

  var capture_device *malgo.DeviceInfo = nil
  for _, info := range infos {
    if info.IsDefault != 0 {
      full, err := ctx.DeviceInfo(malgo.Capture, info.ID, malgo.Shared)
      if err != nil {
        panic(err)
      }
      capture_device = &full
    }
  }

  if capture_device == nil {
    panic("No default capture device")
  }

  inDeviceConfig := malgo.DefaultDeviceConfig(malgo.Capture)
  inDeviceConfig.Capture.Format = malgo.FormatS16
  inDeviceConfig.Capture.Channels = 1
  inDeviceConfig.Capture.DeviceID = capture_device.ID.Pointer()
  inDeviceConfig.SampleRate = 48000
  inDeviceConfig.PeriodSizeInFrames = 960
  inDeviceConfig.Alsa.NoMMap = 1
  inDeviceConfig.Capture.ShareMode = malgo.Shared

  outDeviceConfig := malgo.DefaultDeviceConfig(malgo.Playback)
  outDeviceConfig.Playback.Format = malgo.FormatS16
  outDeviceConfig.Playback.Channels = 1
  outDeviceConfig.Playback.DeviceID = playback_device.ID.Pointer()
  outDeviceConfig.SampleRate = 48000
  outDeviceConfig.PeriodSizeInFrames = 960
  outDeviceConfig.Alsa.NoMMap = 1
  outDeviceConfig.Playback.ShareMode = malgo.Shared

  onSendFrames := func(output_samples []byte, input_samples []byte, framecount uint32) {
    select {
    case pcm := <-speaker:
      for i := 0; i < sample_rate/50; i++ {
        binary.LittleEndian.PutUint16(output_samples[i*2:], uint16(pcm[i]))
      }
    default:
    }
  }

  playbackCallbacks := malgo.DeviceCallbacks{
    Data: onSendFrames,
  }

  outDevice, err := malgo.InitDevice(ctx.Context, outDeviceConfig, playbackCallbacks)
  if err != nil {
    panic(err)
  }

  err = outDevice.Start()
  if err != nil {
    panic(err)
  }

  onRecvFrames := func(output_samples []byte, input_samples []byte, framecount uint32) {
    if encoder != nil {
      pcm := make([]int16, len(input_samples)/2)
      for i := 0; i < len(input_samples)/2; i++ {
        pcm[i] = int16(binary.LittleEndian.Uint16(input_samples[2*i:]))
      }

      data := make([]byte, len(input_samples))
      written, err := encoder.Encode(pcm, data)
      if err != nil {
        panic(err)
      }
      select {
      case mic <- data[:written]:
      default:
      }
    }
  }

  captureCallbacks := malgo.DeviceCallbacks{
    Data: onRecvFrames,
  }

  inDevice, err := malgo.InitDevice(ctx.Context, inDeviceConfig, captureCallbacks)
  if err != nil {
    panic(err)
  }

  err = inDevice.Start()
  if err != nil {
    panic(err)
  }

  return ctx, outDevice, inDevice
}

func get_private_key(path string, generate bool) ed25519.PrivateKey {
  key_file_bytes, err := os.ReadFile(path)
  if err == nil {
    key_pem, _ := pem.Decode(key_file_bytes)
    if key_pem.Type != "PRIVATE KEY" {
      panic("Key file has wrong PEM format")
    }

    private_key, err := x509.ParsePKCS8PrivateKey(key_pem.Bytes)
    if err != nil {
      panic(err)
    }

    var ok bool
    key, ok := private_key.(ed25519.PrivateKey)
    if ok == false {
      panic("Private key is not ed25519.PrivateKey")
    }
    return key
  } else if generate {
    _, key, err := ed25519.GenerateKey(rand.Reader)
    if err != nil {
      panic(err)
    }

    key_pkcs8, err := x509.MarshalPKCS8PrivateKey(key)
    if err != nil {
      panic(err)
    }

    key_pem := pem.EncodeToMemory(&pem.Block{
      Type:  "PRIVATE KEY",
      Bytes: key_pkcs8,
    })

    err = os.WriteFile(path, key_pem, 0o600)
    if err != nil {
      panic(err)
    }
    return key
  } else {
    panic(fmt.Sprintf("Failed to read key from %s", path))
  }
}

func print_decoration(window, title, channels, body ncurses.Window, top_left string) {
  max_y := ncurses.GetMaxY(*ncurses.CurScr)
  max_x := ncurses.GetMaxX(*ncurses.CurScr)

  ncurses.WResize(window, max_y, max_x)
  ncurses.WResize(title, 1, max_x/4-2)
  ncurses.MvWin(title, 1, 1)
  ncurses.WResize(channels, max_y-4, (max_x/4)-1)
  ncurses.MvWin(channels, 3, 1)
  ncurses.WResize(body, max_y-2, max_x-(max_x/4)-2)
  ncurses.MvWin(body, 1, (max_x/4)+1)

  ncurses.MvWAddStr(title, 0, 0, top_left)

  for i := 1; i < max_x-1; i++ {
    ncurses.MvWAddStr(window, 0, i, "═")
  }

  for i := 1; i < max_x/4; i++ {
    ncurses.MvWAddStr(window, 2, i, "═")
  }

  for i := 1; i < max_y-1; i++ {
    ncurses.MvWAddStr(window, i, 0, "║")
  }

  for i := 1; i < max_y-1; i++ {
    ncurses.MvWAddStr(window, i, max_x-1, "║")
  }

  for i := 1; i < max_y-1; i++ {
    ncurses.MvWAddStr(window, i, max_x/4, "║")
  }

  for i := 1; i < max_x-1; i++ {
    ncurses.MvWAddStr(window, max_y-1, i, "═")
  }

  ncurses.MvWAddStr(window, 0, 0, "╔")
  ncurses.MvWAddStr(window, 0, max_x-1, "╗")
  ncurses.MvWAddStr(window, 0, max_x/4, "╦")
  ncurses.MvWAddStr(window, 2, 0, "╠")
  ncurses.MvWAddStr(window, 2, max_x/4, "╣")
  ncurses.MvWAddStr(window, max_y-1, 0, "╚")
  ncurses.MvWAddStr(window, max_y-1, max_x/4, "╩")
  ncurses.MvWAddStr(window, max_y-1, max_x-1, "╝")

  ncurses.WRefresh(window)
  ncurses.WRefresh(title)
  ncurses.WRefresh(channels)
  ncurses.WRefresh(body)
}

type ChannelState struct {
  ID pnyx.ChannelID
  Name string
  Members []pnyx.PeerID
  Modes []pnyx.ModeID
}

type PeerState struct {
  Name string
}

type UIState struct {
  SelectedBody    int
  SelectedPeer    int
  SelectedChannel int
  SelectedArea    int
  Channels        []ChannelState
  Peers           map[pnyx.PeerID]PeerState
}

const (
  AREA_CHANNELS = 0
  AREA_BODY = 1

  BODY_CHANNEL = 0
  BODY_DETAIL = 1
  BODY_QUIT = 2
)

func handle_back(state *UIState) {
  if state.SelectedArea == AREA_BODY {
    if state.SelectedBody == BODY_QUIT {
      state.SelectedBody = BODY_CHANNEL
    } else if state.SelectedBody == BODY_DETAIL {
      state.SelectedBody = BODY_CHANNEL
    }
    state.SelectedArea = AREA_CHANNELS
  }
}

func handle_up(state *UIState) {
  if state.SelectedArea == AREA_CHANNELS {
    if state.SelectedPeer > 0 {
      state.SelectedPeer -= 1
    } else if state.SelectedChannel > 0 {
      state.SelectedChannel -= 1
      state.SelectedPeer = len(state.Channels[state.SelectedChannel].Members)
    }
  }
}

func handle_down(state *UIState) {
  if state.SelectedArea == AREA_CHANNELS {
    if state.SelectedPeer < len(state.Channels[state.SelectedChannel].Members) {
      state.SelectedPeer += 1
    } else if state.SelectedChannel < len(state.Channels) - 1 {
      state.SelectedChannel += 1
      state.SelectedPeer = 0
    }
  }
}

func handle_select(state *UIState) {
  if state.SelectedArea == AREA_CHANNELS {
    state.SelectedArea = AREA_BODY
    state.SelectedBody = BODY_CHANNEL
  }
}

func handle_edit(state *UIState) {
  if state.SelectedArea == AREA_CHANNELS {
    state.SelectedArea = AREA_BODY
    state.SelectedBody = BODY_DETAIL
  }
}

func handle_quit(state *UIState) {
  state.SelectedArea = AREA_BODY
  state.SelectedBody = BODY_QUIT
}

func handle_input(state *UIState, char []byte) {
  switch char[0] {
  case 'q':
    handle_quit(state)
  case 'e':
    handle_edit(state)
  case '\r':
    handle_select(state)
  case 0x1B:
    if len(char) == 1 {
      handle_back(state)
    } else {
      switch char[1] {
      case '[':
        if len(char) == 2 {
        } else {
          switch char[2] {
            case 'A': // Up
              handle_up(state)
            case 'B': // Down
              handle_down(state)
            case 'C': // Right
              handle_select(state)
            case 'D': // Left
              handle_back(state)
          }
        }
      }
    }
  }
}

func draw_channels(state *UIState, channels ncurses.Window) {
  channel_highlight_color := 2
  if state.SelectedArea == AREA_CHANNELS {
    channel_highlight_color = 1
  }

  ncurses.WErase(channels)

  ncurses.WMove(channels, 0, 0)
  for i, channel_state := range(state.Channels) {
    if i == state.SelectedChannel && state.SelectedPeer == 0 {
      ncurses.WAttrOn(channels, ncurses.COLOR_PAIR(channel_highlight_color))
    }

    ncurses.WAddStr(channels, channel_state.Name)
    ncurses.WAddCh(channels, '\n')
    for j, peer_id := range(channel_state.Members) {
      ncurses.WAddStr(channels, "  ↳")
      peer_state, found := state.Peers[peer_id]
      if state.SelectedChannel == i && state.SelectedPeer == (j+1) {
        ncurses.WAttrOn(channels, ncurses.COLOR_PAIR(channel_highlight_color))
      }
      if found {
        ncurses.WAddStr(channels, peer_state.Name)
      } else {
        ncurses.WAddStr(channels, fmt.Sprintf("%s", peer_id))
      }
      if state.SelectedChannel == i && state.SelectedPeer == (j+1) {
        ncurses.WAttrOff(channels, ncurses.COLOR_PAIR(channel_highlight_color))
      }
      ncurses.WAddCh(channels, '\n')
    }

    if i == state.SelectedChannel && state.SelectedPeer == 0 {
      ncurses.WAttrOff(channels, ncurses.COLOR_PAIR(channel_highlight_color))
    }
  }

  ncurses.WRefresh(channels)
}

func draw_body_channel(state *UIState, body ncurses.Window) {
  ncurses.WErase(body)
  ncurses.MvWAddStr(body, 0, 0, "channel")
  ncurses.WRefresh(body)
}

func draw_body_detail(state *UIState, body ncurses.Window) {
  ncurses.WErase(body)
  ncurses.MvWAddStr(body, 0, 0, "detail")
  ncurses.WRefresh(body)
}

func draw_body_quit(state *UIState, body ncurses.Window) {
  ncurses.WErase(body)
  ncurses.MvWAddStr(body, 0, 0, "quit")
  ncurses.WRefresh(body)
}

func draw_gui(state *UIState, channels, body ncurses.Window) {
  draw_channels(state, channels)
  switch state.SelectedBody {
  case BODY_CHANNEL:
    draw_body_channel(state, body)
  case BODY_DETAIL:
    draw_body_detail(state, body)
  case BODY_QUIT:
    draw_body_quit(state, body)
  }
}

func main_loop(client *pnyx.Client, audio_data chan []int16, window ncurses.Window, packet_chan chan pnyx.Payload, user_chan chan []byte, sigwinch_channel chan os.Signal) {
  channels := ncurses.NewWin(0, 0, 0, 0)
  body := ncurses.NewWin(0, 0, 0, 0)
  title := ncurses.NewWin(0, 0, 0, 0)
  state := &UIState{
    Peers: map[pnyx.PeerID]PeerState{
      [32]byte{}: {
        Name: "Test User 0",
      },
      [32]byte{0x01}: {
        Name: "Test User 1",
      },
      [32]byte{0x02}: {
        Name: "Test User 2",
      },
    },
    Channels: []ChannelState{
      {
        ID: 0,
        Name: "Channel 0",
        Members: []pnyx.PeerID{[32]byte{}, [32]byte{0x01}},
      },
      {
        ID: 1,
        Name: "Channel 1",
        Members: []pnyx.PeerID{[32]byte{0x02}},
      },
    },
  }

  ncurses.InitColor(1, 1000, 1000, 1000)
  ncurses.InitColor(2, 300, 300, 300)
  ncurses.InitColor(3, 150, 150, 150)
  ncurses.InitColor(4, 150, 150, 300)
  ncurses.InitPair(1, 1, 2)
  ncurses.InitPair(2, 1, 3)
  ncurses.InitPair(3, 1, 4)
  ncurses.ScrollOk(body, true)

  print_decoration(window, title, channels, body, client.Remote())
  draw_gui(state, channels, body)

  for client.Active() {
    select {
    case <-sigwinch_channel:
      ncurses.EndWin()
      ncurses.WRefresh(window)
      ncurses.WClear(window)
      print_decoration(window, title, channels, body, client.Remote())
      draw_gui(state, channels, body)
    case packet := <-packet_chan:
      switch packet := packet.(type) {
      case pnyx.ChannelCommandPacket:
        if packet.Channel == pnyx.ChannelID(0) {
          if packet.Mode == pnyx.MODE_AUDIO {
            if packet.Command == pnyx.AUDIO_SET_SAMPLE_RATE {
              var new_sample_rate int
              switch pnyx.SampleRate(packet.Data[0]) {
              case pnyx.SAMPLE_RATE_24KHZ:
                new_sample_rate = 24000
              case pnyx.SAMPLE_RATE_48KHZ:
                new_sample_rate = 48000
              default:
                continue
              }

              err := set_sample_rate(audio_data, new_sample_rate)
              if err != nil {
                panic(err)
              }
            }
          }
        }
      case pnyx.PeerPacket:
        if packet.Channel == pnyx.ChannelID(0) {
          decode_chan, exists := decoders[packet.Peer]
          if exists == false {
            if sample_rate != 0 {
              decode_chan = make(chan []byte, 1000)
              decoders[packet.Peer] = decode_chan
              go handle_peer_decode(audio_data, packet.Peer, decoders[packet.Peer], sample_rate)
              decode_chan <- packet.Data
            } else {
              decoders[packet.Peer] = nil
            }
          } else if decode_chan != nil {
            decode_chan <- packet.Data
          }
        }
      }
    case char := <-user_chan:
      handle_input(state, char)
      draw_gui(state, channels, body)
    }
  }
}

func process_mic(client *pnyx.Client, mic chan []byte) {
  for true {
    data := <-mic
    err := client.Send(pnyx.NewDataPacket(pnyx.ChannelID(0), pnyx.MODE_AUDIO, data))
    if err != nil {
      panic(err)
    }
  }
}

func main() {
  key_file_arg := flag.String("key", "${HOME}/.pnyx.key", "Path to the private key file to load/save")
  generate_key_arg := flag.Bool("genkey", false, "Set to generate a key if none exists")
  flag.Parse()

  var audio_data = make(chan []int16, 0)
  var speaker = make(chan []int16, 0)

  go mixer(audio_data, speaker)

  var mic = make(chan []byte, 0)
  ctx, outDevice, inDevice := setup_audio(mic, speaker)

  defer ctx.Free()
  defer ctx.Uninit()
  defer outDevice.Uninit()
  defer outDevice.Stop()
  defer inDevice.Uninit()
  defer inDevice.Stop()

  user_chan := make(chan []byte, 1024)
  packet_chan := make(chan pnyx.Payload, 1024)

  key := get_private_key(os.ExpandEnv(*key_file_arg), *generate_key_arg)
  client, err := pnyx.NewClient(key, flag.Arg(0), func(payload pnyx.Payload) error {
    select {
    case packet_chan <- payload:
      return nil
    default:
      return fmt.Errorf("Channel overflow")
    }
  })
  if err != nil {
    panic(err)
  }

  join_packet := pnyx.NewChannelCommandPacket(uuid.New(), pnyx.ChannelID(0), pnyx.MODE_CHANNEL, pnyx.CHANNEL_COMMAND_JOIN, nil)
  err = client.Send(join_packet)
  if err != nil {
    panic(err)
  }

  get_sample_rate_packet := pnyx.NewChannelCommandPacket(uuid.New(), pnyx.ChannelID(0), pnyx.MODE_AUDIO, pnyx.AUDIO_SET_SAMPLE_RATE, []byte{byte(pnyx.SAMPLE_RATE_48KHZ)})
  err = client.Send(get_sample_rate_packet)
  if err != nil {
    panic(err)
  }

  go process_mic(client, mic)

  ncurses.SetLocale(0, "")

  user_chan, stdin_active := ncurses.UTF8Listener(100, os.Stdin)
  window := ncurses.InitScr()
  ncurses.CursSet(0)
  ret := ncurses.StartColor()
  if ret != 0 {
    panic(ret)
  }

  sigwinch_channel := make(chan os.Signal, 1)
  signal.Notify(sigwinch_channel, syscall.SIGWINCH)

  go main_loop(client, audio_data, window, packet_chan, user_chan, sigwinch_channel)

  os_sigs := make(chan os.Signal, 1)
  signal.Notify(os_sigs, syscall.SIGINT, syscall.SIGABRT)

  <-os_sigs
  stdin_active.Store(false)
  client.Close()
  ncurses.EndWin()
}