// Copyright 2026 International Digital Economy Academy
//
// 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.

///|
priv struct WgslIrFunctionCallGraph {
  direct_calls : Array[Array[Int]]
}

///|
fn WgslIrWriterPlanner::function_order(
  self : WgslIrWriterPlanner,
) -> Array[Int] {
  let call_graph = self.function_call_graph()
  let order = self.generated_import_function_order()
  let reachable = self.reachable_function_order(call_graph)
  for index in reachable {
    if self.function_order_should_include(index, order) {
      order.push(index)
    }
  }
  for index in 0.. WgslIrFunctionCallGraph {
  let direct_calls : Array[Array[Int]] = []
  for function in self.shader_module.functions.items {
    direct_calls.push(wgsl_ir_collect_block_function_calls(function))
  }
  { direct_calls, }
}

///|
fn WgslIrFunctionCallGraph::direct_calls_for(
  self : WgslIrFunctionCallGraph,
  index : Int,
) -> Array[Int] {
  match self.direct_calls.get(index) {
    Some(calls) => calls
    None => []
  }
}

///|
fn WgslIrWriterPlanner::generated_import_function_order(
  self : WgslIrWriterPlanner,
) -> Array[Int] {
  let order : Array[Int] = []
  for event in self.compatibility.import_arena_events {
    for symbol in event.generated_symbols() {
      match symbol.kind() {
        Function =>
          self.push_generated_import_function_by_name(
            order,
            symbol.generated_symbol_name(),
          )
        Type | Constant | Override | GlobalVariable => ()
      }
    }
  }
  for index in 0.. Unit {
  for index in 0.. Array[Int] {
  let order : Array[Int] = []
  for entry_point in self.shader_module.entry_points {
    self.push_source_prefix_before_entry(order, entry_point, call_graph)
    let direct = wgsl_ir_collect_block_function_calls(entry_point.function)
    for function_index in direct {
      self.push_reachable_function_with_dependencies(
        order, function_index, call_graph,
      )
    }
  }
  order
}

///|
fn WgslIrWriterPlanner::push_source_prefix_before_entry(
  self : WgslIrWriterPlanner,
  order : Array[Int],
  entry_point : EntryPoint,
  call_graph : WgslIrFunctionCallGraph,
) -> Unit {
  guard entry_point.source_start >= 0 else { return }
  for index in 0..= 0 &&
      function.source_start < entry_point.source_start &&
      !self.compatibility.is_generated_import(function.name) &&
      self.function_order_should_include(index, order) {
      self.push_reachable_function_with_dependencies(order, index, call_graph)
    }
  }
}

///|
fn WgslIrWriterPlanner::push_function_with_dependencies(
  self : WgslIrWriterPlanner,
  order : Array[Int],
  index : Int,
  call_graph : WgslIrFunctionCallGraph,
) -> Unit {
  let visiting : Array[Int] = []
  self.push_function_with_dependencies_inner(order, visiting, index, call_graph)
}

///|
fn WgslIrWriterPlanner::push_reachable_function_with_dependencies(
  self : WgslIrWriterPlanner,
  order : Array[Int],
  index : Int,
  call_graph : WgslIrFunctionCallGraph,
) -> Unit {
  if !self.function_order_should_include(index, order) {
    return
  }
  let dependencies : Array[Int] = []
  self.collect_transitive_function_dependencies(index, dependencies, call_graph)
  for dependency_index in 0.. index {
      self.push_reachable_function_with_dependencies(
        order, dependency_index, call_graph,
      )
    }
  }
  if self.function_order_should_include(index, order) {
    order.push(index)
  }
}

///|
fn WgslIrWriterPlanner::collect_transitive_function_dependencies(
  self : WgslIrWriterPlanner,
  index : Int,
  dependencies : Array[Int],
  call_graph : WgslIrFunctionCallGraph,
) -> Unit {
  if !self.function_order_should_include(index, []) {
    return
  }
  let direct = call_graph.direct_calls_for(index)
  for dependency in direct {
    if dependency == index {
      continue
    }
    if !dependencies.contains(dependency) {
      dependencies.push(dependency)
      self.collect_transitive_function_dependencies(
        dependency, dependencies, call_graph,
      )
    }
  }
}

///|
fn WgslIrWriterPlanner::push_function_with_dependencies_inner(
  self : WgslIrWriterPlanner,
  order : Array[Int],
  visiting : Array[Int],
  index : Int,
  call_graph : WgslIrFunctionCallGraph,
) -> Unit {
  if !self.function_order_should_include(index, order) ||
    visiting.contains(index) {
    return
  }
  visiting.push(index)
  let direct = call_graph.direct_calls_for(index)
  for dependency_index in direct {
    if dependency_index != index {
      self.push_function_with_dependencies_inner(
        order, visiting, dependency_index, call_graph,
      )
    }
  }
  visiting.pop() |> ignore
  if self.function_order_should_include(index, order) {
    order.push(index)
  }
}

///|
fn WgslIrWriterPlanner::function_order_should_include(
  self : WgslIrWriterPlanner,
  index : Int,
  emitted : Array[Int],
) -> Bool {
  if emitted.contains(index) {
    return false
  }
  if !self.should_emit_function(index) {
    return false
  }
  !wgsl_ir_function_is_entry_point_name(
    self.shader_module.functions.items[index],
    self.shader_module,
  )
}

///|
fn wgsl_ir_function_is_entry_point_name(
  function : Function,
  shader_module : Module,
) -> Bool {
  match function.name {
    Some(name) =>
      for entry_point in shader_module.entry_points {
        if entry_point.name == name {
          return true
        }
      }
    None => ()
  }
  false
}

///|
fn WgslIrWriterPlanner::should_emit_function(
  self : WgslIrWriterPlanner,
  index : Int,
) -> Bool {
  match self.filter {
    Some(filter) => filter.contains_function(index)
    None => true
  }
}

///|
fn wgsl_ir_collect_block_function_calls(function : Function) -> Array[Int] {
  let calls : Array[Int] = []
  wgsl_ir_collect_function_calls_from_block(function, function.body, calls)
  calls
}

///|
fn wgsl_ir_collect_function_calls_from_block(
  function : Function,
  block : Block,
  calls : Array[Int],
) -> Unit {
  for statement in block.statements {
    wgsl_ir_collect_function_calls_from_statement(function, statement, calls)
  }
}

///|
fn wgsl_ir_collect_function_calls_from_statement(
  function : Function,
  statement : Statement,
  calls : Array[Int],
) -> Unit {
  match statement {
    Emit(range) => {
      let mut index = range.start.index()
      while index <= range.end.index() {
        wgsl_ir_collect_function_calls_from_expression(
          function,
          Handle(index),
          calls,
        )
        index = index + 1
      }
    }
    Phony(handle) =>
      wgsl_ir_collect_function_calls_from_expression(function, handle, calls)
    Block(nested) =>
      wgsl_ir_collect_function_calls_from_block(function, nested, calls)
    If(condition, accept, reject) => {
      wgsl_ir_collect_function_calls_from_expression(function, condition, calls)
      wgsl_ir_collect_function_calls_from_block(function, accept, calls)
      wgsl_ir_collect_function_calls_from_block(function, reject, calls)
    }
    Switch(selector, cases) => {
      wgsl_ir_collect_function_calls_from_expression(function, selector, calls)
      for case in cases {
        wgsl_ir_collect_function_calls_from_block(function, case.body, calls)
      }
    }
    Loop(body, continuing, break_if) => {
      wgsl_ir_collect_function_calls_from_block(function, body, calls)
      wgsl_ir_collect_function_calls_from_block(function, continuing, calls)
      match break_if {
        Some(condition) =>
          wgsl_ir_collect_function_calls_from_expression(
            function, condition, calls,
          )
        None => ()
      }
    }
    Return(value) =>
      match value {
        Some(expr) =>
          wgsl_ir_collect_function_calls_from_expression(function, expr, calls)
        None => ()
      }
    ImplicitReturn => ()
    ConstAssert(expr) =>
      wgsl_ir_collect_function_calls_from_expression(function, expr, calls)
    Store(pointer, value) => {
      wgsl_ir_collect_function_calls_from_expression(function, pointer, calls)
      wgsl_ir_collect_function_calls_from_expression(function, value, calls)
    }
    ImageStore(image, coordinate, array_index, value) => {
      wgsl_ir_collect_function_calls_from_expression(function, image, calls)
      wgsl_ir_collect_function_calls_from_expression(
        function, coordinate, calls,
      )
      match array_index {
        Some(expr) =>
          wgsl_ir_collect_function_calls_from_expression(function, expr, calls)
        None => ()
      }
      wgsl_ir_collect_function_calls_from_expression(function, value, calls)
    }
    Atomic(pointer, _, value, compare, result) => {
      wgsl_ir_collect_function_calls_from_expression(function, pointer, calls)
      wgsl_ir_collect_function_calls_from_expression(function, value, calls)
      match compare {
        Some(expr) =>
          wgsl_ir_collect_function_calls_from_expression(function, expr, calls)
        None => ()
      }
      match result {
        Some(expr) =>
          wgsl_ir_collect_function_calls_from_expression(function, expr, calls)
        None => ()
      }
    }
    ImageAtomic(image, coordinate, array_index, _, value) => {
      wgsl_ir_collect_function_calls_from_expression(function, image, calls)
      wgsl_ir_collect_function_calls_from_expression(
        function, coordinate, calls,
      )
      match array_index {
        Some(expr) =>
          wgsl_ir_collect_function_calls_from_expression(function, expr, calls)
        None => ()
      }
      wgsl_ir_collect_function_calls_from_expression(function, value, calls)
    }
    WorkGroupUniformLoad(pointer, result) => {
      wgsl_ir_collect_function_calls_from_expression(function, pointer, calls)
      wgsl_ir_collect_function_calls_from_expression(function, result, calls)
    }
    Call(callee, arguments, result) => {
      for argument in arguments {
        wgsl_ir_collect_function_calls_from_expression(
          function, argument, calls,
        )
      }
      match result {
        Some(expr) =>
          wgsl_ir_collect_function_calls_from_expression(function, expr, calls)
        None => ()
      }
      wgsl_ir_push_function_call(calls, callee.index())
    }
    RayQuery(query, _) =>
      wgsl_ir_collect_function_calls_from_expression(function, query, calls)
    RayPipelineFunction(TraceRay(a, b, c)) => {
      wgsl_ir_collect_function_calls_from_expression(function, a, calls)
      wgsl_ir_collect_function_calls_from_expression(function, b, calls)
      wgsl_ir_collect_function_calls_from_expression(function, c, calls)
    }
    SubgroupBallot(value, result) => {
      wgsl_ir_collect_function_calls_from_expression(function, value, calls)
      match result {
        Some(expr) =>
          wgsl_ir_collect_function_calls_from_expression(function, expr, calls)
        None => ()
      }
    }
    SubgroupGather(_, value, result)
    | SubgroupCollectiveOperation(_, _, value, result) => {
      wgsl_ir_collect_function_calls_from_expression(function, value, calls)
      wgsl_ir_collect_function_calls_from_expression(function, result, calls)
    }
    CooperativeStore(pointer, data) => {
      wgsl_ir_collect_function_calls_from_expression(function, pointer, calls)
      wgsl_ir_collect_function_calls_from_expression(
        function,
        data.pointer,
        calls,
      )
      wgsl_ir_collect_function_calls_from_expression(
        function,
        data.stride,
        calls,
      )
    }
    Declare(_)
    | Break
    | Continue
    | Kill
    | ControlBarrier(_)
    | MemoryBarrier(_) => ()
  }
}

///|
fn wgsl_ir_collect_function_calls_from_expression(
  function : Function,
  handle : Handle,
  calls : Array[Int],
) -> Unit {
  match function.expressions.items.get(handle.index()) {
    Some(expression) =>
      wgsl_ir_collect_function_calls_from_expression_inner(
        function, expression, calls,
      )
    None => ()
  }
}

///|
fn wgsl_ir_collect_function_calls_from_expression_inner(
  function : Function,
  expression : Expression,
  calls : Array[Int],
) -> Unit {
  match expression {
    Compose(_, components) =>
      for component in components {
        wgsl_ir_collect_function_calls_from_expression(
          function, component, calls,
        )
      }
    Access(base, index) | Binary(_, base, index) => {
      wgsl_ir_collect_function_calls_from_expression(function, base, calls)
      wgsl_ir_collect_function_calls_from_expression(function, index, calls)
    }
    AccessIndex(base, _)
    | Component(base, _)
    | Load(base)
    | AddressOf(_, base)
    | Unary(_, base)
    | Bitcast(_, base)
    | As(base, _, _)
    | Relational(_, base)
    | ArrayLength(base)
    | WorkGroupUniformLoad(base)
    | Splat(_, base)
    | Swizzle(_, base, _)
    | Derivative(_, _, base)
    | WorkGroupUniformLoadResult(base)
    | RayQueryProceed(base)
    | RayQueryConfirmIntersection(base)
    | RayQueryTerminate(base)
    | RayQueryGetIntersection(base, _) =>
      wgsl_ir_collect_function_calls_from_expression(function, base, calls)
    AtomicResult(_, _) => ()
    AtomicCall(_, arguments) | SubgroupCall(_, arguments) =>
      for argument in arguments {
        wgsl_ir_collect_function_calls_from_expression(
          function, argument, calls,
        )
      }
    FunctionCall(target, arguments) => {
      for argument in arguments {
        wgsl_ir_collect_function_calls_from_expression(
          function, argument, calls,
        )
      }
      wgsl_ir_push_function_call(calls, target.index())
    }
    Select(condition, accept, reject) => {
      wgsl_ir_collect_function_calls_from_expression(function, condition, calls)
      wgsl_ir_collect_function_calls_from_expression(function, accept, calls)
      wgsl_ir_collect_function_calls_from_expression(function, reject, calls)
    }
    ImageSample(
      image,
      sampler,
      _,
      coordinate,
      array_index,
      offset,
      level,
      depth_ref,
      _
    ) => {
      wgsl_ir_collect_function_calls_from_expression(function, image, calls)
      wgsl_ir_collect_function_calls_from_expression(function, sampler, calls)
      wgsl_ir_collect_function_calls_from_expression(
        function, coordinate, calls,
      )
      wgsl_ir_collect_optional_function_call(function, array_index, calls)
      wgsl_ir_collect_optional_function_call(function, offset, calls)
      wgsl_ir_collect_sample_level_function_calls(function, level, calls)
      wgsl_ir_collect_optional_function_call(function, depth_ref, calls)
    }
    ImageLoad(image, coordinate, array_index, sample, level) => {
      wgsl_ir_collect_function_calls_from_expression(function, image, calls)
      wgsl_ir_collect_function_calls_from_expression(
        function, coordinate, calls,
      )
      wgsl_ir_collect_optional_function_call(function, array_index, calls)
      wgsl_ir_collect_optional_function_call(function, sample, calls)
      wgsl_ir_collect_optional_function_call(function, level, calls)
    }
    ImageQuery(image, query) => {
      wgsl_ir_collect_function_calls_from_expression(function, image, calls)
      match query {
        Size(level) =>
          wgsl_ir_collect_optional_function_call(function, level, calls)
        _ => ()
      }
    }
    Math(_, a, b, c, d) => {
      wgsl_ir_collect_function_calls_from_expression(function, a, calls)
      wgsl_ir_collect_optional_function_call(function, b, calls)
      wgsl_ir_collect_optional_function_call(function, c, calls)
      wgsl_ir_collect_optional_function_call(function, d, calls)
    }
    CallResult(target) => wgsl_ir_push_function_call(calls, target.index())
    RayQueryInitialize(query, acceleration, ray_desc) => {
      wgsl_ir_collect_function_calls_from_expression(function, query, calls)
      wgsl_ir_collect_function_calls_from_expression(
        function, acceleration, calls,
      )
      wgsl_ir_collect_function_calls_from_expression(function, ray_desc, calls)
    }
    RayQueryGenerateIntersection(query, intersection_t) => {
      wgsl_ir_collect_function_calls_from_expression(function, query, calls)
      wgsl_ir_collect_function_calls_from_expression(
        function, intersection_t, calls,
      )
    }
    RayQueryVertexPositions(query, _) =>
      wgsl_ir_collect_function_calls_from_expression(function, query, calls)
    SubgroupOperationResult(handle) =>
      wgsl_ir_collect_function_calls_from_expression(function, handle, calls)
    CooperativeLoad(_, _, _, data) => {
      wgsl_ir_collect_function_calls_from_expression(
        function,
        data.pointer,
        calls,
      )
      wgsl_ir_collect_function_calls_from_expression(
        function,
        data.stride,
        calls,
      )
    }
    CooperativeMultiplyAdd(a, b, c) => {
      wgsl_ir_collect_function_calls_from_expression(function, a, calls)
      wgsl_ir_collect_function_calls_from_expression(function, b, calls)
      wgsl_ir_collect_function_calls_from_expression(function, c, calls)
    }
    Literal(_)
    | Constant(_)
    | Override(_)
    | ZeroValue(_)
    | FunctionArgument(_)
    | GlobalVariable(_)
    | LocalVariable(_)
    | RayQueryProceedResult
    | SubgroupBallotResult => ()
  }
}

///|
fn wgsl_ir_collect_sample_level_function_calls(
  function : Function,
  level : SampleLevel,
  calls : Array[Int],
) -> Unit {
  match level {
    Exact(handle) | Bias(handle) =>
      wgsl_ir_collect_function_calls_from_expression(function, handle, calls)
    Gradient(x, y) => {
      wgsl_ir_collect_function_calls_from_expression(function, x, calls)
      wgsl_ir_collect_function_calls_from_expression(function, y, calls)
    }
    Auto | Zero => ()
  }
}

///|
fn wgsl_ir_collect_optional_function_call(
  function : Function,
  handle : Handle?,
  calls : Array[Int],
) -> Unit {
  match handle {
    Some(value) =>
      wgsl_ir_collect_function_calls_from_expression(function, value, calls)
    None => ()
  }
}

///|
fn wgsl_ir_push_function_call(calls : Array[Int], index : Int) -> Unit {
  if !calls.contains(index) {
    calls.push(index)
  }
}