{-# LANGUAGE DeriveDataTypeable #-}

module Hoogle.Type.TypeSig where

import Data.Binary.Defer
import Data.List
import Data.Data
import Data.Generics.UniplateOn


---------------------------------------------------------------------
-- DATA TYPE WITH CONTEXT


-- FULL TYPE
data TypeSig = TypeSig Constraint Type
               deriving (Eq,Ord,Data,Typeable)

type Constraint = [Type]


-- CONSTRICTED TYPE
-- first argument is a list of contexts, (Context,Variable)
type TypeContext = [(String,String)]
data TypeSimp = TypeSimp TypeContext Type
                deriving (Eq,Ord,Data,Typeable)

instance Show TypeSimp where
    show (TypeSimp c t) = show $ TypeSig [TApp (TLit a) [TVar b] | (a,b) <- c] t


---------------------------------------------------------------------
-- DATA TYPES

data Type = TApp Type [Type] -- a list of types, first one being the constructor
          | TLit String -- bound variables, Maybe, ":", "(,)", "(,,)" (tuple)
          | TVar String -- unbound variables, "a"
          | TFun [Type]
          deriving (Eq,Ord,Data,Typeable)


tApp :: Type -> [Type] -> Type
tApp t [] = t
tApp t ts = TApp t ts


fromTFun :: Type -> [Type]
fromTFun (TFun x) = x
fromTFun x = [x]


isTLit, isTVar :: Type -> Bool
isTLit TLit{} = True; isTLit _ = False
isTVar TVar{} = True; isTVar _ = False


fromTApp :: Type -> (Type, [Type])
fromTApp (TApp x y) = (x,y)
fromTApp x = (x,[])


isTLitTuple :: String -> Bool
isTLitTuple x = ',' `elem` x


insertTApp, removeTApp :: Type -> Type
insertTApp = transform f
    where
        f (TApp (TApp x []) y) = TApp x y
        f (TApp x y) = TApp x y
        f (TFun x) = TFun x
        f x = TApp x []

removeTApp = transform f
    where
        f (TApp x []) = x
        f x = x


---------------------------------------------------------------------
-- UNIPLATE INSTANCES

instance Uniplate Type where
    uniplate (TApp x xs) = (x:xs, \(x:xs) -> TApp x xs)
    uniplate (TFun xs) = (xs, TFun)
    uniplate x = ([], \[] -> x)

onTypeSig :: BiplateType TypeSig Type
onTypeSig (TypeSig xs x) = (x:xs, \(x:xs) -> TypeSig xs x)

transformSig = transformOn onTypeSig
universeSig = universeOn onTypeSig


variables :: Type -> [String]
variables x = [v | TVar v <- universe x]

variablesSig :: TypeSig -> [String]
variablesSig x = [v | TVar v <- universeSig x]


---------------------------------------------------------------------
-- BINARYDEFER INSTANCES

instance BinaryDefer TypeSig where
    put (TypeSig a b) = put2 a b
    get = get2 TypeSig

instance BinaryDefer Type where
    put (TApp a b) = putByte 0 >> put2 a b
    put (TLit a)   = putByte 1 >> put1 a
    put (TVar a)   = putByte 2 >> put1 a
    put (TFun a)   = putByte 3 >> put1 a

    get = do
        i <- getByte
        case i of
            0 -> get2 TApp
            1 -> get1 TLit
            2 -> get1 TVar
            3 -> get1 TFun


---------------------------------------------------------------------
-- SHOW INSTANCES

showConstraint :: Constraint -> String
showConstraint [] = ""
showConstraint [x] = show x ++ " => "
showConstraint xs = "(" ++ intercalate ", " (map show xs) ++ ") => "


-- TODO: show (TLit ":+:") should be "(:+:)"
instance Show Type where
    showsPrec i x = showString $ f i x
        where
            -- Show lists and tuples specially
            f i (TApp (TLit "[]") [x]) = "[" ++ show x ++ "]"
            f i (TApp (TLit ('(':tup)) xs)
                | not (null tup) && last tup == ')' && all (== ',') (init tup) && length tup == length xs
                = b True $ intercalate ", " $ map show xs

            -- Should parallel lists and unboxed tuples specially
            f i (TApp (TLit "[::]") [x]) = "[:" ++ show x ++ ":]"
            f i (TApp (TLit ('(':'#':tup)) xs)
                | "#)" `isSuffixOf` tup && all (== ',') (drop 2 $ reverse tup) && length tup - 1 == length xs
                = "(# " ++ intercalate ", " (map show xs) ++ " #)"

            
            f i (TLit x) = x
            f i (TVar x) = x
            
            f i (TApp x xs) = b (i > 1) $ unwords $ map (f 2) (x:xs)
            f i (TFun xs)   = b (i > 0) $ intercalate " -> " $ map (f 1) xs
            
            b True x = "(" ++ x ++ ")"
            b False x = x


instance Show TypeSig where
    show (TypeSig x xs) = showConstraint x ++ show xs


-- shows an element within a function
-- to get brackets right after splitFun
showFun :: Type -> String
showFun x = showsPrec 1 x ""


---------------------------------------------------------------------
-- OPERATIONS

normaliseTypeSig :: TypeSig -> TypeSig
normaliseTypeSig = transformOn onTypeSig normaliseType


normaliseType :: Type -> Type
normaliseType = transform f
    where
        f (TApp x []) = x
        f (TApp (TLit "->") xs) = f $ TFun xs
        f (TFun [x]) = x
        f (TFun xs) = TFun $ g xs
        f x = x

        g [] = []
        g [TFun xs] = g xs
        g (x:xs) = x : g xs


splitFun :: Type -> [Type]
splitFun (TFun xs) = xs
splitFun x = [x]


renameVars :: (String -> String) -> TypeSig -> TypeSig
renameVars f = transformOn onTypeSig g
    where
        g (TVar x) = TVar $ f x
        g x = x