diff --git a/README.md b/README.md
new file mode 100644
--- /dev/null
+++ b/README.md
@@ -0,0 +1,22 @@
+# hspec-setup
+[![Hackage](https://img.shields.io/hackage/v/hspec-setup.svg)](http://hackage.haskell.org/package/hspec-setup)
+[![Hackage Deps](https://img.shields.io/hackage-deps/v/hspec-setup.svg)](http://hackage.haskell.org/package/hspec-setup)
+[![Build Status](https://travis-ci.org/yamadapc/haskell-hspec-setup.svg?branch=master)](https://travis-ci.org/yamadapc/haskell-hspec-setup)
+- - -
+Add an hspec test-suite in one command
+
+![](/demo.gif)
+
+## Installation
+```
+$ stack install hspec-setup
+```
+
+## Usage
+```
+$ hspec-setup
+```
+
+## License
+This code is licensed under the MIT license. Please refer to the
+[LICENSE](/LICENSEE) file for more information
diff --git a/bin/Main.hs b/bin/Main.hs
new file mode 100644
--- /dev/null
+++ b/bin/Main.hs
@@ -0,0 +1,7 @@
+module Main
+  where
+
+import qualified Test.Hspec.Setup
+
+main :: IO ()
+main = Test.Hspec.Setup.main
diff --git a/hspec-setup.cabal b/hspec-setup.cabal
--- a/hspec-setup.cabal
+++ b/hspec-setup.cabal
@@ -1,32 +1,96 @@
-name:                hspec-setup
-version:             0.1.1.1
-synopsis:            Add an hspec test-suite in one command
-description: @hspec-setup@ is a command-line tool for adding an hspec
-             test-suite with minimal work. See the
-             <https://github.com/yamadapc/haskell-hspec-setup GitHub README>
-             for more information.
-homepage:            https://github.com/yamadapc/haskell-hspec-setup
-license:             MIT
-license-file:        LICENSE
-author:              Pedro Tacla Yamada
-maintainer:          tacla.yamada@gmail.com
-copyright:           Copyright (c) 2015 Pedro Tacla Yamada
-category:            Development
-build-type:          Simple
-cabal-version:       >=1.10
+-- This file has been generated from package.yaml by hpack version 0.14.0.
+--
+-- see: https://github.com/sol/hpack
 
-tested-with: GHC >= 7.10
+name:           hspec-setup
+version:        0.2.0.0
+synopsis:       Add an hspec test-suite in one command
+description:    @hspec-setup@ is a command-line tool for adding an hspec
+                test-suite with minimal work. See the
+                <https://github.com/yamadapc/haskell-hspec-setup GitHub README>
+                for more information.
+homepage:       https://github.com/yamadapc/haskell-hspec-setup
+license:        MIT
+license-file:   LICENSE
+author:         Pedro Tacla Yamada
+maintainer:     tacla.yamada@gmail.com
+copyright:      Copyright (c) 2015 Pedro Tacla Yamada
+category:       Development
+build-type:     Simple
+cabal-version:  >= 1.10
 
+extra-source-files:
+    README.md
+
 source-repository head
   type: git
   location: git://github.com/yamadapc/haskell-hspec-setup
 
+library
+  hs-source-dirs:
+      src
+  build-depends:
+      ansi-wl-pprint
+    , pretty >=1.1.2.0
+    , base >=4 && <5
+    , directory
+    , directory-tree
+    , filepath
+    , process >=1.2
+    , projectroot >=0.2
+    , haskell-src-exts >=1.18.2 && <1.19
+    , split
+  exposed-modules:
+      Test.Hspec.Setup
+      System.Directory.Recursive
+  other-modules:
+      Paths_hspec_setup
+  default-language: Haskell2010
+
 executable hspec-setup
-  main-is:             Main.hs
-  build-depends:       base >=4 && <5
-                     , directory
-                     , filepath
-                     , process >= 1.2
-                     , projectroot >= 0.2
-  hs-source-dirs:      src
-  default-language:    Haskell2010
+  main-is: Main.hs
+  build-depends:
+      ansi-wl-pprint
+    , pretty >=1.1.2.0
+    , base >=4 && <5
+    , directory
+    , directory-tree
+    , filepath
+    , process >=1.2
+    , projectroot >=0.2
+    , haskell-src-exts >=1.18.2 && <1.19
+    , split
+    , hspec-setup
+  hs-source-dirs:
+      bin
+  default-language: Haskell2010
+
+test-suite hspec
+  main-is: Spec.hs
+  type: exitcode-stdio-1.0
+  build-depends:
+      ansi-wl-pprint
+    , pretty >=1.1.2.0
+    , base >=4 && <5
+    , directory
+    , directory-tree
+    , filepath
+    , process >=1.2
+    , projectroot >=0.2
+    , haskell-src-exts >=1.18.2 && <1.19
+    , split
+    , QuickCheck
+    , language-dockerfile >=0.3.4
+    , base
+    , directory
+    , bytestring
+    , cryptohash
+    , hspec
+    , hspec-setup
+  other-modules:
+      SanitySpec
+      System.Directory.RecursiveSpec
+      Test.Hspec.SetupSpec
+  hs-source-dirs:
+      test
+  default-language: Haskell2010
diff --git a/src/Main.hs b/src/Main.hs
deleted file mode 100644
--- a/src/Main.hs
+++ /dev/null
@@ -1,76 +0,0 @@
-module Main
-  where
-
-import           Control.Monad
-import           Data.List
-import           Data.Monoid
-import           System.Directory
-import           System.Directory.ProjectRoot
-import           System.Exit
-import           System.FilePath
-import           System.IO
-import           System.Process
-
-main :: IO ()
-main = getProjectRootCurrent >>= \mpr -> case mpr of
-    Nothing -> error "Couldn't find the project root"
-    Just pr -> do
-        fs <- getDirectoryContents pr
-        case find ((".cabal" ==) . takeExtension) fs of
-            Nothing -> error "Couldn't find your cabal file."
-            Just fp -> hspecSetup pr (pr </> fp)
-
-hspecTestSuite :: String
-hspecTestSuite = unlines [ ""
-                         , "test-suite hspec"
-                         , "  main-is: Spec.hs"
-                         , "  type: exitcode-stdio-1.0"
-                         , "  build-depends: base"
-                         , "               , hspec"
-                         , "               , QuickCheck"
-                         , "  hs-source-dirs: test"
-                         , "  default-language: Haskell2010"
-                         ]
-
-hspecDiscoveryFile :: String
-hspecDiscoveryFile = "{-# OPTIONS_GHC -F -pgmF hspec-discover #-}"
-
-hspecSanitySpec :: String
-hspecSanitySpec = unlines [ "module SanitySpec where"
-                          , ""
-                          , "import Test.Hspec"
-                          , ""
-                          , "spec = describe \"when I have tests\" $"
-                          , "    it \"I have sanity\" $ True `shouldBe` True"
-                          ]
-
-hspecSetup :: FilePath -> FilePath -> IO ()
-hspecSetup pr fp = do
-    c <- getCurrentDirectory
-
-    putStrLn $ "Adding test-suite to " <> makeRelative c fp <> "..."
-    cabalContents <- readFile fp
-
-    when ("type: exitcode-stdio-1.0" `isInfixOf` cabalContents) $ do
-        hPutStrLn stderr "File already has test-suite. Exiting..."
-        exitFailure
-    appendFile fp hspecTestSuite
-
-    putStrLn "Creating test directory..."
-    createDirectoryIfMissing False (pr </> "test")
-
-    putStrLn "Creating test/Spec.hs discovery file..."
-    writeFile (pr </> "test" </> "Spec.hs") hspecDiscoveryFile
-
-    putStrLn "Creating test/SanitySpec.hs..."
-    writeFile (pr </> "test" </> "SanitySpec.hs") hspecSanitySpec
-
-    stackInited <- doesFileExist (pr </> "stack.yaml")
-    unless stackInited $ do
-        putStrLn "No `stack.yaml` found. Running `stack init` for you..."
-        callCommand "stack init"
-
-    putStrLn "Running tests for the first time..."
-    callCommand "stack test"
-
-    return ()
diff --git a/src/System/Directory/Recursive.hs b/src/System/Directory/Recursive.hs
new file mode 100644
--- /dev/null
+++ b/src/System/Directory/Recursive.hs
@@ -0,0 +1,41 @@
+module System.Directory.Recursive
+    (
+      getDirectoryContentsRecursive
+    , isHaskellSource
+    , pretty
+    , putDoc
+    , module System.Directory.Tree
+    )
+  where
+
+import           Debug.Trace
+import           System.Directory
+import           System.Directory.Tree
+import           System.FilePath
+import           Text.PrettyPrint.ANSI.Leijen hiding ((<$>))
+import qualified Text.PrettyPrint.ANSI.Leijen ((<$>))
+-- import           Text.PrettyPrint.HughesPJClass
+
+(<++>) = (Text.PrettyPrint.ANSI.Leijen.<$>)
+
+getDirectoryContentsRecursive
+  :: FilePath -> IO (AnchoredDirTree ())
+getDirectoryContentsRecursive = readDirectoryWith (const $ return ())
+
+instance Pretty (DirTree a) where
+    pretty (Dir dn ds) = cyan (text dn) <$$>
+        vcat (map (hang 2 . (text (replicate 2 ' ') <>) . pretty) ds)
+    pretty (File fn _) = white (text fn) <> red (text " <-")
+    pretty (Failed fn e) = red (text fn <> text (show e))
+
+isHaskellSource (File fn _) =
+    takeExtension fn == ".hs"
+isHaskellSource (Dir _ []) = False
+isHaskellSource (Dir ".stack-work" _) = False
+isHaskellSource (Dir _ ds) = -- traceShow (ds, any isHaskellSource ds) $
+    any isHaskellSource ds
+isHaskellSource _ = False
+
+recursiveLink = undefined
+
+recursiveMirror = undefined
diff --git a/src/Test/Hspec/Setup.hs b/src/Test/Hspec/Setup.hs
new file mode 100644
--- /dev/null
+++ b/src/Test/Hspec/Setup.hs
@@ -0,0 +1,206 @@
+{-# LANGUAGE LambdaCase  #-}
+module Test.Hspec.Setup
+  where
+
+import           Control.Monad
+import           Data.List
+import           Data.List.Split
+import           Data.Maybe
+import           Data.Monoid
+import           Language.Haskell.Exts
+import           System.Directory
+import           System.Directory.ProjectRoot
+import           System.Directory.Recursive
+import           System.Environment
+import           System.Exit
+import           System.FilePath
+import           System.IO
+import           System.Process
+
+main :: IO ()
+main = do
+    mfp <- getData
+    case mfp of
+        Left e -> panic e
+        Right fp -> executeCommand fp
+  where
+    panic e = do
+        hPutStrLn stderr e
+        exitFailure
+    executeCommand (pr, fp) = do
+        as <- getArgs
+        case as of
+            ("--generate":rest) -> hspecGenerate pr (listToMaybe rest)
+            ("-g":rest) -> hspecGenerate pr (listToMaybe rest)
+            [] -> hspecSetup pr fp
+            _ -> error $ unlines [ "Usage: hspec-setup [-g|--generate]"
+                                 , "    hspec-setup                   Adds a test-suite to your project"
+                                 , "    hspec-setup --generate [fp]   Generates tests for a certain module"
+                                 , "    hspec-setup -g"
+                                 ]
+
+data ManifestFilePath = CabalFile FilePath
+                      | HpackFile FilePath
+
+unManifestFilePath :: ManifestFilePath -> FilePath
+unManifestFilePath (CabalFile fp) = fp
+unManifestFilePath (HpackFile fp) = fp
+
+type Options = (FilePath, ManifestFilePath)
+
+getData :: IO (Either String Options)
+getData = getProjectRootCurrent >>=
+    \case
+        Nothing -> return $ Left "Couldn't find the project root"
+        Just pr -> do
+            fs <- getDirectoryContents pr
+            case find ("package.yaml" ==) fs of
+                Just fp -> return $ Right (pr, HpackFile fp)
+                Nothing -> case find ((".cabal" ==) . takeExtension) fs of
+                    Just fp -> return $ Right (pr, CabalFile fp)
+                    Nothing -> return $ Left
+                        "Couldn't find your package manifest file (package.yaml/*.cabal)"
+
+
+hspecTestSuite :: String
+hspecTestSuite = unlines [ ""
+                         , "test-suite hspec"
+                         , "  main-is: Spec.hs"
+                         , "  type: exitcode-stdio-1.0"
+                         , "  build-depends: base"
+                         , "               , hspec"
+                         , "               , QuickCheck"
+                         , "  hs-source-dirs: test"
+                         , "  default-language: Haskell2010"
+                         ]
+
+hspecHpackTestSuite :: String
+hspecHpackTestSuite = unlines [ ""
+                              , "tests:"
+                              , "  hspec:"
+                              , "    main: Spec.hs"
+                              , "    source-dirs: test"
+                              , "    dependencies:"
+                              , "    - base"
+                              , "    - hspec"
+                              , "    - QuickCheck"
+                              ]
+
+hspecDiscoveryFile :: String
+hspecDiscoveryFile = "{-# OPTIONS_GHC -F -pgmF hspec-discover #-}"
+
+hspecSanitySpec :: String
+hspecSanitySpec = unlines [ "module SanitySpec where"
+                          , ""
+                          , "import Test.Hspec"
+                          , ""
+                          , "spec = describe \"when I have tests\" $"
+                          , "    it \"I have sanity\" $ True `shouldBe` True"
+                          ]
+
+hspecGenerate :: FilePath -> Maybe FilePath -> IO ()
+hspecGenerate pr (Just target) = parseFile target >>= \case
+    ParseOk (Module _ (Just (ModuleHead _ (ModuleName _ moduleName) _ (Just (ExportSpecList _ exportSpecs)))) _ _ _) ->
+        go moduleName (exportsFromSpecs exportSpecs)
+    ParseOk (Module _ (Just (ModuleHead _ (ModuleName _ moduleName) _ Nothing)) _ _ moduleDecls) ->
+        go moduleName (exportsFromDecls moduleDecls)
+    ParseOk (Module l Nothing _ _ moduleDecls) -> error (target <> " needs a `module` name")
+    ParseOk h -> error ("Don't know how to generate code from" <> (show h))
+    e -> error (show e)
+  where
+    exportsFromSpecs = mapMaybe mexportFromSpec
+      where
+        mexportFromSpec (EVar _ (UnQual _ n)) = Just (prettyPrint n)
+        mexportFromSpec _ = Nothing
+    exportsFromDecls = mapMaybe mexportFromDecl
+      where
+        mexportFromDecl (FunBind _ ((Match _ n _ _ _):_)) = Just (prettyPrint n)
+        mexportFromDecl _ = Nothing
+    go :: String -> [String] -> IO ()
+    go moduleName moduleExports = do
+        let moduleNameFp = (foldl (</>) "./test" $ splitOn "." (moduleName <> "Spec")) <> ".hs"
+            moduleNameFallbackFp = (foldl (</>) "./test" $ splitOn "." (moduleName <> "AutogenSpec")) <> ".hs"
+            testSuite mn = init $ unlines $
+                        [ "module " <> mn <> " where"
+                        , ""
+                        , "import           " <> moduleName
+                        , ""
+                        , "import           Test.Hspec"
+                        , ""
+                        , "spec :: Spec"
+                        , "spec = do"
+                        ] <> concatMap decToDescribe moduleExports
+
+        createDirectoryIfMissing True (takeDirectory moduleNameFp)
+        e <- doesFileExist moduleNameFp
+
+        when e $ do
+            hPutStrLn stderr ("Refusing to overwrite " <> moduleNameFp)
+            hPutStrLn stderr "Wrote test-suite to:"
+            putStrLn moduleNameFallbackFp
+            writeFile moduleNameFallbackFp (testSuite (moduleName <> "AutogenSpec"))
+            exitSuccess
+
+        writeFile moduleNameFp (testSuite (moduleName <> "Spec"))
+      where
+        decToDescribe decName = [
+            "    describe \"" <> decName <> "\" $ do"
+          , "        it \"works\" pending"
+                                ]
+hspecGenerate pr Nothing = error "Interactive mode not implemented"
+
+hspecSetup :: FilePath -> ManifestFilePath -> IO ()
+hspecSetup pr mfp = do
+    let fp = unManifestFilePath mfp
+
+    c <- getCurrentDirectory
+
+    putStrLn $ "Adding test-suite to " <> makeRelative c fp <> "..."
+
+    appendHspecSuite pr mfp
+
+    putStrLn "Creating test directory..."
+    createDirectoryIfMissing False (pr </> "test")
+
+    putStrLn "Creating test/Spec.hs discovery file..."
+    writeFile (pr </> "test" </> "Spec.hs") hspecDiscoveryFile
+
+    putStrLn "Creating test/SanitySpec.hs..."
+    writeFile (pr </> "test" </> "SanitySpec.hs") hspecSanitySpec
+
+    stackInited <- doesFileExist (pr </> "stack.yaml")
+    unless stackInited $ do
+        putStrLn "No `stack.yaml` found. Running `stack init` for you..."
+        callCommand "stack init"
+
+    putStrLn "Running tests for the first time..."
+    callCommand "stack test"
+
+    return ()
+
+appendHspecSuite :: FilePath -> ManifestFilePath -> IO ()
+appendHspecSuite pr (HpackFile fp) = do
+    hpackContents <- readFile (pr </> fp)
+    let hpackLines = lines hpackContents
+        mi = elemIndex "tests:" hpackLines
+    case mi of
+        Nothing -> do
+            when ("hspec" `isInfixOf` hpackContents) $ do
+                hPutStrLn stderr "File already has test-suite. Exiting..."
+                exitFailure
+            appendFile (pr </> fp) hspecHpackTestSuite
+        Just i -> do
+            let (beforeTests, afterTests) = splitAt i hpackLines :: ([String], [String])
+            writeFile (pr </> fp) $ unlines [ unlines beforeTests
+                                            , unlines (drop 1 (lines hspecHpackTestSuite))
+                                            , unlines afterTests
+                                            ]
+appendHspecSuite pr (CabalFile fp) = do
+    cabalContents <- readFile (pr </> fp)
+
+    when ("type: exitcode-stdio-1.0" `isInfixOf` cabalContents) $ do
+        hPutStrLn stderr "File already has test-suite. Exiting..."
+        exitFailure
+
+    appendFile (pr </> fp) hspecTestSuite
+
diff --git a/test/SanitySpec.hs b/test/SanitySpec.hs
new file mode 100644
--- /dev/null
+++ b/test/SanitySpec.hs
@@ -0,0 +1,7 @@
+module SanitySpec where
+
+import           Test.Hspec
+
+spec :: Spec
+spec = describe "when I have tests" $
+    it "I have sanity" $ True `shouldBe` True
diff --git a/test/Spec.hs b/test/Spec.hs
new file mode 100644
--- /dev/null
+++ b/test/Spec.hs
@@ -0,0 +1,1 @@
+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}
diff --git a/test/System/Directory/RecursiveSpec.hs b/test/System/Directory/RecursiveSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/System/Directory/RecursiveSpec.hs
@@ -0,0 +1,16 @@
+module System.Directory.RecursiveSpec where
+
+import           System.Directory.Recursive
+
+import           Test.Hspec
+
+spec :: Spec
+spec = do
+    describe "getDirectoryContentsRecursive" $ do
+        it "works" pending
+    describe "isHaskellSource" $ do
+        it "works" pending
+    describe "pretty" $ do
+        it "works" pending
+    describe "putDoc" $ do
+        it "works" pending
diff --git a/test/Test/Hspec/SetupSpec.hs b/test/Test/Hspec/SetupSpec.hs
new file mode 100644
--- /dev/null
+++ b/test/Test/Hspec/SetupSpec.hs
@@ -0,0 +1,105 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE PackageImports    #-}
+module Test.Hspec.SetupSpec
+  where
+
+import           Control.Exception
+import           "cryptohash" Crypto.Hash
+import qualified Data.ByteString.Char8 as ByteString (pack)
+import           Data.Maybe
+import           Data.Monoid
+import           Language.Dockerfile
+import           System.Directory
+import           System.Exit
+import           System.Process
+
+import           Test.Hspec
+
+-- | An image with hspec-setup and an empty project
+newProjectImage :: EDockerfileM ()
+newProjectImage = do
+    from "haskell:7.10.3"
+    run "cabal update"
+
+    add "./hspec-setup.cabal" "/hspec-setup/"
+    workdir "/hspec-setup"
+    run "cabal install --only-dependencies -j"
+
+    add "." "/hspec-setup"
+    run "cabal install -j"
+
+    run "mkdir /app"
+    workdir "/app"
+    run "cabal init -m -n"
+    run "echo 'main = print 10' > /app/Main.hs"
+    run "touch /app/LICENSE"
+    run "stack init"
+    run "stack build"
+
+eitherToMaybe (Left _) = Nothing
+eitherToMaybe (Right r) = r
+
+-- | Builds the image defined in this EDockerfileM block by writting it to a
+-- temporary file
+dockerBuild :: EDockerfileM () -> IO String
+dockerBuild img = do
+    let imgPre = "hspec-setup-tests-autogen"
+        imgStr = toDockerfileStr img
+        imgHash = show (hash (ByteString.pack imgStr) :: Digest MD5)
+        imgName = imgPre <> ":" <> imgHash
+        imgFp = "./" <> imgPre <> "." <> imgHash <> ".dockerfile"
+    mimg <- eitherToMaybe <$> (try (
+            fmap (take 2) . listToMaybe . map words . lines <$>
+            readCreateProcess (shell ("docker images | grep " <> imgPre)) "") :: IO (Either SomeException (Maybe [String])))
+            :: IO (Maybe [String])
+
+    -- print mimg
+    -- print (Just [imgPre, imgHash])
+
+    if mimg == Just [imgPre, imgHash]
+        then return imgName
+        else do
+            bracket_
+                (writeFile imgFp imgStr)
+                (removeFile imgFp)
+                (callCommand $ "docker build -t " <> imgName <> " -f " <> imgFp <> " .")
+            return imgName
+
+-- | Runs the image with the tag passed-in
+dockerRun :: String -> IO ()
+dockerRun imgName = callCommand $ "docker run --rm " <> imgName
+
+withImage :: EDockerfileM () -> (String -> IO b) -> IO b
+withImage img action = do
+    {-callCommand $ unwords [ "docker images |"-}
+                          {-, "grep hspec-setup-tests-autogen |"-}
+                          {-, "awk '{print $3}' |"-}
+                          {-, "tail -n +3 |"-}
+                          {-, "xargs docker rmi -f"-}
+                          {-]-}
+    bracket
+        (dockerBuild img)
+        (const (return ()))
+        -- (\imgTag -> callCommand ("docker rmi -f " <> imgTag))
+        action
+
+spec :: Spec
+spec = do
+    describe "given an empty cabal project" $
+        it "works" $ do
+            let img = newProjectImage >> cmd "hspec-setup"
+            withImage img $ \tag ->
+                dockerRun tag
+
+    describe "given an empty hpack project" $
+        it "works" $ do
+            let img = newProjectImage >> do
+                    run "stack install hpack-convert"
+                    run "hpack-convert"
+                    run "rm *.cabal"
+                    run "echo 'library:\\n  dependencies:\\n  - base' >> package.yaml"
+                    run "cat package.yaml"
+                    run "hspec-setup"
+                    cmd "cat package.yaml | grep 'tests:'"
+            withImage img $ \tag ->
+                dockerRun tag