{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE OverloadedStrings #-}

module ClassFileParser
    ( parseJavaClass
    , parseOnlyJavaClass
    , javaClass
    ) where


import JavaClass

import Data.Array (array, (!))
import Data.Attoparsec.ByteString
import Data.Binary.IEEE754 (wordToFloat, wordToDouble)
import Data.Bits
import Data.ByteString (ByteString, foldl', pack, elem, cons)
import Data.Char (ord)
import Prelude hiding (take, takeWhile, lookup, elem)

type BString = Data.ByteString.ByteString

parseJavaClass :: BString -> Result JavaClass
parseJavaClass = parse $ javaClass <* endOfInput

parseOnlyJavaClass :: BString -> Either String JavaClass
parseOnlyJavaClass = parseOnly $ javaClass <* endOfInput

javaClass :: Parser JavaClass
javaClass = magic >> versions >>= pools
    where
        pools v = cPool >>= jcparser v
        jcparser v cp = JavaClass v cp
                <$> accFlags
                <*> index
                <*> index
                <*> itfcs
                <*> flds cp
                <*> mthds cp
                <*> attrs cp

magic :: Parser BString
magic = string magicVal <?> "not a class file"

versions :: Parser Version
versions = Version <$> u2 <*> u2

-- constant pool

cPool :: Parser ConstantPool
cPool = index >>= elements
    where
        elements n = array (1, n - 1) <$> ctelem n 1

ctelem :: Int -> Int -> Parser [(Int, CpInfo)]
ctelem n m
    | n > m = cpinfo >>= check
    | otherwise = return []
    where
        check x@(Long _) = ([(m, x), (m + 1, None)] ++) <$> ctelem n (m + 2)
        check x@(Double _) = ([(m, x), (m + 1, None)] ++) <$> ctelem n (m + 2)
        check x = ((m ,x) :) <$> ctelem n (m + 1)

cpinfo :: Parser CpInfo
cpinfo = choice
    [ word8 0x7 >> cclass
    , word8 0x9 >> cfieldref
    , word8 0xa >> cmethodref
    , word8 0xb >> cinterfaceMethodref
    , word8 0x8 >> cstring
    , word8 0x3 >> cinteger
    , word8 0x4 >> cfloat
    , word8 0x5 >> clong
    , word8 0x6 >> cdouble
    , word8 0xc >> cnameAndType
    , word8 0x1 >> cutf8
    , word8 0xf >> cmethodHandle
    , word8 0x10 >> cmethodType
    , word8 0x12 >> cinvokeDynamic
    ]

cclass :: Parser CpInfo
cclass = Class <$> index

cfieldref :: Parser CpInfo
cfieldref = Fieldref <$> index <*> index

cmethodref :: Parser CpInfo
cmethodref = Methodref <$> index <*> index

cinterfaceMethodref :: Parser CpInfo
cinterfaceMethodref = InterfaceMethodref <$> index <*> index

cstring :: Parser CpInfo
cstring = String <$> index

cinteger :: Parser CpInfo
cinteger = Integer . convert <$> u4

cfloat :: Parser CpInfo
cfloat = Float . wordToFloat <$> u4

clong :: Parser CpInfo
clong = Long . convert <$> u8

cdouble :: Parser CpInfo
cdouble = Double . wordToDouble <$> u8

cnameAndType :: Parser CpInfo
cnameAndType = NameAndType <$> index <*> index

cutf8 :: Parser CpInfo
cutf8 = index >>= utf8bytes
    where
        utf8bytes n = Utf8 <$> take n

cmethodHandle :: Parser CpInfo
cmethodHandle = MethodHandle <$> u1 <*> index

cmethodType :: Parser CpInfo
cmethodType = MethodType <$> index

cinvokeDynamic :: Parser CpInfo
cinvokeDynamic = InvokeDynamic <$> index <*> index

accFlags :: Parser [AccessFlag]
accFlags = checkMask <$> u2

checkMask :: U2 -> [AccessFlag]
checkMask x = foldr ff [] masks
    where
        ff (a, n) acc
            | x .&. n /= 0 = a : acc
            | otherwise = acc

-- end

itfcs :: Parser [Int]
itfcs = index >>= iidxs
    where
        iidxs n = count n index

flds :: ConstantPool -> Parser [FieldInfo]
flds cp = index >>= flip count (finfo cp)

finfo :: ConstantPool -> Parser FieldInfo
finfo cp = FieldInfo <$> flags <*> index <*> index <*> attrs cp
    where
        flags = checkMask <$> u2

mthds :: ConstantPool -> Parser [MethodInfo]
mthds cp = index >>= flip count (minfo cp)

minfo :: ConstantPool -> Parser MethodInfo
minfo cp = MethodInfo <$> flags <*> index <*> index <*> attrs cp
    where
        flags = checkMask <$> u2

-- attributes

attrs :: ConstantPool -> Parser [AttributeInfo]
attrs cp = index >>= flip count (attrinfo cp)

attrinfo :: ConstantPool -> Parser AttributeInfo
attrinfo cp = index >>= alength
    where
        alength = getConstructors cp

{-
attrinfo = AttributeInfo <$> index <*> ainfo
    where
        ainfo = u4 >>= take . fromIntegral
-}

getConstructors :: ConstantPool -> Int -> Parser AttributeInfo
getConstructors cp n = constructors s
    where
        s = getStr $ cp ! n
            where
                getStr (Utf8 x) = x
                getStr _ = error "not a UTF8"
        constructors "ConstantValue" = aconstvalue
        constructors "Code" = acode cp
        constructors "StackMapTable" = astackmaptable
        constructors "Exceptions" = aexceptions
        constructors "InnerClasses" = ainnerclass
        constructors "EnclosingMethod" = aenclosingmethod
        constructors "Synthetic" = asynthetic
        constructors "Signature" = asignature
        constructors "SourceFile" = asourcefile
        constructors "SourceDebugExtension" = asourcedebugextension
        constructors "LineNumberTable" = alinenumbertable
        constructors "LocalVariableTable" = alocalvariabletable
        constructors "LocalVariableTypeTable" = alocalvariabletypetable
        constructors "Deprecated" = adeprecated
        constructors "RuntimeVisibleAnnotations" = aruntimevisibleannotation
        constructors "RuntimeInvisibleAnnotations" = aruntimeinvisibleannotation
        constructors "RuntimeVisibleParameterAnnotations" = aruntimevisibleparameterannotation
        constructors "RuntimeInvisibleParameterAnnotations" = aruntimeinvisibleparameterannotation
        constructors "RuntimeVisibleTypeAnnotations" = aruntimevisibletypeannotation
        constructors "RuntimeInvisibleTypeAnnotations" = aruntimeinvisibletypeannotation
        constructors "AnnotationDefault" = aannotationdefault
        constructors "BootstrapMethods" = abootstrapmethods
        constructors "MethodParameters" = amethodparameters
        constructors _ = OtherAttribute n <$> (u4 >>= take . fromIntegral)

aconstvalue :: Parser AttributeInfo
aconstvalue = ConstantValue <$> (u4 >> index)

acode :: ConstantPool -> Parser AttributeInfo
acode cp = u4 >> acflds
    where
        acflds = Code
            <$> index
            <*> index
            <*> (u4 >>= acodearray . fromIntegral)
            <*> (index >>= flip count execinfo)
            <*> attrs cp

execinfo :: Parser Exception
execinfo = Exception <$> index <*> index <*> index <*> index

astackmaptable :: Parser AttributeInfo
astackmaptable = u4 >> index >>= asmt
    where
        asmt n = StackMapTable <$> count n astackmapframe

astackmapframe :: Parser StackMapFrameUnion
astackmapframe = u1 >>= checkTag
    where
        checkTag t
            | 0 <= t && t < 64 = return $ SameFrame t
            | 63 < t && t < 128 = SameLocalsOneStackItemFrame t <$> averificationtype
            | t == 247 = SameLocalsOneStackItemFrameExtended t <$> u2 <*> averificationtype
            | 247 < t && t < 251 = ChopFrame t <$> u2
            | t == 251 = SameFrameExtended t <$> u2
            | 251 < t && t < 255 = AppendFrame t <$> u2 <*> count (fromIntegral t - 251) averificationtype
            | t == 255 = FullFrame t <$> u2 <*> locals <*> stack
            | otherwise = error "stack_map_frame: reserved tag"
            where
                locals = index >>= flip count averificationtype
                stack = index >>= flip count averificationtype

averificationtype :: Parser VerificationTypeUnion
averificationtype = u1 >>= checkTag
    where
        checkTag t
            | t == 0 = return TopVariable
            | t == 1 = return IntegerVariable
            | t == 2 = return FloatVariable
            | t == 3 = return DoubleVariable
            | t == 4 = return LongVariable
            | t == 5 = return NullVariable
            | t == 6 = return UninitializedThisVariable
            | t == 7 = ObjectVariable <$> index
            | t == 8 = UninitializedVariable <$> u2
            | otherwise = error "verification_type_info: unused tag"

aexceptions :: Parser AttributeInfo
aexceptions = u4 >> index >>= aexcf
    where
        aexcf n = Exceptions <$> count n index

ainnerclass :: Parser AttributeInfo
ainnerclass = u4 >> index >>= aincs
    where
        aincs n = InnerClass <$> count n f
        f = InClass <$> index <*> index <*> index <*> u2

aenclosingmethod :: Parser AttributeInfo
aenclosingmethod = u4 >> aemethod
    where
        aemethod = EnclosingMethod <$> index <*> index

asynthetic :: Parser AttributeInfo
asynthetic = u4 >> return SyntheticA

asignature :: Parser AttributeInfo
asignature = u4 >> asig
    where
        asig = Signature <$> index

asourcefile :: Parser AttributeInfo
asourcefile = u4 >> asf
    where
        asf = SourceFile <$> index

asourcedebugextension :: Parser AttributeInfo
asourcedebugextension = u4 >>= asde . fromIntegral
    where
        asde n = SourceDebugExtension <$> take n

alinenumbertable :: Parser AttributeInfo
alinenumbertable = u4 >> index >>= alnt
    where
        alnt n = LineNumberTable <$> count n alinfo
        alinfo = LineNumberInfo <$> index <*> index

alocalvariabletable :: Parser AttributeInfo
alocalvariabletable = u4 >> index >>= alvt
    where
        alvt n = LocalVariableTable <$> count n alvinfo
        alvinfo = LocalVariableInfo
            <$> index
            <*> index
            <*> index
            <*> index
            <*> index

alocalvariabletypetable :: Parser AttributeInfo
alocalvariabletypetable = u4 >> index >>= alvtt
    where
        alvtt n = LocalVariableTypeTable <$> count n alvtinfo
        alvtinfo = LocalVariableTypeInfo
            <$> index
            <*> index
            <*> index
            <*> index
            <*> index

adeprecated :: Parser AttributeInfo
adeprecated = u4 >> return Deprecated

ara :: ([AnnotationInfo] -> AttributeInfo) -> Parser AttributeInfo
ara f = u4 >> index >>= arva
    where
        arva n = f <$> count n arvainfo

arvainfo :: Parser AnnotationInfo
arvainfo = AnnotationInfo <$> index <*> elemarray
    where
        elemarray = index >>= flip count aelem
        aelem = (,) <$> index <*> aelemvalue

aelemvalue :: Parser ElementValue
aelemvalue = u1 >>= checkTag
    where
        toU1 = fromIntegral . ord :: (Char -> U1)
        checkTag t
            | t `elem` "BCDFIJSZs" = ElementValue t <$> (ConstValue <$> index)
            | t == toU1 'e' = ElementValue t <$> (EnumConst <$> index <*> index)
            | t == toU1 'c' = ElementValue t <$> (ClassInfo <$> index)
            | t == toU1 '@' = ElementValue t <$> (AnnotationValue <$> arvainfo)
            | otherwise = ElementValue t <$> (ArrayValue <$> elemarray)
            where
                elemarray = index >>= flip count aelemvalue

aruntimevisibleannotation :: Parser AttributeInfo
aruntimevisibleannotation = ara RuntimeVisibleAnnotations

aruntimeinvisibleannotation :: Parser AttributeInfo
aruntimeinvisibleannotation = ara RuntimeInvisibleAnnotations

arpa :: ([[AnnotationInfo]] -> AttributeInfo) -> Parser AttributeInfo
arpa f = u4 >> u1 >>= aparray . fromIntegral
    where
        aparray n = f <$> count n apearray
        apearray = index >>= flip count arvainfo

aruntimevisibleparameterannotation :: Parser AttributeInfo
aruntimevisibleparameterannotation = arpa RuntimeVisibleParameterAnnotations

aruntimeinvisibleparameterannotation :: Parser AttributeInfo
aruntimeinvisibleparameterannotation = arpa RuntimeInvisibleParameterAnnotations

arta :: ([TypeAnnotation] -> AttributeInfo) -> Parser AttributeInfo
arta f = u4 >> index >>= atarray
    where
        atarray n = f <$> count n atypeannotation
        atypeannotation = u1 >>= ata
        ata t = TypeAnnotation t
            <$> checkType t
            <*> tpath
            <*> index
            <*> evalpairs
        checkType t
            | t `elem` pack [0x00, 0x01] = TypeParameter <$> u1
            | t == 0x10 = Supertype <$> u2
            | t `elem` pack [0x11, 0x12] = TypeParameterBound <$> u1 <*> u1
            | t `elem` pack [0x13..0x15] = return Empty
            | t == 0x16 = MethodFormalParameter <$> u1
            | t == 0x17 = Throws <$> u2
            | t `elem` pack [0x40, 0x41] = Localvar <$> localvartable
            | t == 0x42 = Catch <$> u2
            | t `elem` pack [0x43..0x46] = Offset <$> u2
            | t `elem` pack [0x47..0x4b] = TypeArgument <$> u2 <*> u1
            | otherwise = error $ "TypeInfoUnion with tag:" ++ show t
            where
                localvartable = index >>= flip count lvtelem
                lvtelem = (,,) <$> index <*> index <*> index
        tpath = u1 >>= tpc . fromIntegral
        tpc n = TypePath <$> count n tpelem
        tpelem = (,) <$> u1 <*> u1
        evalpairs = index >>= flip count evpelem
        evpelem = (,) <$> index <*> aelemvalue

aruntimevisibletypeannotation :: Parser AttributeInfo
aruntimevisibletypeannotation = arta RuntimeVisibleTypeAnnotations

aruntimeinvisibletypeannotation :: Parser AttributeInfo
aruntimeinvisibletypeannotation = arta RuntimeInvisibleTypeAnnotations

aannotationdefault :: Parser AttributeInfo
aannotationdefault = u4 >> (AnnotationDefault <$> aelemvalue)

abootstrapmethods :: Parser AttributeInfo
abootstrapmethods = u4 >> index >>= bms
    where
        bms n = BootstrapMethods <$> count n bmsarray
        bmsarray = (,) <$> index <*> (index >>= flip count index)

amethodparameters :: Parser AttributeInfo
amethodparameters = u4 >> u1 >>= mps . fromIntegral
    where
        mps n = MethodParameters <$> count n parray
        parray = (,) <$> index <*> u2


-- end

-- code
-- if v >= 51, no jsr, jsr_w

acodearray :: Int -> Parser [Instruction]
acodearray n
    | n == 0 = return []
    | n > 0 = parseInst >>= app
    | otherwise = error "acodearray: internal error"
    where
        app (i, l) = (i :) <$> acodearray l
        parseInst = u1 >>= retval
        retval t = (,) <$> inst <*> return (n - m - 1)
            where
                inst = Instruction t s <$> sptake
                (s, m) = opcodes ! t
                sptake
                    | t == 0xc4 = u1 >>= wideCheck -- wide
                    | otherwise = take m
                wideCheck x
                    | x `elem` pack (0xa9 : [0x15..0x19] ++ [0x36..0x3a]) = cons x <$> take 2 -- <x>load, <x>store, ret
                    | otherwise = cons x <$> take 4

-- end

index :: Parser Int
index = fromIntegral <$> u2

u1 :: Parser U1
u1 = anyWord8

u2 :: Parser U2
u2 = cmb <$> take 2

u4 :: Parser U4
u4 = cmb <$> take 4

u8 :: Parser U8
u8 = cmb <$> take 8

cmb :: (Bits a, Num a) => BString -> a
cmb = foldl' ff zeroBits
    where
        ff acc x = shiftL acc 8 .|. fromIntegral x

convert :: (FiniteBits a, Integral a, Enum a) => a -> Int
convert x
    | testBit x (finiteBitSize x - 1) = negate . fromIntegral . complement . pred $ x
    | otherwise = fromIntegral x
