module Main where

import qualified CreateBinding
import qualified Type
import Parser.Signature (Identifier(Identifier))
import CreateBinding
         (formatForeignCall, formatCArrayWrapper, parseHeader,
          foreignFromParameters, arrayTypeFromBool)

import qualified Text.ParserCombinators.Parsec as Parsec
import qualified Text.ParserCombinators.Parsec.Error as ParsecError

import qualified Options.Applicative as OP

import qualified System.IO as IO
import qualified System.Path.PartClass as PathClass
import qualified System.Path.IO as PathIO
import qualified System.Path as Path
import System.Path.Directory (createDirectoryIfMissing)
import System.Path (takeFileName, (</>))

import Text.Printf (printf)

import qualified Data.Csv as Csv
import qualified Data.HashMap.Lazy as HashMap
import qualified Data.Map as Map
import qualified Data.ByteString.Lazy.Char8 as BL
import qualified Data.ByteString.Char8 as B
import qualified Data.Foldable as Fold
import qualified Data.List.HT as ListHT
import qualified Data.List as List
import Data.Traversable (Traversable, for, sequenceA, traverse)
import Data.Foldable (Foldable)
import Data.Vector (Vector)
import Data.ByteString (ByteString)
import Data.HashMap.Lazy (HashMap)
import Data.Map (Map)
import Data.Csv ((.:))
import Data.Tuple.HT (mapSnd, thd3)
import Data.Maybe (isNothing, mapMaybe)
import Data.Bool.HT (if')
import Data.Char (toUpper)
import Data.Monoid ((<>))

import qualified Control.Monad.Exception.Synchronous as ME
import qualified Control.Monad.Trans.Writer as MW
import qualified Control.Applicative.HT as App
import qualified Control.Functor.HT as Func
import Control.Monad (when, join, guard)
import Control.Applicative (Applicative, liftA2, (<$>), (<*>))



class Type.Unification f => Generalize f where
   formatGenericSig :: f Type.Poly -> String

instance Generalize Type.Foreign where
   formatGenericSig = Type.formatForeign

instance Generalize Type.Wrapper where
   formatGenericSig wrapper@(Type.Wrapper inputs _returnType _outputs) =
      concatMap
         (\(typ,dims,dir) ->
            printf "%s -> " (Type.formatArray (arrayTypeFromBool dir) typ dims))
         inputs ++
      "IO " ++ CreateBinding.formatOutputTuple wrapper

prepareClassCall ::
   (Generalize f) => (String, f Type.Poly) -> (String, String, String)
prepareClassCall (name, func) =
   let typeName = map toUpper name
       sig = formatGenericSig func
       newtypeDef =
         printf "newtype %s a = %s {get%s :: %s}" typeName typeName typeName sig
   in  (typeName, sig, newtypeDef)

formatClassCall ::
   (Generalize f) => Char -> Char -> String -> f Type.Poly -> String
formatClassCall singleAbbr doubleAbbr name typ =
   let (typeName, sig, newtypeDef) = prepareClassCall (name, typ)
       typeBranch :: Char -> String
       typeBranch abbr = printf "(%s %c.%s)" typeName abbr name
   in  unlines $
      "" :
      "" :
      newtypeDef :
      "" :
      printf "%s :: Class.Real a => %s" name sig :
      printf "%s = get%s $ Class.switchReal %s %s"
         name typeName (typeBranch singleAbbr) (typeBranch doubleAbbr) :
      []

formatClassCallRC ::
   (Generalize f) => String -> String -> String -> f Type.Poly -> String
formatClassCallRC genName realName complexName typ =
   let (typeName, sig, newtypeDef) = prepareClassCall (genName, typ)
       typeBranch :: Char -> String -> String
       typeBranch abbr name = printf "(%s %c.%s)" typeName abbr name
   in  unlines $
      "" :
      "" :
      newtypeDef :
      "" :
      printf "%s :: Class.Floating a => %s" genName sig :
      printf "%s = get%s $ Class.switchFloating %s %s %s %s" genName typeName
         (typeBranch 'S' realName) (typeBranch 'D' realName)
         (typeBranch 'C' complexName) (typeBranch 'Z' complexName) :
      []


insertExportList :: [String] -> String -> String
insertExportList exports modul =
   (\(prefix,suffixes) ->
      map fst prefix ++
      case suffixes of
         [] -> []
         (_,suffix):_ ->
            "(\n" ++ unlines (map (printf "   %s,") exports) ++
            "   ) " ++ suffix) $
   break (List.isPrefixOf "where" . snd) $
   zip modul (ListHT.tails modul)

writeHeader ::
   (PathClass.AbsRel ar) =>
   Path.Dir ar -> B.ByteString -> Maybe [String] -> IO (Path.File ar)
writeHeader dstDir modul mexports = do
   let makePath =
         either (ioError . userError) (return . (dstDir </>)) . Path.parse
   tmplFile <- makePath $ B.unpack modul ++ "_tmpl.hs"
   dstFile <- makePath $ B.unpack modul ++ ".hs"
   tmpl <- PathIO.readFile tmplFile
   PathIO.writeFile dstFile $
      "-- Do not edit! Automatically generated by create-lapack-ffi.\n" ++
      maybe tmpl (flip insertExportList tmpl) mexports
   return dstFile

indent :: String -> String
indent = unlines . map ("    " ++) . lines

formatFileError :: (PathClass.AbsRel ar) => Path.File ar -> String -> String
formatFileError path msg =
   printf "\n%s:1:1:\n%s" (Path.toString path) (indent msg)

formatParseError :: Parsec.ParseError -> String
formatParseError err =
   let source = Parsec.errorPos err
   in  printf "\n%s:%d:%d:\n%s"
         (Parsec.sourceName source)
         (Parsec.sourceLine source)
         (Parsec.sourceColumn source)
         (indent $ dropWhile ('\n'==) $
          ParsecError.showErrorMessages
            "or" "unknown parse error" "expecting" "unexpected" "end of input"
            (ParsecError.errorMessages err))

printError :: String -> IO ()
printError = IO.hPutStr IO.stderr

printFileError :: (PathClass.AbsRel ar) => Path.File ar -> String -> IO ()
printFileError path msg = printError $ formatFileError path msg


floatColumn, doubleColumn,
   complexFloatColumn, complexDoubleColumn,
   realColumn, complexColumn, genericColumn :: ByteString
floatColumn = B.pack "Float"
doubleColumn = B.pack "Double"
complexFloatColumn = B.pack "ComplexFloat"
complexDoubleColumn = B.pack "ComplexDouble"

realColumn = B.pack "Real"
complexColumn = B.pack "Complex"
genericColumn = B.pack "Generic"

data Unified a =
   Unified {
      unifiedGeneric, unifiedReal, unifiedComplex :: a
   }

instance Foldable Unified where
   foldMap f (Unified g r c) = f g <> f r <> f c

instance Functor Unified where
   fmap f (Unified g r c) = Unified (f g) (f r) (f c)

instance Traversable Unified where
   sequenceA (Unified g r c) = App.lift3 Unified g r c

instance Applicative Unified where
   pure a = Unified a a a
   Unified fg fr fc <*> Unified g r c = Unified (fg g) (fr r) (fc c)

unifiedColumns :: Unified ByteString
unifiedColumns = Unified genericColumn realColumn complexColumn

unifiedColumnsMap :: HashMap ByteString ()
unifiedColumnsMap =
   HashMap.fromList $ map (flip (,) ()) $ Fold.toList unifiedColumns

data FunctionSet =
   FunctionSet (Unified (Bool, String)) (Map ByteString (Bool, Path.RelFile))

newtype RelFile = RelFile {getRelFile :: Maybe (Bool, Path.RelFile)}

instance Csv.FromField RelFile where
   parseField str =
      let parsePath = either fail (return . Just) . Path.parse . B.unpack
      in fmap RelFile $
         case B.uncons str of
            Nothing -> return Nothing
            Just ('!', path) -> fmap ((,) False) <$> parsePath path
            Just _ -> fmap ((,) True) <$> parsePath str

checkExcl :: String -> (Bool, String)
checkExcl ('!':path) = (False, path)
checkExcl path = (True, path)

instance Csv.FromNamedRecord FunctionSet where
   parseNamedRecord m = do
      names <- fmap checkExcl <$> traverse (m.:) unifiedColumns
      when (null $ snd $ unifiedGeneric names) $ fail "empty Generic name"
      set <- Csv.parseNamedRecord $ HashMap.difference m unifiedColumnsMap
      return $ FunctionSet names $ Map.mapMaybe getRelFile set


type URL = String

processFunction ::
   (PathClass.AbsRel ar) =>
   Bool ->
   (Path.File ar, Path.File ar) ->
   URL -> String -> Path.File ar ->
   IO (Maybe (Type.Foreign Type.Mono, Type.Wrapper Type.Mono))
processFunction output (ffiFile, carrayFile) url haskellName srcPath = do
   content <- PathIO.readFile srcPath
   let msignature = do
         sig <-
            fmap MW.runWriter $
            ME.mapException formatParseError $ ME.fromEither $
            Parsec.parse parseHeader (Path.toString srcPath) content
         when
            (elem (thd3 $ fst sig)
               [Just Type.ComplexSingle, Just Type.ComplexDouble]) $
            ME.throw $
               formatFileError srcPath "complex function result not supported"
         return sig
   case msignature of
      ME.Exception msg -> printError msg >> return Nothing
      ME.Success ((Identifier name, params, returnType), msgs) -> do
         mapM_ (printFileError srcPath) msgs
         let haddock =
               if null url
                 then "\n"
                 else
                     printf "\n-- | <%s>\n"
                        (printf url (Path.toString $ takeFileName srcPath)
                           :: String)
         let foreignSig = foreignFromParameters params returnType
         carrayWrapper <- do
            let (mono,poly) =
                  formatCArrayWrapper (haskellName, params, returnType)
            when output $ PathIO.appendFile carrayFile $ haddock ++ mono
            return poly
         when output $ PathIO.appendFile ffiFile $
            haddock ++ formatForeignCall ((haskellName, name), foreignSig)
         return $ Just (foreignSig, carrayWrapper)

unifyFunctions ::
   (Type.Unification f, PathClass.AbsRel ar) =>
   String ->
   Map ByteString (Path.File ar, f Type.Mono) ->
   IO (Unified (Maybe (f Type.Poly)))
unifyFunctions msg monoSigs = do
   let unify uni (xPath,xSig) (yPath,ySig) = do
         let m = Type.unifySignature uni xSig ySig
         when (isNothing m) $ do
            printFileError xPath ""
            printFileError yPath msg
         return m
   real <-
      fmap join $ sequenceA $
      liftA2
         (unify Type.unifyPrecision)
         (Map.lookup floatColumn monoSigs)
         (Map.lookup doubleColumn monoSigs)
   complex <-
      fmap join $ sequenceA $
      liftA2
         (unify Type.unifyPrecision)
         (Map.lookup complexFloatColumn monoSigs)
         (Map.lookup complexDoubleColumn monoSigs)
   let generic =
         join $ liftA2 (Type.unifySignature Type.unifyRealComplex) real complex
   return $ Unified generic real complex


writePolymorphic ::
   (PathClass.AbsRel ar) =>
   Path.Dir ar ->
   ByteString ->
   [(String, String)] ->
   IO ()
writePolymorphic dir name polySigs = do
   file <- writeHeader dir name (Just $ map fst polySigs)
   mapM_ (PathIO.appendFile file . snd) polySigs

writeUnified ::
   (PathClass.AbsRel ar, Generalize f) =>
   Path.Dir ar ->
   Vector (Unified (Bool,String), Unified (Maybe (f Type.Poly))) ->
   IO ()
writeUnified dstDir unifiedSigVec = do
   let unifiedSigs = map maskOutput $ Fold.toList unifiedSigVec
   writePolymorphic dstDir realColumn
      (mapMaybe
         (\(Unified _ (name,func) _) ->
            (,) name . formatClassCall 'S' 'D' name <$> func)
         unifiedSigs)

   writePolymorphic dstDir complexColumn
      (mapMaybe
         (\(Unified _ _ (name,func)) ->
            (,) name . formatClassCall 'C' 'Z' name <$> func)
         unifiedSigs)

   writePolymorphic dstDir genericColumn
      (mapMaybe
         (\(Unified (genericName,func) (realName,_) (complexName,_)) ->
            (,) genericName .
                  formatClassCallRC genericName realName complexName
               <$> func)
         unifiedSigs)

maskOutput ::
   (Unified (Bool,String), Unified (Maybe (f Type.Poly))) ->
   Unified (String, Maybe (f Type.Poly))
maskOutput =
   uncurry $ liftA2 (\(output,name) func -> (name, guard output >> func))


pathArgument ::
   (PathClass.FileDir fd) =>
   OP.Mod OP.ArgumentFields (Path.AbsRel fd) -> OP.Parser (Path.AbsRel fd)
pathArgument = OP.argument (OP.eitherReader Path.parse)

optParser ::
   OP.Parser
      (URL, (Path.AbsRelFile, Path.AbsRelDir, Path.AbsRelDir, Path.AbsRelDir))
optParser =
   App.lift2 (,)
      (OP.strOption $
         OP.long "fortran-url" <>
         OP.metavar "URL" <>
         OP.value "" <>
         OP.help "Printf template for URL of online help") $
   App.lift4 (,,,)
      (pathArgument $
         OP.metavar "CSV" <>
         OP.help "Spreadsheet containing grouped Fortran file names")
      (pathArgument $
         OP.metavar "SRCDIR" <>
         OP.help "Input directory containing Fortran files")
      (pathArgument $
         OP.metavar "FFIDIR" <>
         OP.help "Output directory for Haskell FFI modules")
      (pathArgument $
         OP.metavar "CARRAYDIR" <>
         OP.help "Output directory for Haskell CArray wrapper modules")

optionInfo :: OP.Parser a -> OP.ParserInfo a
optionInfo parser =
   OP.info
      (OP.helper <*> parser)
      (OP.fullDesc <>
       OP.progDesc
         "Automated generation of low-level and mid-level bindings for LAPACK")


main :: IO ()
main = do
   (url, (csvFile, srcDir, ffiDir, carrayDir)) <-
      OP.execParser $ optionInfo optParser
   createDirectoryIfMissing True ffiDir
   createDirectoryIfMissing True carrayDir

   (headers, functionSets) <-
      either (ioError . userError) return . Csv.decodeByName
       =<< BL.readFile (Path.toString csvFile)

   dstFiles <-
      fmap Map.fromList $
      for (filter (not . flip HashMap.member unifiedColumnsMap) $
           Fold.toList headers) $ \columnName ->
         fmap ((,) columnName) $
         liftA2 (,)
            (writeHeader ffiDir columnName Nothing)
            (writeHeader carrayDir columnName Nothing)

   sigs <-
      for functionSets $ \(FunctionSet flaggedNames files) ->
         fmap ((,) flaggedNames . Map.mapMaybe id) $
         sequenceA $
         flip Map.mapWithKey files $ \modul (output,file) -> do
            let names = snd <$> flaggedNames
                name =
                  if' (modul == floatColumn || modul == doubleColumn)
                     (unifiedReal names) $
                  if' (modul == complexFloatColumn ||
                       modul == complexDoubleColumn)
                     (unifiedComplex names) $
                  (unifiedGeneric names)
            let srcPath = srcDir </> file
            fmap ((,) srcPath) <$>
               processFunction output (dstFiles Map.! modul) url name srcPath

   let runUnification dstDir msg select =
         writeUnified dstDir =<<
         traverse (Func.mapSnd (unifyFunctions msg . fmap (mapSnd select))) sigs

   runUnification ffiDir "Cannot generalize FFI functions." fst
   runUnification carrayDir "Cannot generalize CArray functions." snd