// Copyright 2024 Google LLC
//
// 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 chai

import chisel3._
import chisel3.util._
import _root_.circt.stage.ChiselStage

case class Parameters() {
  val sramReadPorts = 1
  val sramWritePorts = 1
  val sramReadWritePorts = 0
  val sramDataBits = 256
  val sramBytes = 4 * 1024 * 1024
  def sramDataEntries(): Int = {
    ((sramBytes * 8) / sramDataBits)
  }
  def sramAddrBits(): Int = {
    log2Ceil(sramDataEntries())
  }
}

object ChAI {
  def apply(p: Parameters): ChAI = {
    return Module(new ChAI(p))
  }
}

class ChAI(p: Parameters) extends RawModule {
  val io = IO(new Bundle {
    val clk_i = Input(Clock())
    val rst_ni = Input(AsyncReset())
    val sram = new Bundle {
      val write_address = Input(UInt(p.sramAddrBits().W))
      val write_enable = Input(Bool())
      val write_data = Input(UInt(p.sramDataBits.W))
    }
    val finish = Output(Bool())
    val fault = Output(Bool())
    val freeze = Input(Bool())

    val uart_rx = Input(Bool())
    val uart_tx = Output(Bool())
  })

  // TODO(atv): Compute that we don't have any overlaps in regions.
  val memoryRegions = Seq(
    new kelvin.MemoryRegion(0, 4 * 1024 * 1024, true, 256), // SRAM
    new kelvin.MemoryRegion(4 * 1024 * 1024, 4 * 1024 * 1024, false, 256) // UART
  )
  val kelvin_p = kelvin.Parameters(memoryRegions)
  val rst_i = (!io.rst_ni.asBool).asAsyncReset

  val u_kelvin = matcha.Kelvin(kelvin_p)
  u_kelvin.clk_i := io.clk_i
  u_kelvin.rst_ni := io.rst_ni
  u_kelvin.clk_freeze := io.freeze
  u_kelvin.ml_reset := 0.U
  u_kelvin.pc_start := 0.U
  u_kelvin.volt_sel := 0.U
  u_kelvin.debug_req := 0.U

  io.finish := u_kelvin.finish
  io.fault := u_kelvin.fault

  withClockAndReset(io.clk_i, rst_i) {
    val tlul_p = new kelvin.TLULParameters()
    val kelvin_to_tlul = chai.KelvinToTlul(tlul_p, kelvin_p)
    kelvin_to_tlul.io.kelvin <> u_kelvin.mem

    val tlul_sram =
      SRAM(p.sramDataEntries(), UInt(p.sramDataBits.W), p.sramReadPorts, p.sramWritePorts, p.sramReadWritePorts)
    val tlul_adapter_sram = Module(new chai.TlulAdapterSram())
    tlul_adapter_sram.io.clk_i := io.clk_i
    tlul_adapter_sram.io.rst_ni := io.rst_ni
    tlul_adapter_sram.io.en_ifetch_i := 9.U // MuBi4False
    tlul_sram.readPorts(0).enable := tlul_adapter_sram.io.req_o
    tlul_sram.readPorts(0).address := tlul_adapter_sram.io.addr_o
    tlul_sram.writePorts(0).enable := Mux(io.freeze, io.sram.write_enable, tlul_adapter_sram.io.we_o)
    tlul_sram.writePorts(0).address := Mux(io.freeze, io.sram.write_address, tlul_adapter_sram.io.addr_o)
    tlul_sram.writePorts(0).data := Mux(io.freeze, io.sram.write_data, tlul_adapter_sram.io.wdata_o)
    tlul_adapter_sram.io.gnt_i := 1.U
    tlul_adapter_sram.io.rdata_i := tlul_sram.readPorts(0).data
    tlul_adapter_sram.io.rvalid_i := 1.U
    tlul_adapter_sram.io.rerror_i := 0.U

    val uart = Module(new chai.Uart(tlul_p))
    uart.io.clk_i := io.clk_i
    uart.io.rst_ni := io.rst_ni
    uart.io.alert_rx_i := 0.U
    uart.io.cio_rx_i := io.uart_rx
    io.uart_tx := uart.io.cio_tx_o

    val crossbar =
      Module(new kelvin.TileLinkUL(tlul_p, kelvin_p.m, /* hosts= */ 1))
    crossbar.io.hosts_a(0) <> kelvin_to_tlul.io.tl_o
    crossbar.io.hosts_d(0) <> kelvin_to_tlul.io.tl_i
    crossbar.io.devices_a(0) <> tlul_adapter_sram.io.tl_i
    crossbar.io.devices_d(0) <> tlul_adapter_sram.io.tl_o
    crossbar.io.devices_a(1) <> uart.io.tl_i
    crossbar.io.devices_d(1) <> uart.io.tl_o
  }
}

object EmitChAI extends App {
  val p = new Parameters()
  ChiselStage.emitSystemVerilogFile(new ChAI(p), args)
}