content: {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE TupleSections #-}
  {-# OPTIONS_GHC -Wno-missing-deriving-strategies #-}
  module CryptolServer.Data.Type
    ( JSONSchema(..)
    , JSONType(..)
    , JSONPType(..)
    ) where
  
  import qualified Data.Aeson as JSON
  import Data.Aeson ((.=), (.:), (.!=), (.:?))
  import Data.Aeson.Types (Parser, typeMismatch)
  import Data.Functor ((<&>))
  import qualified Data.Text as T
  
  import qualified Cryptol.Parser.AST as C
  import Cryptol.Parser.Position (emptyRange)
  import Cryptol.Parser.Selector (ppSelector)
  import Cryptol.Utils.RecordMap (recordFromFields)
  import Cryptol.TypeCheck.PP (NameMap, emptyNameMap, ppWithNames)
  import Cryptol.TypeCheck.Type (Kind(..), PC(..), TC(..), TCon(..),
            TFun(..), TParam(..), Type(..), Schema(..), addTNames, kindOf,
            NominalType(..), NominalTypeDef(..)
            )
  import Cryptol.Utils.Ident (mkIdent)
  import Cryptol.Utils.PP (pp, defaultPPCfg)
  import Cryptol.Utils.RecordMap (canonicalFields)
  
  import qualified Argo.Doc as Doc
  import CryptolServer.AesonCompat
  
  
  newtype JSONSchema = JSONSchema Schema
  
  newtype JSONPType = JSONPType { unJSONPType :: C.Type C.PName }
  
  data JSONType = JSONType NameMap Type
    deriving (Show)
  
  newtype JSONKind = JSONKind Kind
  
  instance JSON.ToJSON JSONSchema where
    toJSON (JSONSchema (Forall vars props ty)) =
      let cfg = defaultPPCfg
          ns = addTNames cfg vars emptyNameMap
      in JSON.object [ "forall" .=
                        [JSON.object
                          [ "name" .= show (ppWithNames ns v)
                          , "kind" .= JSONKind (kindOf v)
                          ]
                        | v@(TParam _uniq _k _flav _info) <- vars
                        ]
                       , "propositions" .= map (JSON.toJSON . JSONType ns) props
                       , "type" .= JSONType ns ty
                       ]
  
  instance JSON.ToJSON JSONKind where
    toJSON (JSONKind k) = convert k
      where
        convert KType = "Type"
        convert KNum = "Num"
        convert KProp = "Prop"
        convert (k1 :-> k2) =
          JSON.object [ "kind" .= T.pack "arrow"
                      , "from" .= convert k1
                      , "to" .= convert k2
                      ]
  
  instance JSON.ToJSON JSONType where
    toJSON (JSONType ns t) = convert t
      where
        convert (TCon (TError {}) _) =
          -- TODO: figure out what to do in this situation
          error "JSON conversion of errors is not yet supported"
        convert (TCon (TC tc) args) =
          JSON.object $
          case (tc, args) of
            (TCNum n, []) ->
              [ "type" .= T.pack "number" , "value" .= n ]
            (TCInf, []) -> ["type" .= T.pack "inf"]
            (TCBit, []) -> ["type" .= T.pack "Bit"]
            (TCInteger, []) -> ["type" .= T.pack "Integer"]
            (TCIntMod, [n]) -> [ "type" .= T.pack "Z"
                               , "modulus" .= JSONType ns n
                               ]
            (TCSeq, [t1, TCon (TC TCBit) []]) ->
              [ "type" .= T.pack "bitvector"
              , "width" .= JSONType ns t1
              ]
            (TCSeq, [t1, t2]) ->
              [ "type" .= T.pack "sequence"
              , "length" .= JSONType ns t1
              , "contents" .= JSONType ns t2
              ]
            (TCFun, [t1, t2]) ->
              [ "type" .= T.pack "function"
              , "domain" .= JSONType ns t1
              , "range" .= JSONType ns t2
              ]
            (TCTuple _ , []) ->
              [ "type" .= T.pack "unit" ]
            (TCTuple _, fs) ->
              [ "type" .= T.pack "tuple"
              , "contents" .= map (JSONType ns) fs
              ]
            (TCRational, []) ->
              [ "type" .= T.pack "Rational"]
            (other, otherArgs) ->
              [ "type" .= T.pack "unknown"
              , "constructor" .= show other
              , "arguments" .= map (JSONType ns) otherArgs
              ]
        convert (TCon (TF tf) args) =
          JSON.object
          [ "type" .= T.pack t'
          , "arguments" .= map (JSONType ns) args
          ]
          where
            t' =
              case tf of
                TCAdd -> "+"
                TCSub -> "-"
                TCMul -> "*"
                TCDiv -> "/"
                TCMod -> "%"
                TCExp -> "^^"
                TCWidth -> "width"
                TCMin -> "min"
                TCMax -> "max"
                TCCeilDiv -> "/^"
                TCCeilMod -> "%^"
                TCLenFromThenTo -> "lengthFromThenTo"
        convert (TCon (PC pc) args) =
          JSON.object $
          case (pc, args) of
            (PEqual, [t1, t2]) ->
              [ "prop" .= T.pack "=="
              , "left" .= JSONType ns t1
              , "right" .= JSONType ns t2
              ]
            (PNeq, [t1, t2]) ->
              [ "prop" .= T.pack "!="
              , "left" .= JSONType ns t1
              , "right" .= JSONType ns t2
              ]
            (PGeq, [t1, t2]) ->
              [ "prop" .= T.pack ">="
              , "greater" .= JSONType ns t1
              , "less" .= JSONType ns t2
              ]
            (PFin, [t']) ->
              [ "prop" .= T.pack "fin"
              , "subject" .= JSONType ns t'
              ]
            (PHas x, [t1, t2]) ->
              [ "prop" .= T.pack "has"
              , "selector" .= show (ppSelector x)
              , "type" .= JSONType ns t1
              , "is"   .= JSONType ns t2
              ]
            (PRing, [t']) ->
              [ "prop" .= T.pack "Ring"
              , "subject" .= JSONType ns t'
              ]
            (PField, [t']) ->
              [ "prop" .= T.pack "Field"
              , "subject" .= JSONType ns t'
              ]
            (PRound, [t']) ->
              [ "prop" .= T.pack "Round"
              , "subject" .= JSONType ns t'
              ]
            (PIntegral, [t']) ->
              [ "prop" .= T.pack "Integral"
              , "subject" .= JSONType ns t'
              ]
            (PCmp, [t']) ->
              [ "prop" .= T.pack "Cmp"
              , "subject" .= JSONType ns t'
              ]
            (PSignedCmp, [t']) ->
              [ "prop" .= T.pack "SignedCmp"
              , "subject" .= JSONType ns t'
              ]
            (PLiteral, [t1, t2]) ->
              [ "prop" .= T.pack "Literal"
              , "size" .= JSONType ns t1
              , "subject" .= JSONType ns t2
              ]
            (PZero, [t']) ->
              [ "prop" .= T.pack "Zero"
              , "subject" .= JSONType ns t'
              ]
            (PLogic, [t']) ->
              [ "prop" .= T.pack "Logic"
              , "subject" .= JSONType ns t'
              ]
            (PTrue, []) ->
              [ "prop" .= T.pack "True"
              ]
            (PAnd, [t1, t2]) ->
              [ "prop" .= T.pack "And"
              , "left" .= JSONType ns t1
              , "right" .= JSONType ns t2
              ]
            (pc', args') ->
              [ "prop" .= T.pack "unknown"
              , "constructor" .= show pc'
              , "arguments" .= map (JSONType ns) args'
              ]
        convert (TVar v) =
          JSON.object
          [ "type" .= T.pack "variable"
          , "kind" .= JSONKind (kindOf v)
          , "name" .= show (ppWithNames ns v)
          ]
        convert (TUser _n _args def) = convert def
        convert (TNominal nt ts) =
          JSON.object [ "type" .= case ntDef nt of
                                    Struct {} -> T.pack "newtype"
                                    Enum {}   -> T.pack "enum"
                                    Abstract  -> T.pack "primitive"
                      , "name" .= show (pp (ntName nt))
                      , "arguments" .= map (JSONType ns) ts
                      ]
        convert (TRec fields) =
          JSON.object
          [ "type" .= T.pack "record"
          , "fields" .=
            JSON.object [ keyFromText (T.pack (show (pp f))) .= JSONType ns t'
                        | (f, t') <- canonicalFields fields
                        ]
          ]
  
  
  instance JSON.FromJSON JSONPType where
    parseJSON val = JSONPType <$> getType val
  
      where
        getType :: JSON.Value -> Parser (C.Type C.PName)
        getType (JSON.Object o) =
              case lookupKM "type" o of
                Just t -> asType t o
                Nothing ->
                  case lookupKM "prop" o of
                    Just p -> asProp p o
                    Nothing -> fail "Expected type or prop key"
        getType other = typeMismatch "object" other
  
        asType "record" = \o -> C.TRecord <$> ((o .: "fields") >>= getFields)
          where
            getFields obj = recordFromFields <$> traverse (\(k, v) -> (mkIdent (keyToText k),) . (emptyRange,) <$> getType v) (toListKM obj)
        asType "variable" = \o -> tUser <$> (name <$> o .: "name") <*> (map unJSONPType <$> (o .:? "arguments" .!= []))
        asType "number" = \o -> C.TNum <$> (o .: "value")
        asType "inf" = const $ pure $ tUser (name "inf") []
        asType "Bit" = const $ pure $ C.TBit
        asType "Integer" = const $ pure $ tUser (name "Integer") []
        asType "Z" = \o -> typeField o "modulus" <&> \m -> tUser (name "Z") [m]
        asType "bitvector" = \o -> typeField o "width" <&> \w -> C.TSeq w C.TBit
        asType "sequence" = binTC C.TSeq "length" "contents"
        asType "function" = binTC C.TFun "domain" "range"
        asType "unit" = const $ pure $ C.TTuple []
        asType "tuple" = \o -> C.TTuple <$> typeListField o "contents"
        asType "Rational" = const $ pure $ tUser (name "Rational") []
        asType "+" = tyFun "+"
        asType "-" = tyFun "-"
        asType "*" = tyFun "*"
        asType "/" = tyFun "/"
        asType "%" = tyFun "%"
        asType "^^" = tyFun "^^"
        asType "width" = tyFun "width"
        asType "min" = tyFun "min"
        asType "max" = tyFun "max"
        asType "/^" = tyFun "/^"
        asType "%^" = tyFun "%^"
        asType "lengthFromThenTo" = tyFun "lengthFromThenTo"
        asType other = const $ fail $ "Didn't understand type tag " <> show other
  
        typeField o fname = (o .: fname) >>= getType
        typeListField o fname = (o .: fname) >>= traverse getType
  
        name = C.UnQual . mkIdent
  
        asProp "==" = binPropF "==" "left" "right"
        asProp "!=" = binPropF "!=" "left" "right"
        asProp ">=" = binPropF ">=" "left" "right"
        asProp "fin" = unaryPropF "fin" "subject"
        asProp "Ring" = unaryPropF "Ring" "subject"
        asProp "Field" = unaryPropF "Field" "subject"
        asProp "Round" = unaryPropF "Round" "subject"
        asProp "Integral" = unaryPropF "Integral" "subject"
        asProp "Cmp" = unaryPropF "Cmp" "subject"
        asProp "SignedCmp" = unaryPropF "SignedCmp" "subject"
        asProp "Literal" = binPropF "Literal" "size" "subject"
        asProp "Zero" = unaryPropF "Zero" "subject"
        asProp "Logic" = unaryPropF "Logic" "subject"
        asProp "True" = const $ pure $ tUser (name "True") []
        asProp "And" = binPropF "And" "left" "right"
        asProp other = const $ fail $ "Didn't understand proposition " ++ show other
  
        binProp prop a b = tUser (name prop) [a, b]
        binPropF prop f1 f2 o = binProp prop <$> typeField o f1 <*> typeField o f2
        unaryProp prop a = tUser (name prop) [a]
        unaryPropF prop f o = unaryProp prop <$> typeField o f
        binTC tc f1 f2 o = tc <$> typeField o f1 <*> typeField o f2
        tyFun tf o = tUser (name tf) <$> typeListField o "arguments"
  
        tUser nm = C.TUser (C.Located emptyRange nm)
  
  instance Doc.Described JSONSchema where
    typeName = "JSON Cryptol Types"
    description =
      [ Doc.Paragraph [Doc.Text "JSON representations of types are type schemas. A type schema has three fields:"]
      , Doc.DescriptionList
          [ (pure $ Doc.Literal "forall",
             Doc.Paragraph [ Doc.Text "Contains an array of objects. Each object has two fields: "
                           , Doc.Literal "name", Doc.Text " is the name of a type variable, and "
                           , Doc.Literal "kind", Doc.Text " is its kind. There are four kind formers: the string "
                           , Doc.Literal "Type", Doc.Text " represents ordinary datatypes, the string "
                           , Doc.Literal "Num", Doc.Text " is the kind of numbers, and "
                           , Doc.Literal "Prop", Doc.Text " is the kind of propositions. "
                           , Doc.Text "Arrow kinds are represented by objects in which the field "
                           , Doc.Literal "kind", Doc.Text " is the string "
                           , Doc.Literal "arrow", Doc.Text ", and the fields "
                           , Doc.Literal "from", Doc.Text " and ", Doc.Literal "to"
                           , Doc.Text " are the kinds on the left and right side of the arrow, respectively."
                           ])
          , (pure $ Doc.Literal "propositions",
             Doc.Paragraph [Doc.Text "A JSON array of the constraints in the type."])
          , (pure $ Doc.Literal "type",
             Doc.Paragraph [ Doc.Text "The type in which the variables from "
                           , Doc.Literal "forall", Doc.Text " are in scope and the constraints in "
                           , Doc.Literal "propositions", Doc.Text " are in effect."
                           ])
          ]
      ]
  