package main

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

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

var decoders = map[pnyx.PeerID]chan[]byte{}
var encoder *opus.Encoder
var sample_rate int = 0
var mic = make(chan []byte, 0)
var speaker = make(chan []int16, 0)
var audio_data = make(chan []int16, 0)

func set_sample_rate(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(peer_id, decoders[peer_id], sample_rate)
  }
  return nil
}

func handle_peer_decode(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 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()

  ctx, err := malgo.InitContext(nil, malgo.ContextConfig{}, nil)
  if err != nil {
    panic(err)
  }

  go mixer(audio_data, speaker)

  defer ctx.Free()
  defer ctx.Uninit()

	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)
  }

  defer outDevice.Uninit()
  defer outDevice.Stop()

	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)
	}

  defer inDevice.Uninit()
  defer inDevice.Stop()

  var key ed25519.PrivateKey = nil

  key_file_path := os.ExpandEnv(*key_file_arg)
  key_file_bytes, err := os.ReadFile(key_file_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")
    }
  } else if *generate_key_arg {
    _, 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(key_file_path, key_pem, 0o600)
    if err != nil {
      panic(err)
    }
  }

  client, err := pnyx.NewClient(key, flag.Arg(0))
  if err != nil {
    panic(err)
  }
  
  go func() {
    var buf [1024]byte
    for true {
      read, _, err := client.Connection.ReadFromUDP(buf[:])
      if err != nil {
        break
      }

      data, err := pnyx.ParseSessionData(&client.Session, buf[pnyx.COMMAND_LENGTH + pnyx.SESSION_ID_LENGTH:read])
      if err != nil {
        continue
      }

      packet, err := pnyx.ParsePacket(data)
      if err != nil {
        continue
      }

      switch packet := packet.(type) {
      case pnyx.PingPacket:
      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(new_sample_rate)
              if err != nil {
                panic(err)
              }
            }
          }
        }
      case pnyx.CommandPacket:
      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(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
          }
        }
      default:
      }
    }
  }()

  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 func(){
    for true {
      data := <- mic
      err = client.Send(pnyx.NewDataPacket(pnyx.ChannelID(0), pnyx.MODE_AUDIO, data))
      if err != nil {
        panic(err)
      }
    }
  }()

  window := ncurses.Init()
  active := atomic.Bool{}
  active.Store(true)

  go func() {
    ncurses.ColorPair(1).Init(ncurses.ColorBlue, ncurses.ColorRed)
    window.AddStr("pnyx client")

    for active.Load() {
      window.Refresh()
      time.Sleep(200*time.Millisecond)
      peers := make([]pnyx.PeerID, 0, len(decoders))
      for peer_id := range(decoders) {
        peers = append(peers, peer_id)
      }

      slices.SortFunc(peers, func(a, b pnyx.PeerID) int {
        return slices.Compare(a[:], b[:])
      })

      for i, peer_id := range(peers) {
        window.MvAddStr(i+1, 0, fmt.Sprintf("%x", peer_id))
      }
    }
  }()

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

  <-os_sigs
  active.Store(false)
  ncurses.EndWin()
}