// Copyright 2025 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.

///|
/// Import Resolution for diago IR
///
/// Handles `@path/to/file.d2` import statements.
/// Imports are resolved by a user-provided resolver function.

///|
/// Import resolver function type
/// Takes an import path and returns the file content, or None if not found
pub struct ImportResolver((String) -> String?)

///|
/// Create a new import resolver from a function
pub fn ImportResolver::new(f : (String) -> String?) -> ImportResolver {
  ImportResolver(f)
}

///|
/// Compile an AST Map with import resolution
pub fn compile_with_imports(
  ast : @ast.Map,
  resolver : ImportResolver,
) -> Map raise CompileError {
  let ir = Map::new()
  compile_map_nodes_with_imports(ir, ast.nodes, resolver, [])
  // Resolve variable substitutions
  ir.resolve_substitutions()
  ir
}

///|
fn compile_map_nodes_with_imports(
  ir : Map,
  nodes : Array[@ast.MapNode],
  resolver : ImportResolver,
  import_stack : Array[String],
) -> Unit raise CompileError {
  for node in nodes {
    match node {
      @ast.MapNode::Comment(_, _) => ()
      @ast.MapNode::BlockComment(_, _) => ()
      @ast.MapNode::Substitution(sub) =>
        ir.add_field(Field::new("", Some(substitution_scalar(sub)), None, []))
      @ast.MapNode::Import(imp) =>
        compile_import(ir, imp, resolver, import_stack)
      @ast.MapNode::Key(key) =>
        compile_key_with_imports(ir, key, resolver, import_stack)
    }
  }
}

///|
/// Compile an import statement
fn compile_import(
  ir : Map,
  imp : @ast.Import,
  resolver : ImportResolver,
  import_stack : Array[String],
) -> Unit raise CompileError {
  let (file_key, selector, display) = split_import_file_and_selector(imp)
  let (content, canonical) = resolve_import_content(resolver, file_key) catch {
    _ =>
      raise CompileError::ImportError("could not resolve import '\{display}'")
  }
  if contains(import_stack, canonical) {
    raise CompileError::ImportError(
      "detected cyclic import chain: \{format_cycle(import_stack, canonical)}",
    )
  }
  let stack2 = import_stack.copy()
  stack2.push(canonical)
  let (import_ast, errors) = @parser.parse(content)
  if !errors.is_empty() {
    raise CompileError::ImportError(
      "parse error in imported file '\{display}': \{errors}",
    )
  }
  let import_ir = Map::new()
  compile_map_nodes_with_imports(import_ir, import_ast.nodes, resolver, stack2)
  if imp.spread {
    // Spread import in map context: only maps are allowed.
    let m2 = select_map_for_spread(import_ir, selector) catch {
      _ =>
        raise CompileError::ImportError("cannot spread import non map into map")
    }
    merge_into_map(ir, m2)
    return
  }

  // Non-spread import at map node level: create field named by the imported file key.
  let (primary, composite) = select_primary_and_composite(import_ir, selector)
  let field_name = extract_filename(display)
  let field0 = ir.ensure_field(field_name)
  ir.set_field(Field::new(field0.name, primary, composite, field0.references))
}

///|
fn split_import_file_and_selector(
  imp : @ast.Import,
) -> (String, Array[String], String) {
  // `@a.b.c` means: import file "a", then select path ["b", "c"].
  // The resolver key is the file part only (d2 parity).
  let file = if imp.path.length() > 0 {
    join_import_pre(imp.pre, imp.path[0].content())
  } else {
    join_import_pre(imp.pre, "")
  }
  let selector : Array[String] = []
  for i = 1; i < imp.path.length(); i = i + 1 {
    selector.push(imp.path[i].content())
  }
  let display = if selector.length() == 0 {
    file
  } else {
    file + "." + selector.join(".")
  }
  (file, selector, display)
}

///|
fn join_import_pre(pre : String, file : String) -> String {
  let combined = pre + file
  let is_abs = combined.length() > 0 && combined.to_array()[0] == '/'
  let parts : Array[String] = []
  for seg in combined.split("/").collect() {
    let s = seg.to_string()
    if s == "" || s == "." {
      continue
    }
    if s == ".." {
      if parts.length() > 0 && parts[parts.length() - 1] != ".." {
        let _ = parts.pop()
      } else if !is_abs {
        parts.push("..")
      }
      continue
    }
    parts.push(s)
  }
  let joined = parts.join("/")
  if is_abs {
    if joined == "" {
      "/"
    } else {
      "/" + joined
    }
  } else {
    joined
  }
}

///|
fn resolve_import_content(
  resolver : ImportResolver,
  file_key : String,
) -> (String, String) raise CompileError {
  match (resolver.0)(file_key) {
    Some(content) => (content, file_key)
    None =>
      if file_key.has_suffix(".d2") {
        raise CompileError::ImportError("missing import")
      } else {
        let k2 = file_key + ".d2"
        match (resolver.0)(k2) {
          Some(content) => (content, k2)
          None => raise CompileError::ImportError("missing import")
        }
      }
  }
}

///|
fn contains(arr : Array[String], s : String) -> Bool {
  for x in arr {
    if x == s {
      return true
    }
  }
  false
}

///|
fn format_cycle(stack : Array[String], next : String) -> String {
  let parts = stack.copy()
  parts.push(next)
  parts.join(" -> ")
}

///|
fn select_primary_and_composite(
  root : Map,
  selector : Array[String],
) -> (Scalar?, Composite?) raise CompileError {
  if selector.length() == 0 {
    return (None, Some(Composite::Map(root)))
  }
  let f = select_field(root, selector)
  (f.primary, f.composite)
}

///|
fn select_map_for_spread(
  root : Map,
  selector : Array[String],
) -> Map raise CompileError {
  if selector.length() == 0 {
    return root
  }
  let f = select_field(root, selector)
  match f.composite {
    Some(Composite::Map(m)) => m
    _ => raise CompileError::ImportError("not a map")
  }
}

///|
fn select_field(
  root : Map,
  selector : Array[String],
) -> Field raise CompileError {
  let mut current = root
  for i = 0; i < selector.length(); i = i + 1 {
    let name = selector[i]
    match current.get_field(name) {
      Some(f) =>
        if i == selector.length() - 1 {
          return f
        } else {
          match f.composite {
            Some(Composite::Map(m)) => current = m
            _ => raise CompileError::ImportError("not a map")
          }
        }
      None => raise CompileError::ImportError("not found")
    }
  }
  raise CompileError::ImportError("not found")
}

///|
fn merge_into_map(dst : Map, src : Map) -> Unit {
  for f in src.fields {
    dst.set_field(f)
  }
  for e in src.edges {
    dst.set_edge(e)
  }
}

///|
/// Extract filename without extension from path
fn extract_filename(path : String) -> String {
  // Find last /
  let mut last_slash = -1
  for i in 0.. Unit raise CompileError {
  // Handle ampersand filters
  if key.ampersand || key.not_ampersand {
    return
  }
  // Handle edges first
  if key.edges.length() > 0 {
    compile_edges_with_imports(ir, key, resolver, import_stack)
    return
  }
  // Handle regular key
  guard key.key is Some(key_path) else { return }
  // Check for glob pattern
  if key_path.has_glob() {
    compile_glob_key(ir, key_path, key)
    return
  }
  let path = key_path.to_strings()
  if path.length() == 0 {
    return
  }
  // Navigate/create the field path
  let mut current_map = ir
  for i = 0; i < path.length() - 1; i = i + 1 {
    let name = path[i]
    let elem = key_path.path[i]
    let field = record_field_reference(
      current_map,
      current_map.ensure_field(name),
      elem,
      i,
      false,
    )
    match field.composite {
      Some(Composite::Map(m)) => current_map = m
      None => {
        let new_map = Map::new()
        let field = Field::new(
          field.name,
          field.primary,
          Some(Composite::Map(new_map)),
          field.references,
        )
        current_map.set_field(field)
        current_map = new_map
      }
      Some(_) =>
        raise CompileError::InvalidKey("field already has non-map value")
    }
  }
  // Set the final field
  let last_index = path.length() - 1
  let final_name = path[last_index]
  let last_elem = key_path.path[last_index]
  let field0 = current_map.ensure_field(final_name)
  let field1 = if key.primary is Some(_) {
    record_field_reference(current_map, field0, last_elem, last_index, true)
  } else {
    field0
  }
  let field = if key.value is Some(_) {
    record_field_reference(current_map, field1, last_elem, last_index, false)
  } else {
    field1
  }
  // Set primary value
  let field = match key.primary {
    Some(scalar) =>
      Field::new(
        field.name,
        Some(compile_scalar(scalar)),
        field.composite,
        field.references,
      )
    None => field
  }
  // Set value (with import handling)
  let field = match key.value {
    Some(value) =>
      compile_value_with_imports(field, value, resolver, import_stack)
    None => field
  }
  current_map.set_field(field)
}

///|
fn compile_value_with_imports(
  field : Field,
  value : @ast.Value,
  resolver : ImportResolver,
  import_stack : Array[String],
) -> Field raise CompileError {
  match value {
    @ast.Value::Scalar(s) =>
      Field::new(
        field.name,
        Some(compile_scalar(s)),
        field.composite,
        field.references,
      )
    @ast.Value::Map(m) => {
      let ir_map = Map::new()
      compile_map_nodes_with_imports(ir_map, m.nodes, resolver, import_stack)
      Field::new(
        field.name,
        field.primary,
        Some(Composite::Map(ir_map)),
        field.references,
      )
    }
    @ast.Value::Array(a) => {
      let ir_array = compile_array_with_imports(a, resolver, import_stack)
      Field::new(
        field.name,
        field.primary,
        Some(Composite::Array(ir_array)),
        field.references,
      )
    }
    @ast.Value::Import(imp) => {
      let (file_key, selector, display) = split_import_file_and_selector(imp)
      let (content, canonical) = resolve_import_content(resolver, file_key) catch {
        _ =>
          raise CompileError::ImportError(
            "could not resolve import '\{display}'",
          )
      }
      if contains(import_stack, canonical) {
        raise CompileError::ImportError(
          "detected cyclic import chain: \{format_cycle(import_stack, canonical)}",
        )
      }
      let stack2 = import_stack.copy()
      stack2.push(canonical)
      let (import_ast, errors) = @parser.parse(content)
      if !errors.is_empty() {
        raise CompileError::ImportError(
          "parse error in imported file '\{display}'",
        )
      }
      let import_ir = Map::new()
      compile_map_nodes_with_imports(
        import_ir,
        import_ast.nodes,
        resolver,
        stack2,
      )
      let (primary, composite) = select_primary_and_composite(
        import_ir, selector,
      )
      Field::new(field.name, primary, composite, field.references)
    }
    @ast.Value::BlockScalar(label, m) => {
      let field = match label {
        @ast.Value::Scalar(s) =>
          Field::new(
            field.name,
            Some(compile_scalar(s)),
            field.composite,
            field.references,
          )
        _ => field
      }
      let ir_map = Map::new()
      compile_map_nodes_with_imports(ir_map, m.nodes, resolver, import_stack)
      Field::new(
        field.name,
        field.primary,
        Some(Composite::Map(ir_map)),
        field.references,
      )
    }
  }
}

///|
fn compile_array_with_imports(
  arr : @ast.ArrayValue,
  resolver : ImportResolver,
  import_stack : Array[String],
) -> IRArray raise CompileError {
  let values : Array[Value] = []
  for node in arr.nodes {
    match node {
      @ast.ArrayNode::Comment(_, _) => ()
      @ast.ArrayNode::BlockComment(_, _) => ()
      @ast.ArrayNode::Substitution(sub) =>
        values.push(Value::Scalar(substitution_scalar(sub)))
      @ast.ArrayNode::Import(imp) => {
        let (file_key, selector, display) = split_import_file_and_selector(imp)
        let (content, canonical) = resolve_import_content(resolver, file_key) catch {
          _ =>
            raise CompileError::ImportError(
              "could not resolve import '\{display}'",
            )
        }
        if contains(import_stack, canonical) {
          raise CompileError::ImportError(
            "detected cyclic import chain: \{format_cycle(import_stack, canonical)}",
          )
        }
        let stack2 = import_stack.copy()
        stack2.push(canonical)
        let (import_ast, errors) = @parser.parse(content)
        if !errors.is_empty() {
          raise CompileError::ImportError(
            "parse error in imported file '\{display}'",
          )
        }
        let import_ir = Map::new()
        compile_map_nodes_with_imports(
          import_ir,
          import_ast.nodes,
          resolver,
          stack2,
        )
        if imp.spread {
          // Only allow spreading arrays into arrays (d2 parity).
          let f = select_field(import_ir, selector)
          match f.composite {
            Some(Composite::Array(a)) =>
              for v in a.values {
                values.push(v)
              }
            _ =>
              raise CompileError::ImportError(
                "can only spread import array into array",
              )
          }
        } else {
          let (primary, composite) = select_primary_and_composite(
            import_ir, selector,
          )
          match (primary, composite) {
            (Some(p), _) => values.push(Value::Scalar(p))
            (None, Some(Composite::Map(m))) => values.push(Value::Map(m))
            (None, Some(Composite::Array(a))) => values.push(Value::Array(a))
            (None, None) => ()
          }
        }
      }
      @ast.ArrayNode::Value(v) => {
        let compiled = compile_ast_value_with_imports(v, resolver, import_stack)
        values.push(compiled)
      }
    }
  }
  IRArray::new(values)
}

///|
fn compile_ast_value_with_imports(
  value : @ast.Value,
  resolver : ImportResolver,
  import_stack : Array[String],
) -> Value raise CompileError {
  match value {
    @ast.Value::Scalar(s) => Value::Scalar(compile_scalar(s))
    @ast.Value::Map(m) => {
      let ir_map = Map::new()
      compile_map_nodes_with_imports(ir_map, m.nodes, resolver, import_stack)
      Value::Map(ir_map)
    }
    @ast.Value::Array(a) =>
      Value::Array(compile_array_with_imports(a, resolver, import_stack))
    @ast.Value::Import(imp) => {
      let (file_key, selector, display) = split_import_file_and_selector(imp)
      let (content, canonical) = resolve_import_content(resolver, file_key) catch {
        _ =>
          raise CompileError::ImportError(
            "could not resolve import '\{display}'",
          )
      }
      if contains(import_stack, canonical) {
        raise CompileError::ImportError(
          "detected cyclic import chain: \{format_cycle(import_stack, canonical)}",
        )
      }
      let stack2 = import_stack.copy()
      stack2.push(canonical)
      let (import_ast, errors) = @parser.parse(content)
      if !errors.is_empty() {
        raise CompileError::ImportError(
          "parse error in imported file '\{display}'",
        )
      }
      let import_ir = Map::new()
      compile_map_nodes_with_imports(
        import_ir,
        import_ast.nodes,
        resolver,
        stack2,
      )
      let (primary, composite) = select_primary_and_composite(
        import_ir, selector,
      )
      match (primary, composite) {
        (Some(p), _) => Value::Scalar(p)
        (None, Some(Composite::Map(m))) => Value::Map(m)
        (None, Some(Composite::Array(a))) => Value::Array(a)
        (None, None) => Value::Scalar(Scalar::null())
      }
    }
    @ast.Value::BlockScalar(_, m) => {
      let ir_map = Map::new()
      compile_map_nodes_with_imports(ir_map, m.nodes, resolver, import_stack)
      Value::Map(ir_map)
    }
  }
}

///|
fn compile_edges_with_imports(
  ir : Map,
  key : @ast.Key,
  resolver : ImportResolver,
  import_stack : Array[String],
) -> Unit raise CompileError {
  for edge in key.edges {
    let src_paths = edge_endpoint_paths(ir, edge.src)
    let dst_paths = edge_endpoint_paths(ir, edge.dst)
    let src_has_glob = edge.src.has_glob()
    let dst_has_glob = edge.dst.has_glob()
    if src_paths.is_empty() || dst_paths.is_empty() {
      continue
    }
    fn next_index_for_group(ir : Map, base_id : EdgeID) -> Int {
      let existing = ir.find_edges(base_id)
      let mut max_idx = -1
      for e in existing {
        match e.id.index {
          Some(i) => if i > max_idx { max_idx = i }
          None => if max_idx < 0 { max_idx = 0 }
        }
      }
      max_idx + 1
    }

    fn find_exact_edge(ir : Map, id : EdgeID) -> Edge? {
      for e in ir.edges {
        if e.id == id {
          return Some(e)
        }
      }
      None
    }

    fn compile_edge_value(
      existing : Edge?,
      value : @ast.Value?,
      resolver : ImportResolver,
      import_stack : Array[String],
    ) -> (Scalar?, Map?) raise CompileError {
      match (existing, value) {
        (Some(e), Some(v)) =>
          match v {
            @ast.Value::Scalar(s) => (Some(compile_scalar(s)), e.map)
            @ast.Value::Map(m) => {
              let ir_map = Map::new()
              compile_map_nodes_with_imports(
                ir_map,
                m.nodes,
                resolver,
                import_stack,
              )
              (e.primary, Some(ir_map))
            }
            @ast.Value::BlockScalar(label, m) => {
              let updated_primary = match label {
                @ast.Value::Scalar(s) => Some(compile_scalar(s))
                _ => e.primary
              }
              let ir_map = Map::new()
              compile_map_nodes_with_imports(
                ir_map,
                m.nodes,
                resolver,
                import_stack,
              )
              (updated_primary, Some(ir_map))
            }
            _ => (e.primary, e.map)
          }
        (Some(e), None) => (e.primary, e.map)
        (None, Some(v)) =>
          match v {
            @ast.Value::Scalar(s) => (Some(compile_scalar(s)), None)
            @ast.Value::Map(m) => {
              let ir_map = Map::new()
              compile_map_nodes_with_imports(
                ir_map,
                m.nodes,
                resolver,
                import_stack,
              )
              (None, Some(ir_map))
            }
            @ast.Value::BlockScalar(label, m) => {
              let mut primary : Scalar? = None
              match label {
                @ast.Value::Scalar(s) => primary = Some(compile_scalar(s))
                _ => ()
              }
              let ir_map = Map::new()
              compile_map_nodes_with_imports(
                ir_map,
                m.nodes,
                resolver,
                import_stack,
              )
              (primary, Some(ir_map))
            }
            _ => (None, None)
          }
        (None, None) => (None, None)
      }
    }

    for src_path in src_paths {
      for dst_path in dst_paths {
        if (src_has_glob || dst_has_glob) &&
          string_path_equal(src_path, dst_path) {
          continue
        }
        let base_id = EdgeID::new(
          src_path,
          dst_path,
          edge.src_arrow == "<",
          edge.dst_arrow == ">",
        )
        match key.edge_index {
          Some(ei) =>
            if ei.glob {
              let existing = ir.find_edges(base_id)
              for e in existing {
                e.references.push(
                  EdgeReference::new(edge.range, e.id.index.unwrap_or(0), true),
                )
                let (primary, edge_map) = compile_edge_value(
                  Some(e),
                  key.value,
                  resolver,
                  import_stack,
                )
                ir.set_edge(Edge::new(e.id, primary, edge_map, e.references))
              }
            } else {
              let idx = ei.index.unwrap_or(0)
              let edge_id = EdgeID::with_index(
                base_id.src_path,
                base_id.dst_path,
                base_id.src_arrow,
                base_id.dst_arrow,
                idx,
              )
              match find_exact_edge(ir, edge_id) {
                Some(e) => {
                  e.references.push(EdgeReference::new(edge.range, idx, true))
                  let (primary, edge_map) = compile_edge_value(
                    Some(e),
                    key.value,
                    resolver,
                    import_stack,
                  )
                  ir.set_edge(Edge::new(e.id, primary, edge_map, e.references))
                }
                None => {
                  let (primary, edge_map) = compile_edge_value(
                    None,
                    key.value,
                    resolver,
                    import_stack,
                  )
                  let new_edge = Edge::new(edge_id, primary, edge_map, [
                    EdgeReference::new(edge.range, idx, true),
                  ])
                  ir.add_edge(new_edge)
                  ensure_field_path(ir, src_path)
                  ensure_field_path(ir, dst_path)
                }
              }
            }
          None => {
            let idx = next_index_for_group(ir, base_id)
            let edge_id = EdgeID::with_index(
              base_id.src_path,
              base_id.dst_path,
              base_id.src_arrow,
              base_id.dst_arrow,
              idx,
            )
            let (primary, edge_map) = compile_edge_value(
              None,
              key.value,
              resolver,
              import_stack,
            )
            let new_edge = Edge::new(edge_id, primary, edge_map, [
              EdgeReference::new(edge.range, idx, true),
            ])
            ir.add_edge(new_edge)
            ensure_field_path(ir, src_path)
            ensure_field_path(ir, dst_path)
          }
        }
      }
    }
  }
}