#include "atmega328p_def_inc.S"

.org 0
   rjmp main

 main:
    ; reset system status
    ldi r16,0
    out SREG,r16

    ; init stack pointer
    ldi r16,lo8(RAMEND)
    out SPL,r16
    ldi r16,hi8(RAMEND)
    out SPH,r16

    rcall uart_init
    rjmp mainloop

;-----------------------------------------------------------
; UART
;-----------------------------------------------------------
uart_init:
  ; set baud rate to 1 Mbps
  ldi r16, 1<<U2X0
  sts UCSR0A, r16
  ldi r16, 0
  sts UBRR0H, r16
  ldi r16, 1
  sts UBRR0L, r16

  ; enable receive and transmit, without interrupts
  ldi r16, (1<<RXEN0)|(1<<TXEN0)
  sts UCSR0B,r16

  ; frame format: 8 data bits, 1 stop bit, no parity.
  ldi r16, (1<<UCSZ01)|(1<<UCSZ00)
  sts UCSR0C,r16

  ret

; sends one byte over the serial line.
; data must reside in r17.
uart_send_byte:
  ; wait for send register to become empty
  lds r16, UCSR0A
  sbrs r16, UDRE0
  rjmp uart_send_byte

  ; write data to transmit register
  sts UDR0, R17

  ; wait for send complete flag (MAY not be necessary)
;_wait_send_flag:
;  lds r16, UCSR0A
;  sbrs r16, TXC0
;  rjmp _wait_send_flag

  ret

; receives one byte from the serial line.
; data received is moved to r17.
uart_receive_byte:
  ; wait for receive complete flag
  lds r16, UCSR0A
  sbrs r16, RXC0
  rjmp uart_receive_byte

  ; write data to r17
  lds r17, UDR0

  ret



;-----------------------------------------------------------
; main loop
;-----------------------------------------------------------
#define rcv_count r18
#define snd_count r19

.equ RCV_BUF_ADDR, 512
.equ RCV_BUF_MAX_LEN, 64

mainloop:
  ; load buffer base addr into pointer register
  ldi r30, lo8(RCV_BUF_ADDR)
  ldi r31, hi8(RCV_BUF_ADDR)

  ; reset receive and send counts
  ldi rcv_count, 0
  ldi snd_count, 0

rcv_byte:
  rcall uart_receive_byte         ; fetch next byte into r17
  cpi r17, '\n'                   ; echo if newline received
  breq echo
  cpi rcv_count, RCV_BUF_MAX_LEN  ; echo if buf_max_len bytes received
  breq echo

  inc rcv_count
  st Z+, r17       ; update receive buffer
  rjmp rcv_byte


echo:
  ; load buffer base addr into pointer register
  ldi r30, lo8(RCV_BUF_ADDR)
  ldi r31, hi8(RCV_BUF_ADDR)

snd_byte:
  ld r17, Z+            ; load what pointer register points to into r17 and increment pointer
  rcall uart_send_byte  ; send byte in r17
  inc snd_count

  cp snd_count, rcv_count ; compare send/receive amounts
  breq mainloop           ; if they are equal goto mainloop
  rjmp snd_byte           ; if not, send next byte