import gleam/int
import gleam/list
import gleam/option
import gleam/result
import gleam/string
import gleedoc/parse.{type CodeBlock}
import gleedoc/scan
import simplifile
import snag

/// Configuration for test generation.
pub type Config {
  Config(
    /// Directory to write generated tests to, typically "test"
    output_dir: String,
  )
}

/// Generate test files from extracted code blocks.
pub fn generate_tests(
  blocks: List(CodeBlock),
  config: Config,
) -> Result(List(String), snag.Snag) {
  // Group blocks by the file they came from
  let by_file = group_by_file(blocks)

  by_file
  |> list.try_map(fn(pair) {
    let #(file, file_blocks) = pair
    let test_file_name = test_file_name(file)
    let output_dir = config.output_dir <> "/gleedoc"
    let test_path = output_dir <> "/" <> test_file_name
    let module_name = module_name_from_file(file)

    use _ <- result.try(
      simplifile.create_directory_all(output_dir)
      |> result.map_error(fn(err) {
        snag.new(
          "Failed to create directory: "
          <> output_dir
          <> " - "
          <> string.inspect(err),
        )
      }),
    )

    let block_imports = file_blocks |> list.flat_map(fn(b) { b.imports })
    let auto_imports = unique_imports(file, module_name)
    let all_imports = merge_imports(list.append(block_imports, auto_imports))

    let test_functions = generate_test_functions(file_blocks)

    let content =
      string.join(
        ["// Generated by gleedoc - do not edit manually", "", ..all_imports]
          |> list.append(test_functions),
        "\n",
      )

    use _ <- result.try(
      simplifile.write(test_path, content)
      |> result.map_error(fn(err) {
        snag.new(
          "Failed to write test file: "
          <> test_path
          <> " - "
          <> string.inspect(err),
        )
      }),
    )

    Ok(test_path)
  })
}

fn group_by_file(blocks: List(CodeBlock)) -> List(#(String, List(CodeBlock))) {
  blocks
  |> list.fold([], fn(groups, block) {
    let file = block.source.file
    case find_group(groups, file) {
      Ok(#(_, existing)) -> {
        list.map(groups, fn(pair) {
          case pair.0 == file {
            True -> #(file, [block, ..existing])
            False -> pair
          }
        })
      }
      Error(Nil) -> [#(file, [block]), ..groups]
    }
  })
  |> list.map(fn(pair) { #(pair.0, list.reverse(pair.1)) })
}

fn find_group(
  groups: List(#(String, List(CodeBlock))),
  file: String,
) -> Result(#(String, List(CodeBlock)), Nil) {
  list.find(groups, fn(pair) { pair.0 == file })
}

fn module_name_from_file(file: String) -> String {
  string.split_once(file, "/")
  |> result.map(fn(pair) { pair.1 })
  |> result.unwrap(file)
  |> string.replace(".gleam", "")
}

fn test_file_name(source_file: String) -> String {
  let name =
    string.split_once(source_file, "/")
    |> result.map(fn(pair) { pair.1 })
    |> result.unwrap(source_file)
    |> string.replace(".gleam", "")
    |> string.replace("/", "_")

  name <> "_gleedoc_test.gleam"
}

type Import {
  Import(module: String, names: List(String))
}

fn parse_import(line: String) -> Import {
  let rest = line |> string.trim |> string.drop_start(7) |> string.trim_start
  case string.split_once(rest, ".{") {
    Ok(#(module, names_part)) -> {
      let names =
        names_part
        |> string.replace("}", "")
        |> string.split(",")
        |> list.map(string.trim)
        |> list.filter(fn(s) { s != "" })
      Import(module, names)
    }
    Error(Nil) -> Import(rest, [])
  }
}

fn import_to_string(imp: Import) -> String {
  case imp.names {
    [] -> "import " <> imp.module
    names -> {
      let types =
        names
        |> list.filter(fn(n) { string.starts_with(n, "type ") })
        |> list.sort(string.compare)
      let values =
        names
        |> list.filter(fn(n) { !string.starts_with(n, "type ") })
        |> list.sort(string.compare)
      let names_str = string.join(list.append(types, values), ", ")
      "import " <> imp.module <> ".{" <> names_str <> "}"
    }
  }
}

fn merge_imports(imports: List(String)) -> List(String) {
  imports
  |> list.map(parse_import)
  |> list.fold([], fn(acc, imp) {
    case list_find_import(acc, imp.module) {
      Ok(existing) -> {
        let merged_names = list.unique(list.append(existing.names, imp.names))
        list.map(acc, fn(existing_imp) {
          case existing_imp.module == imp.module {
            True -> Import(..existing_imp, names: merged_names)
            False -> existing_imp
          }
        })
      }
      Error(Nil) -> [imp, ..acc]
    }
  })
  |> list.map(import_to_string)
  |> list.sort(string.compare)
}

fn list_find_import(
  imports: List(Import),
  module: String,
) -> Result(Import, Nil) {
  list.find(imports, fn(imp) { imp.module == module })
}

fn unique_imports(file: String, module_name: String) -> List(String) {
  // Try to read the source file and extract public names for unqualified imports.
  // If parsing fails, fall back to a simple qualified import.
  case simplifile.read(file) {
    Ok(source) -> {
      // Build the import for the target module itself.
      // Unwrap to [] on error or empty, then format accordingly.
      let names =
        scan.public_names(file, source)
        |> result.unwrap([])
      let target_import = case names {
        [] -> "import " <> module_name
        _ -> "import " <> module_name <> ".{" <> string.join(names, ", ") <> "}"
      }
      // Also carry over the source module's own top-level imports so that
      // doc examples can use them without having to restate them in the snippet.
      let source_imports =
        scan.module_imports(file, source)
        |> result.unwrap([])
      [target_import, ..source_imports]
    }
    Error(_) -> ["import " <> module_name]
  }
}

fn generate_test_functions(blocks: List(CodeBlock)) -> List(String) {
  blocks
  |> list.index_map(generate_test_function)
}

fn generate_test_function(block: CodeBlock, index: Int) -> String {
  let target_name = option.unwrap(block.source.target, "module")

  let test_func_name =
    sanitize_name(target_name) <> "_" <> int.to_string(index + 1) <> "_test"

  let source_info =
    "// From: "
    <> block.source.file
    <> ":"
    <> int.to_string(block.source.start_line + block.doc_line_offset - 1)

  let code = block.code |> string.trim

  string.join(
    [
      "",
      source_info,
      "pub fn " <> test_func_name <> "() {",
      "  " <> string.replace(code, "\n", "\n  "),
      "}",
    ],
    "\n",
  )
}

/// Replace characters that aren't valid in function names
fn sanitize_name(name: String) -> String {
  name
  |> string.replace(".", "_")
  |> string.replace("-", "_")
}

/// Clean up generated test files.
pub fn clean_generated(output_dir: String) -> Result(Nil, snag.Snag) {
  let dir = output_dir <> "/gleedoc"
  case simplifile.read_directory(dir) {
    Ok(files) -> {
      files
      |> list.filter(fn(f) {
        string.ends_with(f, "_gleedoc_test.gleam")
      })
      |> list.each(fn(f) {
        let path = dir <> "/" <> f
        let _ = simplifile.delete(path)
        Nil
      })
      Ok(Nil)
    }
    Error(err) -> {
      // If the output dir doesn't exist, that's fine
      case string.inspect(err) {
        "Enoent" -> Ok(Nil)
        _ ->
          Error(snag.new(
            "Failed to clean generated files: " <> string.inspect(err),
          ))
      }
    }
  }
}