.include "ps2_keyboard.h65"
.include "string.h65"
.include "lcd.h65"
.include "parity.h65"
Export ps2kb, init, begin_receive, scancode, status, scancode_handler
Export ps2kb, _receive_irq_shift_reg_handler, _receive_irq_timer_handler, _send_byte, _send_irq_shift_reg_handler, _send_irq_timer_handler,
Export ps2kb, send_command, send_cmd, send_data, cmd_response, response_length, FMT_CMD_FAIL

.bss
status:                 .res 1  ; current status
prev_status:            .res 1  ; status before sending command
send_last_bits:         .res 1  ; last bits to load after 8 bits were shifted out
send_cmd:               .res 1  ; command to send/last sent
send_data:              .res 1  ; data byte to send/last sent or NO_DATA
expect_data_length:     .res 1  ; number of data bytes to expect from keyboard
response_length:        .res 1  ; number of response bytes from keyboard
cmd_response:           .res 3  ; responses from keyboard
key_read:               .res 2  ; first 8 bits, last 3 bits from scancode
scancode:               .res 1  ; last received scancode
scancode_handler:       .res 2  ; pointer to a function that handles new scancodes

.code

; these are macros the save time during the interrupt handler (no jsr, rts)
;;********************************************************************************
;; @macro Enable the clock signal from the keyboard
;; @modifies: A
;; @details
;;  Stop pulling the keyboards clock low
;;********************************************************************************
.macro _EnableClock
    ; set pin to input
    lda ps2kb::VIA + ps2kb::PULL_DDR
    and #<~ps2kb::PULL_MASK_CLK
    sta ps2kb::VIA + ps2kb::PULL_DDR
.endmacro

;;********************************************************************************
;; @macro Disable the clock signal from the keyboard
;; @modifies: A
;; @details
;;  Pulls the keyboards clock low
;;********************************************************************************
.macro _DisableClock
    ; set pin to output
    lda ps2kb::VIA + ps2kb::PULL_DDR
    ora #ps2kb::PULL_MASK_CLK
    sta ps2kb::VIA + ps2kb::PULL_DDR
    ; set pin low
    lda ps2kb::VIA + ps2kb::PULL_REG
    and #<~ps2kb::PULL_MASK_CLK
    sta ps2kb::VIA + ps2kb::PULL_REG
.endmacro


;;********************************************************************************
;; @macro Stop pulling the keyboard data pin low
;; @modifies: A
;;********************************************************************************
.macro _StopPullDataLow
    ; set pin to input
    lda ps2kb::VIA + ps2kb::PULL_DDR
    and #<~ps2kb::PULL_MASK_DAT
    sta ps2kb::VIA + ps2kb::PULL_DDR
.endmacro

;;********************************************************************************
;; @macro Pull the keyboard data pin low
;; @modifies: A
;;********************************************************************************
.macro _PullDataLow
    ; set pin to output
    lda ps2kb::VIA + ps2kb::PULL_DDR
    ora #ps2kb::PULL_MASK_DAT
    sta ps2kb::VIA + ps2kb::PULL_DDR
    ; set pin low
    lda ps2kb::VIA + ps2kb::PULL_REG
    and #<~ps2kb::PULL_MASK_DAT
    sta ps2kb::VIA + ps2kb::PULL_REG
.endmacro


.proc init
    _DisableClock
    stz key_read 
    stz key_read+1
    stz scancode
    stz status
    ; T2 oneshot mode, clear shift reg
    lda ps2kb::VIA + IO::ACR
    and #<~IO::ACR_MASK::SR
    ora #IO::ACR::T2_IRQ_LOAD
    sta ps2kb::VIA + IO::ACR
    ; load this rts as scancode_handler
    StoreDByte @rts, scancode_handler
@rts:
    rts
.endproc


;;********************************************************************************
;; @function Start receiving data from the keyboard
;; @details
;;  - use the shift register interrupts to read the first 8 bits
;;  - configure timer for timing the read of the last 3 bits
;; @modifies: A
;;********************************************************************************
.proc begin_receive
    ; disable timer interrupts (this might be called while waiting on the last 3 bits)
    IO_DisableIRQ ps2kb::VIA, IO::IRQ::T2
    ; (re)set shift register to shift in under external clock on CB1
    lda ps2kb::VIA + IO::ACR
    and #<~IO::ACR_MASK::SR
    sta ps2kb::VIA + IO::ACR    ; important, disabling SR resets the count
    ora #IO::ACR::SR_SIN_PHIE
    sta ps2kb::VIA + IO::ACR
    stz key_read
    stz key_read+1
    stz scancode

    ; the 3 last bits take about 230us, at @1MHz => wait 230 cycles and then the shift register
    ; (this could be shorter since the it takes a few cycles after the interrupt)
    lda #<ps2kb::TIMER_RECV
    sta ps2kb::VIA + IO::T2CL

    ; set to RECEIVE_KEYS only if RECEIVE_ANSWER is not set
    bit status
    bvs @status_set
    lda #ps2kb::STATUS::RECEIVE_KEYS
    sta status
@status_set:
    ; todo setup irq handlers
    StoreDByte _receive_irq_shift_reg_handler, $3000
    StoreDByte _receive_irq_timer_handler, $3060

    _EnableClock

    IO_EnableIRQ ps2kb::VIA, IO::IRQ::SR
    ; reset and start SR
    stz ps2kb::VIA + IO::SR
    rts
.endproc


;;********************************************************************************
;; @function Read the first 8 bits an
;; @modifies: A
;; @details
;;  - read shift register
;;  - disable shift register interrupts
;;  - reset shift register
;;  - enable timer 2 interrupts
;;  - start timer 2
;;  IO::SR has to be read before the next bit is shifted in, which happens ~75us after the irq
;;  at 1MHz, handling this interrupt takes about 50us (without any additional debug code),
;;  so it should work
;;********************************************************************************
.proc _receive_irq_shift_reg_handler
    lda ps2kb::VIA + IO::SR
    sta key_read
    stz ps2kb::VIA + IO::SR

    IO_DisableIRQ ps2kb::VIA, IO::IRQ::SR
    IO_EnableIRQ ps2kb::VIA, IO::IRQ::T2
    ; start timer, low order count already in latch after begin_receive
    lda #>ps2kb::TIMER_RECV
    sta ps2kb::VIA + IO::T2CH
    rts
.endproc


;;********************************************************************************
;; @function Read the last 3 bits after after timer 2 is up
;; @modifies: A
;; @details
;;  - read shift register
;;  - disable timer 2 interrupts
;;  - enable shift register interrupts
;;  - reset shift register
;;  - jump to the subroutine pointed to by scancode_handler
;;********************************************************************************
.proc _receive_irq_timer_handler
    lda ps2kb::VIA + IO::SR
    sta key_read+1

    lda ps2kb::VIA + IO::T2CL  ; clear interrupt flag

    IO_DisableIRQ ps2kb::VIA, IO::IRQ::T2
    ; reset SR
    ; disabling shifting in acr seems necessary to reset - otherwise
    ; it continues counting and interrupts after the first 5 bits of the next keypress
    lda #(IO::ACR::SR_SIN_PHIE | IO::ACR::T2_IRQ_LOAD)
    trb ps2kb::VIA + IO::ACR
    tsb ps2kb::VIA + IO::ACR
    stz ps2kb::VIA + IO::SR
    IO_EnableIRQ ps2kb::VIA, IO::IRQ::SR

    lda key_read+1
    ror             ; stop bit -> C
    sta key_read+1  ; parity in bit 0, D7 in bit 1
    ror             ; parity -> C
    ror             ; D7 -> C
    lda key_read  ; not affecting carry
    rol             ; C  -> bit 0, startbit -> C
    Reverse A
    sta scancode

    ; check parity
    CalculateOddParity
    eor key_read+1  ; if bit 0 is 1 - parity error
    and #1          ; bit 1 is still D7
    bne @parity_error
@handle_scancode:

    ; check what to do with the scancode
    bit status
    bpl @status_not_receive_keys    ; RECEIVE_KEYS
    jmp (scancode_handler)
@status_not_receive_keys:
    bvc @status_dont_handle         ; RECEIVE_ANSWER
    jmp _process_cmd_answer 
@status_dont_handle:
    rts
@parity_error:      ; TODO handle somehow
    lda #$fe        ; fe means error/resend for commands
    sta scancode
    DEBUG_LED_ON 2
    bra @handle_scancode
.endproc


;;********************************************************************************
;; @function Send a command to the keyboard
;; @details
;;  No checks are done on the validity of the command and data byte.
;;  You can send anything you want to the keyboard.
;;  `send_command` will return immediately, as sending the command is asynchronous.
;;  The command is done if `ps2kb::status & ps2kb::STATUS::SEND == 0`.
;;  @section Response
;;   The answers will be in the byte array ps2kb::cmd_response.
;;   If the command did not have a data byte, cmd_response is:
;;       - 0: command response
;;       - 1: NO_RESPONSE or additional response byte (identify keyboard command only)
;;       - 2: NO_RESPONSE or additional response byte (identify keyboard command only)

;;   If the command did have a data byte, cmd_response is:
;;       - 0: 0xFA (ACK) or 0xFE (Resend/Error)
;;       - 1: 0xFA (ACK) or 0xFE (Resend/Error)
;;       - 2: NO_RESPONSE or additional response byte
;;
;;   If ANY of the bytes in cmd_response is 0xFE, the command failed.
;;  
;; @modifies: A,X,Y
;; @param A: The command byte
;; @param X: The data byte or NO_DATA if only the command byte should be sent
;; @param Y: The number of data bytes expected to receive. Must be one of {0, 1, 2}
;;********************************************************************************
.proc send_command
    sty expect_data_length
    stx send_data
    stz response_length
    ldy #ps2kb::NO_RESPONSE
    sty cmd_response
    sty cmd_response+1
    sty cmd_response+2
    ldy status
    sty prev_status
    ldy #ps2kb::STATUS::SEND_CMD
    sty status
    sta send_cmd        ; store if it needs to be resent
    jmp _send_byte
.endproc


;;********************************************************************************
;; @function Send a byte to the keyboard
;; @modifies: A,X,Y
;; @param A: The byte to send
;; @details
;;  - pull clock low to stop keyboard transmissions
;;  - setup SR to shift out
;;  - pull data low (by shifting out a zero)
;;  - reset SR
;;  - put (reversed) byte in SR
;;  - put parity + stopbit(s) in send_last_bits variable
;;  - setup interrupt handlers
;;  - enable clock
;;********************************************************************************
.proc _send_byte
    pha
        IO_DisableIRQ ps2kb::VIA, (IO::IRQ::T2 | IO::IRQ::SR)

        _DisableClock

        ; set shift register to output
        lda ps2kb::VIA + IO::ACR
        and #<~IO::ACR_MASK::SR
        ora #IO::ACR::SR_SOUT_PHIE
        sta ps2kb::VIA + IO::ACR
        stz ps2kb::VIA + IO::SR
        ; shift out the startbit = data low
        _EnableClock    
        _DisableClock
        ; reset shift register to start the count at 0 again
        lda #IO::ACR::SR_SOUT_PHIE
        trb ps2kb::VIA + IO::ACR
        tsb ps2kb::VIA + IO::ACR

        ; setup the low timer byte
        lda #<ps2kb::TIMER_SEND
        sta ps2kb::VIA + IO::T2CL

        ; setup interrupt handlers
        StoreDByte _send_irq_shift_reg_handler, $3000
        StoreDByte _send_irq_timer_handler, $3060
    pla
    Reverse A
    pha
        CalculateOddParity
        ror     ; Parity -> C
        lda #$ff
        ror     ; Parity -> bit 7, rest = 1
        sta send_last_bits  ; loaded into SR by irq handler
    pla
    sta ps2kb::VIA + IO::SR
    IO_EnableIRQ ps2kb::VIA, IO::IRQ::SR
    _EnableClock    
    rts
.endproc


;;********************************************************************************
;; @function Send the lasts 3 bits 
;; @modifies: A
;; @details
;;  - load the remaining 3 bits of the transmission into the shift register
;;  - disable shift register interrupts
;;  - enable timer 2 interrupts
;;  - start timer 2
;;********************************************************************************
.proc _send_irq_shift_reg_handler
    lda send_last_bits
    sta ps2kb::VIA + IO::SR

    IO_DisableIRQ ps2kb::VIA, IO::IRQ::SR
    IO_EnableIRQ ps2kb::VIA, IO::IRQ::T2
    ; start timer, low order count already in latch after init
    lda #>ps2kb::TIMER_SEND
    sta ps2kb::VIA + IO::T2CH
    rts
.endproc


;;********************************************************************************
;; @function Setup VIA to receive the keyboard's answer
;; @modifies: A
;; @details
;;  - disable timer 2 interrupts
;;  - pull clock low
;;  - or status with RECEIVE_ANSWER
;;  - begin receive
;;********************************************************************************
.proc _send_irq_timer_handler
    lda ps2kb::VIA + IO::T2CL   ; clear interrupt flag
    IO_DisableIRQ ps2kb::VIA, IO::IRQ::T2

    ; disable shift register
    lda ps2kb::VIA + IO::ACR
    and #<~IO::ACR_MASK::SR
    sta ps2kb::VIA + IO::ACR

    _DisableClock               ; disable keyboard while setting up receive
    lda status
    ora #ps2kb::STATUS::RECEIVE_ANSWER
    sta status
    jmp ps2kb::begin_receive
.endproc


;;********************************************************************************
;; @function Process the response of a command
;; @details
;;  Stores the answer in the `ps2kb::cmd_response` array.
;;  If the response is $FE or command related transmissions are done (see below),
;;  no further data will be sent or received and the `ps2kb::status` takes the previous value.
;; 
;;  - store response
;;  - if response == 0xFE -> stop
;;  - if status == SEND_CMD -> send databyte
;;  - if statusk
;;  You can send anything you want to the keyboard.
;;  `send_command` will return immediately, as sending the command is asynchronous.
;;  The command is done if `ps2kb::status & ps2kb::STATUS::SEND == 0`.
;;  @section Response
;;   The answers will be in the byte array ps2kb::cmd_response.
;;   If the command did not have a data byte, cmd_response is:
;;       - 0: command response
;;       - 1: NO_RESPONSE or additional response byte (identify keyboard command only)
;;       - 2: NO_RESPONSE or additional response byte (identify keyboard command only)

;;   If the command did have a data byte, cmd_response is:
;;       - 0: 0xFA (ACK) or 0xFE (Resend/Error)
;;       - 1: 0xFA (ACK) or 0xFE (Resend/Error)
;;       - 2: NO_RESPONSE or additional response byte
;;
;;   If ANY of the bytes in cmd_response is 0xFE, the command failed.
;;  
;; @modifies: A,X,Y
;; @param A: The command byte
;; @param X: The data byte or NO_DATA if only the command byte should be sent
;; @param Y: The number of data bytes expected to receive. Must be one of {0, 1, 2}
;;********************************************************************************
.proc _process_cmd_answer 
@store_response:
    ldx response_length
    lda scancode
    sta cmd_response,x
    inx
    stx response_length
    stz scancode
    ; check for resend
    cmp #$fe
    beq @cmd_fail
    ; received something useful, check state
    lda status
    bit #ps2kb::STATUS::SEND_CMD
    bne @cmd_sent
    bit #ps2kb::STATUS::SEND_DATA
    bne @everything_sent
    ; status must be SEND_RECV

@receive_data_response:
    dec expect_data_length
    bmi @cmd_done           ; no more expected data
    rts                     ; wait for another data byte

@cmd_sent:
    ; check if a data byte needs to be sent
    lda send_data
    beq @send_data          ; if 0 (would fall through the cmp check TODO would it??)
    cmp #ps2kb::NO_DATA
    beq @everything_sent    ; if not NO_DATA
@send_data:
    lda #ps2kb::STATUS::SEND_DATA
    sta status
    lda send_data
    jmp ps2kb::_send_byte          ; send the data

@everything_sent:           ; check if additonal bytes are expected
    dec expect_data_length
    bmi @cmd_done           ; expect_data_length was 0
    lda #(ps2kb::STATUS::SEND_RECV | ps2kb::STATUS::RECEIVE_ANSWER)
    sta status
    rts

@cmd_fail:                  ; keyboard wont expect data byte/send an additional response byte
@cmd_done:
    ; restore previous status
    lda prev_status
    sta status
    ; disable the clock if the previous was not RECEIVE_KEYS
    cmp #ps2kb::STATUS::RECEIVE_KEYS
    beq @rts
    _DisableClock
@rts:
    rts
.endproc

.rodata
;; Can be used by other programs
FMT_CMD_FAIL: .asciiz "PS/2 Cmd fail:  %x-%x > %x%x%x  "