import.go

package convert

import (
    "errors"
    "fmt"
    "go/ast"
)

/*
 * Given the root node of an AST, returns the node containing the
 * import statements for the file.
 */
func importsForRootNode(rootNode *ast.File) (imports *ast.GenDecl, err error) {
    for _, declaration := range rootNode.Decls {
        decl, ok := declaration.(*ast.GenDecl)
        if !ok || len(decl.Specs) == 0 {
            continue
        }

        _, ok = decl.Specs[0].(*ast.ImportSpec)
        if ok {
            imports = decl
            return
        }
    }

    err = errors.New(fmt.Sprintf("Could not find imports for root node:\n\t%#v\n", rootNode))
    return
}

/*
 * Removes "testing" import, if present
 */
func removeTestingImport(rootNode *ast.File) {
    importDecl, err := importsForRootNode(rootNode)
    if err != nil {
        panic(err.Error())
    }

    var index int
    for i, importSpec := range importDecl.Specs {
        importSpec := importSpec.(*ast.ImportSpec)
        if importSpec.Path.Value == "\"testing\"" {
            index = i
            break
        }
    }

    importDecl.Specs = append(importDecl.Specs[:index], importDecl.Specs[index+1:]...)
}

/*
 * Adds import statements for onsi/ginkgo, if missing
 */
func addGinkgoImports(rootNode *ast.File) {
    importDecl, err := importsForRootNode(rootNode)
    if err != nil {
        panic(err.Error())
    }

    if len(importDecl.Specs) == 0 {
        // TODO: might need to create a import decl here
        panic("unimplemented : expected to find an imports block")
    }

    needsGinkgo := true
    for _, importSpec := range importDecl.Specs {
        importSpec, ok := importSpec.(*ast.ImportSpec)
        if !ok {
            continue
        }

        if importSpec.Path.Value == "\"github.com/onsi/ginkgo\"" {
            needsGinkgo = false
        }
    }

    if needsGinkgo {
        importDecl.Specs = append(importDecl.Specs, createImport(".", "\"github.com/onsi/ginkgo\""))
    }
}

/*
 * convenience function to create an import statement
 */
func createImport(name, path string) *ast.ImportSpec {
    return &ast.ImportSpec{
        Name: &ast.Ident{Name: name},
        Path: &ast.BasicLit{Kind: 9, Value: path},
    }
}