.include "skirmish.inc"

.export INIT_GAME
.export UPDATE_GAME
.export REBUILD_OCCUPANCY
.export UNIT_ATTACK
.export ENSURE_SELECTED_UNIT

.import ENEMY_AI_TURN

.importzp input_flags
.importzp game_state
.importzp selected_unit
.importzp turn_counter
.importzp player_units
.importzp enemy_units
.importzp move_dx
.importzp move_dy
.importzp action_result
.importzp temp0
.importzp temp1
.importzp temp2
.importzp temp3
.importzp temp4
.importzp temp5

.import unit_x
.import unit_y
.import unit_hp
.import unit_team
.import grid_occupancy

.segment "CODE"
INIT_GAME:
    lda #STATE_PLAYER
    sta game_state
    lda #0
    sta turn_counter
    sta action_result
    lda #0
    sta selected_unit

    lda #PLAYER_COUNT
    sta player_units
    lda #ENEMY_COUNT
    sta enemy_units

    ldx #0
clear_units:
    lda #0
    sta unit_x,x
    sta unit_y,x
    sta unit_hp,x
    sta unit_team,x
    inx
    cpx #UNIT_COUNT
    bne clear_units

    lda #0
    sta unit_x+0
    sta unit_x+1
    sta unit_x+2
    lda #7
    sta unit_x+3
    sta unit_x+4
    sta unit_x+5

    lda #1
    sta unit_y+0
    lda #3
    sta unit_y+1
    lda #5
    sta unit_y+2
    lda #1
    sta unit_y+3
    lda #3
    sta unit_y+4
    lda #5
    sta unit_y+5

    lda #3
    ldx #0
hp_loop:
    sta unit_hp,x
    inx
    cpx #UNIT_COUNT
    bne hp_loop

    lda #TEAM_PLAYER
    sta unit_team+0
    sta unit_team+1
    sta unit_team+2
    lda #TEAM_ENEMY
    sta unit_team+3
    sta unit_team+4
    sta unit_team+5

    jsr REBUILD_OCCUPANCY
    rts

UPDATE_GAME:
    jsr REBUILD_OCCUPANCY

    lda game_state
    cmp #STATE_PLAYER
    beq update_player
    cmp #STATE_ENEMY
    beq update_enemy
    cmp #STATE_CHECK
    beq update_check
    jmp update_end

update_player:
    jsr ENSURE_SELECTED_UNIT

    lda input_flags
    and #INPUT_NEXT
    beq :+
    jsr CYCLE_SELECTED_UNIT
    rts
:
    lda input_flags
    and #INPUT_FIRE
    beq :+
    jsr CYCLE_SELECTED_UNIT
    rts
:
    lda input_flags
    and #INPUT_WAIT
    beq :+
    lda #STATE_ENEMY
    sta game_state
    rts
:
    lda #0
    sta move_dx
    sta move_dy

    lda input_flags
    and #INPUT_UP
    beq :+
    lda #$FF
    sta move_dy
    jmp do_player_action
:
    lda input_flags
    and #INPUT_DOWN
    beq :+
    lda #1
    sta move_dy
    jmp do_player_action
:
    lda input_flags
    and #INPUT_LEFT
    beq :+
    lda #$FF
    sta move_dx
    jmp do_player_action
:
    lda input_flags
    and #INPUT_RIGHT
    beq :+
    lda #1
    sta move_dx
    jmp do_player_action
:
    rts

do_player_action:
    jsr UNIT_TRY_ACTION
    lda action_result
    beq :+
    lda #STATE_ENEMY
    sta game_state
:
    rts

update_enemy:
    jsr ENEMY_AI_TURN
    inc turn_counter
    lda #STATE_CHECK
    sta game_state
    rts

update_check:
    jsr CHECK_WIN_LOSE
    lda game_state
    cmp #STATE_CHECK
    bne :+
    lda #STATE_PLAYER
    sta game_state
:
    rts

update_end:
    lda input_flags
    and #INPUT_RESTART
    beq :+
    jsr INIT_GAME
:
    rts

UNIT_TRY_ACTION:
    lda #0
    sta action_result

    ldx selected_unit
    cpx #PLAYER_COUNT
    bcs unit_try_done
    lda unit_hp,x
    beq unit_try_done

    lda unit_x,x
    clc
    adc move_dx
    sta temp0
    cmp #8
    bcs unit_try_done

    lda unit_y,x
    clc
    adc move_dy
    sta temp1
    cmp #8
    bcs unit_try_done

    lda temp1
    asl
    asl
    asl
    clc
    adc temp0
    tay
    lda grid_occupancy,y
    cmp #EMPTY_CELL
    beq move_player

    tay
    lda unit_team,y
    cmp #TEAM_ENEMY
    bne unit_try_done

    ldx selected_unit
    jsr UNIT_ATTACK_XY
    lda #1
    sta action_result
    jmp unit_try_done

move_player:
    ldx selected_unit
    lda temp0
    sta unit_x,x
    lda temp1
    sta unit_y,x
    lda #1
    sta action_result

unit_try_done:
    rts

REBUILD_OCCUPANCY:
    ldx #0
clear_grid:
    lda #EMPTY_CELL
    sta grid_occupancy,x
    inx
    cpx #64
    bne clear_grid

    ldx #0
grid_loop:
    lda unit_hp,x
    beq next_grid_unit
    lda unit_y,x
    asl
    asl
    asl
    clc
    adc unit_x,x
    tay
    txa
    sta grid_occupancy,y
next_grid_unit:
    inx
    cpx #UNIT_COUNT
    bne grid_loop
    rts

CHECK_WIN_LOSE:
    lda enemy_units
    bne :+
    lda #STATE_WIN
    sta game_state
    rts
:
    lda player_units
    bne :+
    lda #STATE_LOSE
    sta game_state
    rts
:
    lda #STATE_CHECK
    sta game_state
    rts

ENSURE_SELECTED_UNIT:
    lda selected_unit
    cmp #PLAYER_COUNT
    bcs choose_first_player
    tax
    lda unit_hp,x
    bne ensure_done

choose_first_player:
    ldx #0
find_alive_player:
    lda unit_hp,x
    bne found_player
    inx
    cpx #PLAYER_COUNT
    bne find_alive_player
    lda #$FF
    sta selected_unit
    rts

found_player:
    stx selected_unit
ensure_done:
    rts

CYCLE_SELECTED_UNIT:
    lda player_units
    beq no_cycle_target

    lda selected_unit
    cmp #PLAYER_COUNT
    bcc :+
    lda #$FF
:
    clc
    adc #1
    cmp #PLAYER_COUNT
    bcc :+
    lda #0
:
    sta temp0
    ldx #PLAYER_COUNT
cycle_search:
    ldy temp0
    lda unit_hp,y
    bne cycle_found
    inc temp0
    lda temp0
    cmp #PLAYER_COUNT
    bcc :+
    lda #0
    sta temp0
:
    dex
    bne cycle_search

no_cycle_target:
    lda #$FF
    sta selected_unit
    rts

cycle_found:
    sty selected_unit
    rts

UNIT_ATTACK:
    jsr UNIT_ATTACK_XY
    rts

UNIT_ATTACK_XY:
    lda unit_hp,y
    beq attack_done
    sec
    sbc #1
    sta unit_hp,y
    bne attack_done
    lda unit_team,y
    beq dead_player
    dec enemy_units
    jmp attack_done

dead_player:
    dec player_units

attack_done:
    rts