{-# LANGUAGE Haskell2010
    , GeneralizedNewtypeDeriving
    , DeriveDataTypeable
 #-}
{-# OPTIONS
    -Wall
    -fno-warn-name-shadowing
 #-}

-- | INTERNAL module is used to infer Haskell types from Java types.
module CodeGen.Typomatic (
    runTypomatic,

    methodType,
    getterType,
    setterType,
    constructorType,

    ArgInfo (..),

    printJniType,

    dataTName,
    dataCName,
    newtTName,
    newtCName,
    tyclTName
) where

-- import Utils

import Foreign.Java.Utils
import Language.Java.Reflect
import qualified Language.Java.Reflect.Types as Types

import Control.Monad.State hiding (void)
import qualified Control.Monad.State as State
import Data.Functor.Identity

import Data.Generics
import Data.Strings
import qualified Data.Set as Set
import Data.Set (Set)
import qualified Data.List as List

import Haskell.X.Ops


dataTName, dataCName, newtTName, newtCName, tyclTName :: String -> String

dataTName = (++ "''")
dataCName = id -- (++ "")
newtTName = (++ "'")
newtCName = (++ "'")
tyclTName = id -- (++ "")

-- | This is the information which is ultimately
-- gathered by the use of this module.
data ArgInfo = ArgInfo {
    fSignature :: String,
    fArguments :: [TVar],
    fReturnType :: TVar,
    fArgNames :: [String],
    fJavaSignature :: String,
    fJavaReturnType :: String,
    fJniSignature :: String
  }

tr :: Eq a => a -> a -> [a] -> [a]
tr a b (x:xs)
    | a == x    = b : tr a b xs
    | otherwise = x : tr a b xs
tr _ _ [] = []

breakLast :: [a] -> ([a], a)
breakLast [a] = ([], a)
breakLast (a:as) =
    let (init', last') = breakLast as
    in  (a:init', last')
breakLast _ = error "Foreign.Java.Util.breakLast: empty list"


--------------------
-- Type variables --
--------------------

data TVar = TVar String | TVars [TVar]
    deriving (Eq, Ord, Show, Data, Typeable)

printTVar :: TVar -> String
printTVar var = case var of
    (TVar v) -> v
    (TVars vs) -> concat ["(", tail (concatMap ((' ':) . printTVar) vs), ")"]


---------------
-- Utilities --
---------------

printJniType :: JavaType -> String
printJniType t = case t of
    JBoolean -> "JNI.boolean"
    JChar    -> "JNI.char"
    JByte    -> "JNI.byte"
    JShort   -> "JNI.short"
    JInt     -> "JNI.int"
    JLong    -> "JNI.long"
    JFloat   -> "JNI.float"
    JDouble  -> "JNI.double"
    JObj n   -> "JNI.object \"" ++ n ++ "\""
    -- arrays are treated as objects:
    JArr c   -> "JNIS.object' \"[" ++ printJniRawType c ++ "\""


printJniRawType :: JavaType -> String
printJniRawType t = case t of
    JBoolean -> "Z"
    JChar    -> "C"
    JByte    -> "B"
    JShort   -> "S"
    JInt     -> "I"
    JLong    -> "J"
    JFloat   -> "F"
    JDouble  -> "D"
    JObj n   -> 'L' : tr '.' '/' n ++ ";"
    JArr c   -> '[' : printJniRawType c


printJniSignature :: JavaMethod -> String
printJniSignature method = show name ++ " JNI.::= " ++ args ++ ret
  where
    name = methodName method
    args = concatMap ((++ " --> ") . printJniType) (map fst $ methodParams method)
    ret  = maybe "JNI.void" printJniType (fst $ methodReturnType method)

typeclassFor :: String -> String
typeclassFor name = name ++ "__." ++ tyclTName (takeClassName name)

newtypeFor :: String -> String
newtypeFor name = name ++ "__." ++ newtTName (takeClassName name)

jtypeForClass :: JavaClass -> JavaType
jtypeForClass clazz = JObj (className clazz)

gtypeForClass :: JavaClass -> JavaGenericType
gtypeForClass clazz
    | null (classTypeParams clazz) = NotSoGeneric
    | otherwise = Parameterized {
        jgtBasetype = className clazz,
        jgtParameters = map (TypeVarReference . paramName) (classTypeParams clazz)
      }

constructorType :: JavaClass -> JavaConstructor -> (JavaClass -> String) -> Typomatic ArgInfo
constructorType clazz constr modNameFor = do
    let method = JavaMethod {
            methodName = "", methodName' = "",
            methodParams = constructorParams constr,
            methodReturnType = (Just (jtypeForClass clazz), gtypeForClass clazz),
            methodExceptions = constructorExceptions constr,
            methodTypeParams = constructorTypeParams constr,
            methodStatic = True, methodFinal = True,
            methodAbstract = False, methodNative = False, methodSynchronized = False
          }
    result <- methodType clazz method modNameFor
    let jniSig = fJniSignature result ->> strSplitAll " --> " ->> init ->> strJoin " --> "
    return $ result {
        fJniSignature = drop 11 jniSig
      }

getterType :: JavaClass -> JavaField -> (JavaClass -> String) -> Typomatic ArgInfo
getterType clazz field modNameFor = do
    let isStatic = fieldStatic field
        (jtype, gtype) = fieldType field
        method = JavaMethod {
            methodName = "", methodName' = "",
            methodParams = [],
            methodReturnType = (Just jtype, gtype),
            methodExceptions = [], methodTypeParams = [],
            methodStatic = isStatic, methodFinal = fieldFinal field,
            methodAbstract = False, methodNative = False, methodSynchronized = False
          }
    methodType clazz method modNameFor

setterType :: JavaClass -> JavaField -> (JavaClass -> String) -> Typomatic ArgInfo
setterType clazz field modNameFor = do
    let isStatic = fieldStatic field
        method = JavaMethod {
            methodName = "", methodName' = "",
            methodParams = [fieldType field],
            methodReturnType = (Just (JObj "@()"), NotSoGeneric),
            methodExceptions = [], methodTypeParams = [],
            methodStatic = isStatic, methodFinal = fieldFinal field,
            methodAbstract = False, methodNative = False, methodSynchronized = False
          }
    methodType clazz method modNameFor
    
methodType :: JavaClass -> JavaMethod -> (JavaClass -> String) -> Typomatic ArgInfo
methodType clazz method_ modNameFor = do

    let className = Types.className clazz
        classParams = classTypeParams clazz

        -- sanitize tyVars tied to the method by distinguishing
        -- from tyVars tied to the class by adding an apostrophe
        -- to the name of the tyVar if it is tied to the method.
        methodTypeVars = map paramName (methodTypeParams method_)
        safe var@(TyVar str)
            | var `elem` methodTypeVars = TyVar (str ++ "'")
            | otherwise = var
        method = everywhere (mkT safe) method_

        -- the following three functions create the list of
        -- argument parameters, including @this@ (if the method
        -- is not static) and the return type. The return type
        -- is separated later on again.
        thisParam = if null classParams
            then NotSoGeneric
            else Parameterized {
                    jgtBasetype = className,
                    jgtParameters = (map (TypeVarReference . paramName) classParams)
                  }
        params = (if methodStatic method then [] else [thisParam])
            ++ map snd (methodParams method)
            ++ [maybe (TypeVarReference (TyVar "()"))
                      (const $ snd $ methodReturnType method)
                      (fst $ methodReturnType method)]
        jtypes = (if methodStatic method then [] else [JObj className])
            ++ map fst (methodParams method)

        -- turns a JavaGenericType definition into type variables.
        -- The names are taken from the monad via 'newVar'.
        tvar param = case jgtType param of
            WildcardT -> do
                name <- newVar
                return $ TVar name
            ParameterizedT -> do
                name <- newVar
                params <- mapM tvar (jgtParameters param)
                return $ TVars $ TVar name : params
            GenericArrayT -> do
                name <- newVar
                return $ TVar name
            TypeVarReferenceT -> do
                let name = (tyVarName (jgtName param))
                return $ TVar name
            NotSoGenericT -> do
                name <- newVar
                return $ TVar name

        -- creates a haskell signature (-> String)
        signature typeVars returnVar = do
            -- retrieve the context and turn each variable into a String.
            contexts <- getContext >>= mapM (\(tvar, context) -> do
                return $ context ++ " " ++ printTVar tvar)

            let argTypes = concatMap ((++ " -> ") . printTVar) typeVars
                -- the final type is wrapped in the Java monad
                returnType = "JNI.Java " ++ printTVar returnVar
                -- finally assemble the conetext.
                context = if null contexts then "" else
                    "(" ++ concat (List.intersperse ", " contexts) ++ ") => "
            
            -- return the full signature, consisting of the context,
            -- the type of the arguments, and the return type.
            return $ concat [context, argTypes, returnType]

        -- create a java signature (-> String)
        javaSignature = do
            let name = methodName method
                args = map printJavaType $ map fst $ methodParams method
            return $ name ++ "(" ++ concat (List.intersperse ", " args) ++ ")"
            

    -- if this is not a static method the first argument
    -- is /this/. This merely pushed the name into the list
    -- of type variable names in the monad.
    when (not $ methodStatic method) (pushVar "this")

    -- get type variables for all arguments, including the
    -- return type (as the return type may be the same as
    -- one of the argument types).
    --
    -- Split the result into arguments and return var again,
    -- since the return variable will get special treatment
    -- henceforth.
    (typeVars, returnVar_) <- breakLast <$> mapM tvar params

    -- augment the return type variable, i.e. if it is not
    -- a type variable at all, replace the variable name by
    -- a constant reference to a specific type.
    --
    -- This is only the case with parameterized type variables
    -- and not-so-generic ones.
    returnVar <- case fst (methodReturnType method) of
        Just (JObj typeName) -> case jgtType (snd $ methodReturnType method) of
            ParameterizedT -> do
                clazz <- getClass typeName
                let (TVars (TVar _ : ts)) = returnVar_
                return $ TVars (TVar (newtypeFor (modNameFor clazz)) : ts)
            NotSoGenericT -> do
                clazz <- getClass typeName
                return $ TVar (newtypeFor (modNameFor clazz))
            _ -> return returnVar_
        _ -> return returnVar_

    -- Add contexts for all arguments and augment types
    let makeContext typeVar jtype = case jtype of
            JObj ('@':tv) -> return (TVar tv)
            JObj name -> do
                clazz <- getClass name
                addContext typeVar $ typeclassFor $ modNameFor clazz
                return typeVar
            JArr componentType -> do
                addContext typeVar "JNIS.Array"
                return typeVar
            JBoolean -> do
                addContext typeVar "JNIS.JBoolean"
                return typeVar
            JChar -> do
                addContext typeVar "JNIS.JChar"
                return typeVar
            JByte -> do
                addContext typeVar "JNIS.JByte"
                return typeVar
            JShort -> do
                addContext typeVar "JNIS.JShort"
                return typeVar
            JInt -> do
                addContext typeVar "JNIS.JInt"
                return typeVar
            JLong -> do
                addContext typeVar "JNIS.JLong"
                return typeVar
            JFloat -> do
                addContext typeVar "JNIS.JFloat"
                return typeVar
            JDouble -> do
                addContext typeVar "JNIS.JDouble"
                return typeVar

    -- Here makeContexts is applied (see above). In the same pass a new
    -- set of typeVars (typeVars') is generated, since makeContexts
    -- migth further investigate array and create type variables for
    -- their component types.
    typeVars' <- mapM (uncurry makeContext) (zip typeVars jtypes)

    -- Create contexts for the ultimate return type.
    returnVar' <- case fst (methodReturnType method) of
            Nothing -> do
                let tvar = TVar "void"
                addContext tvar "JNIS.VoidResult"
                return tvar
            Just t -> case t of
                JObj ('@':tv) -> return (TVar tv)
                JObj _   -> do
                    let tvar = TVars [TVar "object", returnVar]
                    addContext tvar "JNIS.ObjectResult"
                    return tvar
                JArr _   -> do
                    let tvar = TVar "JNI.JObject"
--                    addContext tvar "JNIS.ArrayResult"
                    return tvar
                JBoolean -> do
                    let tvar = TVar "boolean"
                    addContext tvar "JNIS.BooleanResult"
                    return tvar
                JChar    -> do
                    let tvar = TVar "char"
                    addContext tvar "JNIS.CharResult"
                    return tvar
                JByte    -> do
                    let tvar = TVar "byte"
                    addContext tvar "JNIS.ByteResult"
                    return tvar
                JShort   -> do
                    let tvar = TVar "short"
                    addContext tvar "JNIS.ShortResult"
                    return tvar
                JInt     -> do
                    let tvar = TVar "int"
                    addContext tvar "JNIS.IntResult"
                    return tvar
                JLong    -> do
                    let tvar = TVar "long"
                    addContext tvar "JNIS.LongResult"
                    return tvar
                JFloat   -> do
                    let tvar = TVar "float"
                    addContext tvar "JNIS.FloatResult"
                    return tvar
                JDouble  -> do
                    let tvar = TVar "double"
                    addContext tvar "JNIS.DoubleResult"
                    return tvar

    -- generate the Haskell signature (a String)
    sig <- signature typeVars' returnVar'

    -- generate the Java signature (a String).
    -- This is used for documentation purposed later on
    -- (i.e. inserted as haddock docstring).
    jsig <- javaSignature

    -- generate the names of the arguments
    let argNames = (if methodStatic method then id else ("this":) . init)
            $ zipWith (\_ i -> 'a' : show i) typeVars [(1 :: Integer)..]

    -- assemble and return all the calculated information
    return $ ArgInfo {
        fArguments  = typeVars,
        fReturnType = returnVar',
        fArgNames   = argNames,
        fSignature  = sig,
        fJavaSignature = jsig,
        fJavaReturnType = maybe "void" printJavaType (fst $ methodReturnType method),
        fJniSignature = printJniSignature method
      }

--------------------------------------------
-- The following are utilities for the monad
--------------------------------------------

-- | The state of the monad.
data TypomaticState = TypomaticState {
    tVars :: [String],
    tContext :: Set (TVar, String),
    tParams :: [String],
    tClasses :: String -> JavaClass
  }

-- | Retrieve the definition of a class.
--
-- The monad has an internal store of class names
-- and their definitions. See 'tClasses'.
getClass :: String -> Typomatic JavaClass
getClass name = do
    state <- State.get
    return (tClasses state name)

-- | Get the current context as a list.
getContext :: Typomatic [(TVar, String)]
getContext = State.get >>= return . Set.toList . tContext

-- | Add a context for a specific type variable.
addContext :: TVar -> String -> Typomatic ()
addContext tvar string = do
    state <- State.get
    State.put (state {tContext = ((tvar, string) `Set.insert` tContext state)})

-- | Introduce a new name.
--
-- This simply takes the next element in the infinite
-- 'tVars' list and stores the tail back in the monad.
newVar :: Typomatic String
newVar = do
    state <- State.get
    let (v:vs) = tVars state
    State.put (state { tVars = vs})
    return v

-- | Push a new name in the front of the available names.
pushVar :: String -> Typomatic ()
pushVar name = do
    state <- State.get
    State.put (state {tVars = (name : tVars state)})

-- | The monad.
newtype Typomatic a = Typomatic { _runTypomatic :: StateT TypomaticState Identity a }
    deriving (Monad, MonadState TypomaticState, Functor)

-- | Run a computation in the monad.
runTypomatic :: (String -> JavaClass) -> Typomatic a -> a
runTypomatic infoFunc =
    let state = TypomaticState { -- initial state
        tVars = map (('v':) . show) [(1 :: Integer)..],
        tContext = Set.empty,
        tParams = [],
        tClasses = infoFunc
    } in fst . runIdentity . flip runStateT state . _runTypomatic