// SPDX-License-Identifier: Apache-2.0

package chisel3.reflect

import chisel3._
import chisel3.internal._
import chisel3.internal.firrtl._
import chisel3.experimental.BaseModule
import scala.reflect.ClassTag

object DataMirror {
  def widthOf(target:              Data): Width = target.width
  def specifiedDirectionOf(target: Data): SpecifiedDirection = target.specifiedDirection
  def directionOf(target: Data): ActualDirection = {
    requireIsHardware(target, "node requested directionality on")
    target.direction
  }

  private def hasBinding[B <: ConstrainedBinding: ClassTag](target: Data) = {
    // Cannot use isDefined because of the ClassTag
    target.topBindingOpt match {
      case Some(b: B) => true
      case _ => false
    }
  }

  /** Check if a given `Data` is an IO port
    * @param x the `Data` to check
    * @return `true` if x is an IO port, `false` otherwise
    */
  def isIO(x: Data): Boolean = hasBinding[PortBinding](x)

  /** Check if a given `Data` is a Wire
    * @param x the `Data` to check
    * @return `true` if x is a Wire, `false` otherwise
    */
  def isWire(x: Data): Boolean = hasBinding[WireBinding](x)

  /** Check if a given `Data` is a Reg
    * @param x the `Data` to check
    * @return `true` if x is a Reg, `false` otherwise
    */
  def isReg(x: Data): Boolean = hasBinding[RegBinding](x)

  /** Check if two Chisel types are the same type.
    * Internally, this is dispatched to each Chisel type's
    * `typeEquivalent` function for each type to determine
    * if the types are intended to be equal.
    *
    * For most types, different parameters should ensure
    * that the types are different.
    * For example, `UInt(8.W)` and `UInt(16.W)` are different.
    * Likewise, Records check that both Records have the same
    * elements with the same types.
    *
    * Equivalent to being structural, alignment, and width type equivalent
    *
    * @param x First Chisel type
    * @param y Second Chisel type
    * @return true if the two Chisel types are equal.
    */
  def checkTypeEquivalence(x: Data, y: Data): Boolean = x.typeEquivalent(y)

  /** Check if two Chisel types have the same alignments for all matching members
    *
    * This means that for matching members in Aggregates, they must have matching member alignments relative to the parent type
    * For matching non-aggregates, they must be the same alignment to their parent type.
    *
    * @param x First Chisel type
    * @param y Second Chisel type
    * @return true if the two Chisel types have alignment type equivalence.
    */
  def checkAlignmentTypeEquivalence(x: Data, y: Data): Boolean = {
    //TODO(azidar): Perhaps there is a better pattern of `iterateOverMatches` that we can support
    collectMembersOverMatches(connectable.Alignment(x, true), connectable.Alignment(y, true)) {
      case (a, b) => a.alignment == b.alignment
    }(AlignmentMatchingZipOfChildren).forall(r => r)
  }

  /** Returns the ports of a module
    * {{{
    * class MyModule extends Module {
    *   val io = IO(new Bundle {
    *     val in = Input(UInt(8.W))
    *     val out = Output(Vec(2, UInt(8.W)))
    *   })
    *   val extra = IO(Input(UInt(8.W)))
    *   val delay = RegNext(io.in)
    *   io.out(0) := delay
    *   io.out(1) := delay + extra
    * }
    * val mod = Module(new MyModule)
    * DataMirror.modulePorts(mod)
    * // returns: Seq(
    * //   "clock" -> mod.clock,
    * //   "reset" -> mod.reset,
    * //   "io" -> mod.io,
    * //   "extra" -> mod.extra
    * // )
    * }}}
    */
  def modulePorts(target: BaseModule): Seq[(String, Data)] = target.getChiselPorts

  /** Returns a recursive representation of a module's ports with underscore-qualified names
    * {{{
    * class MyModule extends Module {
    *   val io = IO(new Bundle {
    *     val in = Input(UInt(8.W))
    *     val out = Output(Vec(2, UInt(8.W)))
    *   })
    *   val extra = IO(Input(UInt(8.W)))
    *   val delay = RegNext(io.in)
    *   io.out(0) := delay
    *   io.out(1) := delay + extra
    * }
    * val mod = Module(new MyModule)
    * DataMirror.fullModulePorts(mod)
    * // returns: Seq(
    * //   "clock" -> mod.clock,
    * //   "reset" -> mod.reset,
    * //   "io" -> mod.io,
    * //   "io_out" -> mod.io.out,
    * //   "io_out_0" -> mod.io.out(0),
    * //   "io_out_1" -> mod.io.out(1),
    * //   "io_in" -> mod.io.in,
    * //   "extra" -> mod.extra
    * // )
    * }}}
    * @note The returned ports are redundant. An [[Aggregate]] port will be present along with all
    *       of its children.
    * @see [[DataMirror.modulePorts]] for a non-recursive representation of the ports.
    */
  def fullModulePorts(target: BaseModule): Seq[(String, Data)] = {
    def getPortNames(name: String, data: Data): Seq[(String, Data)] = Seq(name -> data) ++ (data match {
      case _: Element => Seq()
      case r: Record =>
        r.elements.toSeq.flatMap {
          case (eltName, elt) =>
            if (r._isOpaqueType) { getPortNames(s"${name}", elt) }
            else { getPortNames(s"${name}_${eltName}", elt) }
        }
      case v: Vec[_] => v.zipWithIndex.flatMap { case (elt, index) => getPortNames(s"${name}_${index}", elt) }
    })
    modulePorts(target).flatMap {
      case (name, data) =>
        getPortNames(name, data).toList
    }
  }

  /** Returns the parent module within which a module instance is instantiated
    *
    * @note Top-level modules in any given elaboration do not have a parent
    * @param target a module instance
    * @return the parent of the `target`, if one exists
    */
  def getParent(target: BaseModule): Option[BaseModule] = target._parent

  // Internal reflection-style APIs, subject to change and removal whenever.
  object internal {
    def isSynthesizable(target: Data): Boolean = target.isSynthesizable
    // For those odd cases where you need to care about object reference and uniqueness
    def chiselTypeClone[T <: Data](target: Data): T = {
      target.cloneTypeFull.asInstanceOf[T]
    }
  }

  /** Return all expanded components, including intermediate aggregate nodes
    *
    * @param d Component to find leafs if aggregate typed. Intermediate fields/indicies ARE included
    */
  def getIntermediateAndLeafs(d: Data): Seq[Data] = collectMembers(d) { case x => x }.toSeq

  /** Collects all fields selected by collector within a data and all recursive children fields
    * Accepts a collector partial function, rather than a collector function
    *
    * @param data Data to collect fields, as well as all children datas it directly and indirectly instantiates
    * @param collector Collector partial function to pick which components to collect
    * @tparam T Type of the component that will be collected
    */
  def collectMembers[T](d: Data)(collector: PartialFunction[Data, T]): Iterable[T] = {
    val myItems = collector.lift(d)
    val deepChildrenItems = d match {
      case a: Aggregate => a.getElements.flatMap { x => collectMembers(x)(collector) }
      case other => Nil
    }
    myItems ++ deepChildrenItems
  }

  // Alignment-aware collections
  import connectable.{AlignedWithRoot, Alignment, FlippedWithRoot}
  // Implement typeclass to enable collecting over Alignment
  implicit val AlignmentMatchingZipOfChildren = new HasMatchingZipOfChildren[Alignment] {
    def matchingZipOfChildren(
      left:  Option[Alignment],
      right: Option[Alignment]
    ): Seq[(Option[Alignment], Option[Alignment])] =
      Alignment.matchingZipOfChildren(left, right)
  }

  /** Collects all members of base who are aligned w.r.t. base
    * Accepts a collector partial function, rather than a collector function
    *
    * @param base Data from whom aligned members (w.r.t. base) are collected
    * @param collector Collector partial function to pick which components to collect
    * @tparam T Type of the component that will be collected
    */
  def collectAlignedDeep[T](base: Data)(pf: PartialFunction[Data, T]): Seq[T] = {
    collectMembersOverAllForAny(Some(Alignment(base, true)), None) {
      case (Some(x: AlignedWithRoot), _) => (pf.lift(x.member), None)
    }.map(_._1).flatten
  }

  /** Collects all members of base who are flipped w.r.t. base
    * Accepts a collector partial function, rather than a collector function
    *
    * @param base Data from whom flipped members (w.r.t. base) are collected
    * @param collector Collector partial function to pick which components to collect
    * @tparam T Type of the component that will be collected
    */
  def collectFlippedDeep[T](base: Data)(pf: PartialFunction[Data, T]): Seq[T] = {
    collectMembersOverAllForAny(Some(Alignment(base, true)), None) {
      case (Some(x: FlippedWithRoot), _) => (pf.lift(x.member), None)
    }.map(_._1).flatten
  }

  /** Typeclass trait to use collectMembersOverMatches, collectMembersOverAll, collectMembersOverAllForAny, collectMembersOverAllForAnyFunction */
  trait HasMatchingZipOfChildren[T] {
    def matchingZipOfChildren(left: Option[T], right: Option[T]): Seq[(Option[T], Option[T])]
  }

  /** Collects over members left and right who have structurally corresponding members in both left and right
    * Accepts a collector partial function, rather than a collector function
    *
    * @param left Data from whom members are collected
    * @param right Data from whom members are collected
    * @param collector Collector partial function to pick which components from left and right to collect
    * @tparam T Type of the thing being collected
    */
  def collectMembersOverMatches[D: HasMatchingZipOfChildren, T](
    left:      D,
    right:     D
  )(collector: PartialFunction[(D, D), T]
  ): Seq[T] = {
    def newCollector(lOpt: Option[D], rOpt: Option[D]): Option[(Option[T], Option[Unit])] = {
      (lOpt, rOpt) match {
        case (Some(l), Some(r)) =>
          collector.lift((l, r)) match {
            case Some(x: T) => Some((Some(x), None))
            case None => None
          }
        case other => None
      }
    }
    collectMembersOverAllForAnyFunction(Some(left), Some(right)) {
      case (Some(l), Some(r)) =>
        collector.lift((l, r)) match {
          case Some(x: T) => Some((Some(x), None))
          case None => None
        }
      case other => None
    }.collect {
      case (Some(x: T), None) => (x)
    }
  }

  /** Collects over members left and right who have structurally corresponding members in either left and right
    * Accepts a collector partial function, rather than a collector function
    *
    * @param left Data from whom members are collected
    * @param right Data from whom members are collected
    * @param collector Collector partial function to pick which components from left, right, or both to collect
    * @tparam T Type of the thing being collected
    */
  def collectMembersOverAll[D: HasMatchingZipOfChildren, T](
    left:      D,
    right:     D
  )(collector: PartialFunction[(Option[D], Option[D]), T]
  ): Seq[T] = {
    collectMembersOverAllForAnyFunction(Some(left), Some(right)) {
      case (lOpt: Option[D], rOpt: Option[D]) =>
        collector.lift((lOpt, rOpt)) match {
          case Some(x: T) => Some((Some(x), None))
          case None => None
        }
    }.collect {
      case (Some(x: T), None) => x
    }
  }

  /** Collects over members left and right who have structurally corresponding members in either left and right
    * Can return an optional value for left, right, both or neither
    * Accepts a collector partial function, rather than a collector function
    *
    * @param left Data from whom members are collected
    * @param right Data from whom members are collected
    * @param collector Collector partial function to pick which components from left, right, or both to collect
    * @tparam L Type of the thing being collected from the left
    * @tparam R Type of the thing being collected from the right
    */
  def collectMembersOverAllForAny[D: HasMatchingZipOfChildren, L, R](
    left:       Option[D],
    right:      Option[D]
  )(pcollector: PartialFunction[(Option[D], Option[D]), (Option[L], Option[R])]
  ): Seq[(Option[L], Option[R])] = {
    collectMembersOverAllForAnyFunction(left, right)(pcollector.lift)
  }

  /** Collects over members left and right who have structurally corresponding members in either left and right
    * Can return an optional value for left, right, both or neither
    * Accepts a full function
    *
    * @param left Data from whom members are collected
    * @param right Data from whom members are collected
    * @param collector Collector full function to pick which components from left, right, or both to collect
    * @tparam L Type of the thing being collected from the left
    * @tparam R Type of the thing being collected from the right
    */
  def collectMembersOverAllForAnyFunction[D: HasMatchingZipOfChildren, L, R](
    left:      Option[D],
    right:     Option[D]
  )(collector: ((Option[D], Option[D])) => Option[(Option[L], Option[R])]
  ): Seq[(Option[L], Option[R])] = {
    val myItems = collector((left, right)) match {
      case None               => Nil
      case Some((None, None)) => Nil
      case Some(other)        => Seq(other)
    }
    val matcher = implicitly[HasMatchingZipOfChildren[D]]
    val childItems = matcher.matchingZipOfChildren(left, right).flatMap {
      case (l, r) => collectMembersOverAllForAnyFunction(l, r)(collector)
    }
    myItems ++ childItems
  }
}
