// Copyright 2021 Howard Lau
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package bus

import chisel3._
import chisel3.experimental.ChiselEnum
import chisel3.util._
import riscv.Parameters

object AXI4Lite {
  val protWidth = 3
  val respWidth = 2
}

class AXI4LiteWriteAddressChannel(addrWidth: Int) extends Bundle {

  val AWVALID = Output(Bool())
  val AWREADY = Input(Bool())
  val AWADDR = Output(UInt(addrWidth.W))
  val AWPROT = Output(UInt(AXI4Lite.protWidth.W))

}

class AXI4LiteWriteDataChannel(dataWidth: Int) extends Bundle {
  val WVALID = Output(Bool())
  val WREADY = Input(Bool())
  val WDATA = Output(UInt(dataWidth.W))
  val WSTRB = Output(UInt((dataWidth / 8).W))
}

class AXI4LiteWriteResponseChannel extends Bundle {
  val BVALID = Input(Bool())
  val BREADY = Output(Bool())
  val BRESP = Input(UInt(AXI4Lite.respWidth.W))
}

class AXI4LiteReadAddressChannel(addrWidth: Int) extends Bundle {
  val ARVALID = Output(Bool())
  val ARREADY = Input(Bool())
  val ARADDR = Output(UInt(addrWidth.W))
  val ARPROT = Output(UInt(AXI4Lite.protWidth.W))
}

class AXI4LiteReadDataChannel(dataWidth: Int) extends Bundle {
  val RVALID = Input(Bool())
  val RREADY = Output(Bool())
  val RDATA = Input(UInt(dataWidth.W))
  val RRESP = Input(UInt(AXI4Lite.respWidth.W))
}

class AXI4LiteInterface(addrWidth: Int, dataWidth: Int) extends Bundle {
    val AWVALID = Output(Bool())
    val AWREADY = Input(Bool())
    val AWADDR = Output(UInt(addrWidth.W))
    val AWPROT = Output(UInt(AXI4Lite.protWidth.W))
    val WVALID = Output(Bool())
    val WREADY = Input(Bool())
    val WDATA = Output(UInt(dataWidth.W))
    val WSTRB = Output(UInt((dataWidth / 8).W))
    val BVALID = Input(Bool())
    val BREADY = Output(Bool())
    val BRESP = Input(UInt(AXI4Lite.respWidth.W))
    val ARVALID = Output(Bool())
    val ARREADY = Input(Bool())
    val ARADDR = Output(UInt(addrWidth.W))
    val ARPROT = Output(UInt(AXI4Lite.protWidth.W))
    val RVALID = Input(Bool())
    val RREADY = Output(Bool())
    val RDATA = Input(UInt(dataWidth.W))
    val RRESP = Input(UInt(AXI4Lite.respWidth.W))
}

class AXI4LiteChannels(addrWidth: Int, dataWidth: Int) extends Bundle {
  val write_address_channel = new AXI4LiteWriteAddressChannel(addrWidth)
  val write_data_channel = new AXI4LiteWriteDataChannel(dataWidth)
  val write_response_channel = new AXI4LiteWriteResponseChannel()
  val read_address_channel = new AXI4LiteReadAddressChannel(addrWidth)
  val read_data_channel = new AXI4LiteReadDataChannel(dataWidth)
}

class AXI4LiteSlaveBundle(addrWidth: Int, dataWidth: Int) extends Bundle {
  val read = Output(Bool())                 // tell slave device to read
  val write = Output(Bool())                // tell slave device to write
  val read_data = Input(UInt(dataWidth.W))  // data read from slave device
  val read_valid = Input(Bool())            // indicates if read_data is valid
  val write_data = Output(UInt(dataWidth.W))
  val write_strobe = Output(Vec(Parameters.WordSize, Bool()))
  val address = Output(UInt(addrWidth.W))
}

class AXI4LiteMasterBundle(addrWidth: Int, dataWidth: Int) extends Bundle {
  val read = Input(Bool())                  // request a read transaction
  val write = Input(Bool())                 // request a write transaction
  val read_data = Output(UInt(dataWidth.W))
  val write_data = Input(UInt(dataWidth.W))
  val write_strobe = Input(Vec(Parameters.WordSize, Bool()))
  val address = Input(UInt(addrWidth.W))

  val busy = Output(Bool())                 // if busy, master is not ready to accept new transactions
  val read_valid = Output(Bool())           // indicates read transaction done successfully and asserts for ONLY 1 cycle.
  val write_valid = Output(Bool())          // indicates write transaction done successfully and asserts for ONLY 1 cycle.
}

object AXI4LiteStates extends ChiselEnum {
  val Idle, ReadAddr, ReadDataWait, ReadData, WriteAddr, WriteData, WriteResp = Value
}

class AXI4LiteSlave(addrWidth: Int, dataWidth: Int) extends Module {
  val io = IO(new Bundle {
    val channels = Flipped(new AXI4LiteChannels(addrWidth, dataWidth))
    val bundle = new AXI4LiteSlaveBundle(addrWidth, dataWidth)
  })
  val state = RegInit(AXI4LiteStates.Idle)
  val addr = RegInit(0.U(dataWidth.W))
  io.bundle.address := addr
  val read = RegInit(false.B)
  io.bundle.read := read
  val write = RegInit(false.B)
  io.bundle.write := write
  val write_data = RegInit(0.U(dataWidth.W))
  io.bundle.write_data := write_data
  val write_strobe = RegInit(VecInit(Seq.fill(Parameters.WordSize)(false.B)))
  io.bundle.write_strobe := write_strobe


  val ARREADY = RegInit(false.B)
  io.channels.read_address_channel.ARREADY := ARREADY
  val RVALID = RegInit(false.B)
  io.channels.read_data_channel.RVALID := RVALID
  val RRESP = RegInit(0.U(AXI4Lite.respWidth))
  io.channels.read_data_channel.RRESP := RRESP

  val AWREADY = RegInit(false.B)
  io.channels.write_address_channel.AWREADY := AWREADY
  val WREADY = RegInit(false.B)
  io.channels.write_data_channel.WREADY := WREADY
  write_data := io.channels.write_data_channel.WDATA
  val BVALID = RegInit(false.B)
  io.channels.write_response_channel.BVALID := BVALID
  val BRESP = WireInit(0.U(AXI4Lite.respWidth))
  io.channels.write_response_channel.BRESP := BRESP
  //lab4(BUS)
  // 读数据寄存器，保证 RDATA 在 RVALID 为高时稳定
  val rdataReg = RegInit(0.U(dataWidth.W))
  io.channels.read_data_channel.RDATA := rdataReg

  switch(state) {
    is(AXI4LiteStates.Idle) {
      // 默认信号
      ARREADY := false.B
      RVALID := false.B
      AWREADY := false.B
      WREADY := false.B
      BVALID := false.B
      read := false.B
      write := false.B



      when(io.channels.read_address_channel.ARVALID) {
        addr := io.channels.read_address_channel.ARADDR
        ARREADY := true.B
        state := AXI4LiteStates.ReadAddr
      }.elsewhen(io.channels.write_address_channel.AWVALID) {
        addr := io.channels.write_address_channel.AWADDR
        AWREADY := true.B
        state := AXI4LiteStates.WriteAddr
      }
    }

    is(AXI4LiteStates.ReadAddr) {
      // 已经完成 ARVALID/ARREADY 握手，将地址送给从设备并发起一次读
      ARREADY := false.B
      read := true.B
      state := AXI4LiteStates.ReadData
    }

    is(AXI4LiteStates.ReadData) {
      // 保持对设备的读请求，等待设备返回 read_valid
      read := true.B

      when(io.bundle.read_valid && !RVALID) {
        // 设备已经准备好数据，将其锁存并让RVALID为1
        rdataReg := io.bundle.read_data
        RVALID := true.B
      }.elsewhen(RVALID && io.channels.read_data_channel.RREADY) {
        // 主机已经接收数据，完成一次读事务
        RVALID := false.B
        read := false.B
        state := AXI4LiteStates.Idle
      }
    }

    is(AXI4LiteStates.WriteAddr) {
      // 写地址握手结束，等待写数据
      AWREADY := false.B
      when(io.channels.write_data_channel.WVALID) {
        // 收到写数据，锁存写数据及写选通
        write_data := io.channels.write_data_channel.WDATA
        val wstrb = io.channels.write_data_channel.WSTRB
        when(wstrb === 0.U) {
          write_strobe := VecInit(Seq.fill(Parameters.WordSize)(true.B))
        }.otherwise {
          write_strobe := VecInit(wstrb.asBools)
        }
        WREADY := true.B
        write := true.B
        state := AXI4LiteStates.WriteData
      }
    }

    is(AXI4LiteStates.WriteData) {
      // 向设备发出一次写请求，WREADY 在一个周期后拉低
      WREADY := false.B
      write := false.B
      BVALID := true.B
      state := AXI4LiteStates.WriteResp
    }

    is(AXI4LiteStates.WriteResp) {
      // 等待主机对写响应握手
      when(io.channels.write_response_channel.BREADY) {
        BVALID := false.B
        state := AXI4LiteStates.Idle
      }
    }
  }

}

class AXI4LiteMaster(addrWidth: Int, dataWidth: Int) extends Module {
  val io = IO(new Bundle {
    val channels = new AXI4LiteChannels(addrWidth, dataWidth)
    val bundle = new AXI4LiteMasterBundle(addrWidth, dataWidth)
  })
  val state = RegInit(AXI4LiteStates.Idle)
  io.bundle.busy := state =/= AXI4LiteStates.Idle

  val addr = RegInit(0.U(dataWidth.W))
  val read_valid = RegInit(false.B)
  io.bundle.read_valid := read_valid
  val write_valid = RegInit(false.B)
  io.bundle.write_valid := write_valid
  val write_data = RegInit(0.U(dataWidth.W))
  val write_strobe = RegInit(VecInit(Seq.fill(Parameters.WordSize)(false.B)))
  val read_data = RegInit(0.U(dataWidth.W))

  io.channels.read_address_channel.ARADDR := 0.U
  val ARVALID = RegInit(false.B)
  io.channels.read_address_channel.ARVALID := ARVALID
  io.channels.read_address_channel.ARPROT := 0.U
  val RREADY = RegInit(false.B)
  io.channels.read_data_channel.RREADY := RREADY

  io.bundle.read_data := read_data
  val AWVALID = RegInit(false.B)
  io.channels.write_address_channel.AWADDR := 0.U
  io.channels.write_address_channel.AWVALID := AWVALID
  val WVALID = RegInit(false.B)
  io.channels.write_data_channel.WVALID := WVALID
  io.channels.write_data_channel.WDATA := write_data
  io.channels.write_address_channel.AWPROT := 0.U
  io.channels.write_data_channel.WSTRB := write_strobe.asUInt
  val BREADY = RegInit(false.B)
  io.channels.write_response_channel.BREADY := BREADY
  //lab4(BUS)
  // read_valid / write_valid 只拉高一个周期，在switch之前清零以便switch可以覆盖
  when(read_valid) { read_valid := false.B }
  when(write_valid) { write_valid := false.B }

  switch(state) {
    is(AXI4LiteStates.Idle) {
      ARVALID := false.B
      RREADY := false.B
      AWVALID := false.B
      WVALID := false.B
      BREADY := false.B

      when(io.bundle.read) {
        // 发起读事务
        addr := io.bundle.address
        state := AXI4LiteStates.ReadAddr
        ARVALID := true.B
      }.elsewhen(io.bundle.write) {
        // 发起写事务
        addr := io.bundle.address
        write_data := io.bundle.write_data
        write_strobe := io.bundle.write_strobe
        state := AXI4LiteStates.WriteAddr
        AWVALID := true.B
      }
    }

    is(AXI4LiteStates.ReadAddr) {
      // 发送读地址，直到从机准备好
      ARVALID := true.B
      io.channels.read_address_channel.ARADDR := addr
      when(io.channels.read_address_channel.ARREADY) {
        ARVALID := false.B
        RREADY := true.B
        state := AXI4LiteStates.ReadData
      }
    }

    is(AXI4LiteStates.ReadData) {
      // 保持 RREADY 高，等待从机 RVALID
      RREADY := true.B
      when(io.channels.read_data_channel.RVALID) {
        read_data := io.channels.read_data_channel.RDATA
        RREADY := false.B
        read_valid := true.B
        state := AXI4LiteStates.Idle
      }
    }

    is(AXI4LiteStates.WriteAddr) {
      // 发送写地址，直到从机准备好
      AWVALID := true.B
      io.channels.write_address_channel.AWADDR := addr
      when(io.channels.write_address_channel.AWREADY) {
        AWVALID := false.B
        WVALID := true.B
        state := AXI4LiteStates.WriteData
      }
    }

    is(AXI4LiteStates.WriteData) {
      // 发送写数据，直到从机接受
      WVALID := true.B
      io.channels.write_data_channel.WDATA := write_data
      io.channels.write_data_channel.WSTRB := write_strobe.asUInt
      when(io.channels.write_data_channel.WREADY) {
        WVALID := false.B
        BREADY := true.B
        state := AXI4LiteStates.WriteResp
      }
    }

    is(AXI4LiteStates.WriteResp) {
      // 等待从机写响应
      BREADY := true.B
      when(io.channels.write_response_channel.BVALID) {
        BREADY := false.B
        write_valid := true.B
        state := AXI4LiteStates.Idle
      }
    }
  }

}