{-
Copyright (c) 2011, Justin Pombrio All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.

Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-}

import Data.Char
import Data.List
import Control.Monad
import Control.Monad.State
import System.IO
import qualified Data.Set as Set
import qualified Data.Map as Map

import SExpr
import Debug.Trace

{-
Usage:
  Take a CPSA-style protocol definition, and compile it into a 
  geometric logic theory.
  > cat prot.scm | ./compiler > prot.glt

Overview:
  The translation pipeline is as follows:

  String  ---->  SExpr  ---->  Protocol  ---->  Theory  ---->  String
          read         compile         translate     show, format

  If you want to change the output formatting (to make use of, say,
  another implementation of the chase), see the Formatting code. The
  chase implementation should be able to handle variable names with
  whatever characters are used in variables in the protocol definition.

Limitations:
  Needs a formula "regular(n) => init_1(n) | ...".
  Interspersed comments cause parse error.
  Syntax errors will be shown as incomplete patterns.
  Protocol algorithm ignored.
  Do not support labels.
  Do not support node pairs.

Terminology:
    
  init_0(n) & init_x(n, x) -> send(n, pair(x, x))

  'init'         principal (princ)
  '0'            index
  'n'            node
  'x'            parameter (param)
  'pair(x, x)'   mesg
  'init_0'       node relation
  'init_0(n)'    node atom
  'init_x'       parameter relation
  'init_x(n, x)' parameter atom
  'send'         event relation
  'send(n, x)'   send atom
-}


{- Protocol -}

type Var = String

type Sort = String

data Mesg = Simple Var
          | Compound String [Mesg]

type Direction = String

data Event = Event Direction Mesg

type Trace = [Event]

type Predicate = String

type HTMesg = (Int, Mesg)

data Assumption = Assumption Predicate HTMesg

data SkelAssum = SkelAssum Predicate Mesg

data Precedes = Precedes (Int, Int) (Int, Int)

data Role = Role String (Map.Map Var Sort) Trace [Assumption]

type Node = (Princ, Int)

type Maplet = (Var, Mesg)

data Strand = Strand Node [Maplet]
            | Listener Mesg

data Form = Protocol String [Role]
          | Skeleton (Map.Map Var Sort) [Strand] [SkelAssum] [Precedes]


{- Main -}

readSExprs :: Handle -> IO [SExpr Pos]
readSExprs handle = do
  pos <- posHandle "" handle
  loop pos
  where
    loop pos = do
      sexpr <- load pos
      case sexpr of
        Nothing -> return []
        Just expr -> do
          exprs <- loop pos
          return (expr : exprs)

testShowTheory = putStrLn $ show $ Theory "theory"
                 [Rule
                  [Relation "R" [term_x]]
                  (Exists ["y", "z"]
                   [Relation "P" [term_x,
                                  Function "f" [term_x,
                                                Variable "y"]],
                    Equation
                    (Function "g" [term_x])
                    (Function "g" [Variable "z"])])]


main = do
  exprs <- readSExprs stdin
  putStrLn $ show $ translate "" $ compile exprs


{- Compilation -}

compile :: [SExpr a] -> [Form]
compile [] = [] -- 'sequence' ought just ignore Nothings...
compile (form : forms) = case compileForm form of
  Nothing -> compile forms
  Just result -> result : compile forms

compileForm :: SExpr a -> Maybe Form
compileForm expr @ (L _ (S _ head : body)) = case head of
  "defprotocol" -> Just $ compileProtocol expr
  "defskeleton" -> Just $ compileSkeleton expr
  "herald" -> Nothing
  "comment" -> Nothing

compileProtocol :: SExpr a -> Form
compileProtocol (L _ (S _ "defprotocol" : S _ id : _ : rest)) =
  let (roles, alist) = partition isRole rest in
  Protocol id (map compileRole roles)
    where
      isRole (L _ (S _ "defrole" : _)) = True
      isRole _ = False

compileRole :: SExpr a -> Role
compileRole (L _ (S _ "defrole" : S _ id : vars : trace : assumptions)) =
  Role id (compileVars vars) (compileTrace trace)
  (concatMap compileAssumption assumptions)

compileVars :: SExpr a -> Map.Map Var Sort
compileVars (L _ (S _ "vars" : vars)) =
  Map.unions $ map compileDecl vars

compileDecl :: SExpr a -> Map.Map Var Sort
compileDecl (L _ decl) =
  let sort = case last decl of S _ s -> s
      vars = map compileVar (init decl) in
  Map.fromList (zip vars (repeat sort))

compileVar :: SExpr a -> Var
compileVar (S _ var) = var

compileTrace :: SExpr a -> Trace
compileTrace (L _ (S _ "trace" : events)) =
  map compileEvent events

compileEvent :: SExpr a -> Event
compileEvent (L _ [S _ direction, mesg]) =
  Event direction (compileMesg mesg)

compileAssumption :: SExpr a -> [Assumption]
compileAssumption (L _ (S _ "comment" : _)) = []
compileAssumption (L _ (S _ pred : mesgs)) =
  map (Assumption pred . compileHTMesg) mesgs

compileHTMesg :: SExpr a -> HTMesg
compileHTMesg (L _ [N _ index, mesg]) =
  (index, compileMesg mesg)
compileHTMesg expr = (1, compileMesg expr)

compileMesg :: SExpr a -> Mesg
compileMesg (S _ id) = Simple id
compileMesg (L _ (S _ head : body)) =
  Compound head (map compileMesg body)
compileMesg (N _ x) = error $ show x
compileMesg (Q _ x) = error x

compileSkeleton :: SExpr a -> Form
compileSkeleton (L _ (S _ "defskeleton" : S _ id : vars : rest)) =
  let (strands, alist) = partition isStrand rest
      (precs, assums) = partition isPrec alist in
  Skeleton (compileVars vars) (map compileStrand strands)
  (concatMap compileSkelAssum assums)
  (concatMap compilePrec precs)
  where
    isStrand (L _ (S _ "defstrand" : _)) = True
    isStrand _ = False
    isPrec (L _ (S _ "precedes" : _)) = True
    isPrec _ = False

compileSkelAssum :: SExpr a -> [SkelAssum]
compileSkelAssum (L _ (S _ "comment" : _)) = []
compileSkelAssum (L _ (S _ pred : mesgs)) =
  map (\(p, m) -> SkelAssum p m) $
  zip (repeat pred) (map compileMesg mesgs)

compilePrec :: SExpr a -> [Precedes]
compilePrec (L _ (S _ "precedes" : precs)) =
  map compileNodePair precs

compileNodePair :: SExpr a -> Precedes
compileNodePair (L _ [L _ [N _ x1, N _ x2],
                      L _ [N _ y1, N _ y2]]) =
  Precedes (x1, x2) (y1, y2)

compileStrand :: SExpr a -> Strand
compileStrand (L _ [S _ "deflistener", mesg]) =
  Listener $ compileMesg mesg
compileStrand (L _ (S _ "defstrand" : S _ id : N _ size : maplets)) =
  Strand (id, size) (map compileMaplet maplets)

compileMaplet :: SExpr a -> Maplet
compileMaplet (L _ [left, right]) =
  (compileVar left, compileMesg right)


{- Theory -}

type Rel = String -- relation symbol
type Fun = String
type Princ = String
type Param = String

data Term = Variable Var
          | Function Fun [Term]

data Atom = Relation Rel [Term]
          | Equation Term Term
atom rel vars = Relation rel (map Variable vars)

type Conjunction = [Atom]

data Existential = Exists [Var] Conjunction

-- Assuming that disjunctions aren't necessary.
data Rule = Rule Conjunction Existential

data Theory = Theory String [Rule]


freeVars :: Term -> [Var]
freeVars (Variable var) = [var]
freeVars (Function _ subterms) = nub $ concatMap freeVars subterms


{- Translation -}

translate :: String -> [Form] -> Theory
translate name forms =
  Theory name $ concatMap translateForm forms

translateForm :: Form -> [Rule]
translateForm protocol@(Protocol _ _) =
  translateProtocol protocol
translateForm skeleton@(Skeleton _ _ _ _) =
  [translateSkeleton skeleton]


{- Protocol Translation -}

translateProtocol :: Form -> [Rule]
translateProtocol (Protocol names roles) =
  concatMap translateRole roles

translateRole :: Role -> [Rule]
translateRole (Role name vars trace assumptions) =
  let varRules = Map.elems $ Map.mapWithKey (translateVar name) vars
      traceRules = zipWith (translateEvent name) trace [1..]
      assumRules = map (translateAssumption name) assumptions in
  concat (varRules ++ traceRules ++ assumRules)

translateVar :: String -> Var -> Sort -> [Rule]
translateVar princ var sort =
  [paramSortRule princ var sort,
   paramFuncRule princ var,
   paramParentRule princ var]

translateEvent :: String -> Event -> Int -> [Rule]
translateEvent princ event index =
  let rule1 = eventRule (princ, index) event
      rule2 = eventParentRule (princ, index)
      rule3 = eventParamRule (princ, index) event in
  case rule2 of
    Nothing -> [rule1, rule3]
    Just rule -> [rule1, rule, rule3]

translateAssumption :: String -> Assumption -> [Rule]
translateAssumption princ (Assumption pred (index, mesg)) =
  [assumptionRule (princ, index) pred (translateMesg id mesg)]

translateMesg :: (Var -> String) -> Mesg -> Term
translateMesg namer (Simple var) =
  Variable (namer var)
translateMesg namer (Compound "pubk" [Simple var]) =
  Function "Pubk" [Variable $ namer var]
translateMesg namer (Compound "privk" [Simple var]) =
  Function "Privk" [Variable $ namer var]
translateMesg namer (Compound "invk" [Simple var]) =
  Function "Invk" [Variable $ namer var]
translateMesg namer (Compound "ltk" [Simple var1, Simple var2]) =
  Function "Ltk" [Variable $ namer var1, Variable $ namer var2]
translateMesg namer (Compound "cat" [x, y]) =
  Function "Pair" (map (translateMesg namer) [x, y])
-- Ensure that Cats and Encs have exactly two subterms.
translateMesg namer (Compound "cat" (x : y : ys)) =
  Function "Pair"
  [translateMesg namer x,
   translateMesg namer (Compound "cat" (y : ys))]
translateMesg namer (Compound "enc" [plaintext, key]) =
  Function "Enc" (map (translateMesg namer) [plaintext, key])
translateMesg namer (Compound "enc" subterms) | length subterms >= 3 =
  Function "Enc" [translateMesg namer (Compound "cat" (init subterms)),
                  translateMesg namer (last subterms)]


{- Rules -}

-- "init1(n) & b_init(n, b) -> non-orig(Privk<b>)"
assumptionRule :: Node -> Predicate -> Term -> Rule
assumptionRule (princ, index) pred term =
  let vars = freeVars term in
  Rule (nodeAtom (princ, index) term_n :
        paramConj princ term_n (zip vars (map Variable vars)))
  (Exists [] [assumAtom pred term])

-- "init1(n) & a_init(n, a) & n1_init(n, n1) -> send(n, Enc<n1, Pubk<a>>)"
eventRule :: Node -> Event -> Rule
eventRule (princ, index) (Event dir mesg) =
  let term = translateMesg id mesg
      params = freeVars term in
  Rule (nodeAtom (princ, index) term_n :
        paramConj princ term_n (zip params (map Variable params)))
  (Exists [] [eventAtom dir term_n term])

-- "init2(n) -> Exists m: Parent(m, n) & init1(m)"
eventParentRule :: Node -> Maybe Rule
eventParentRule (princ, 1) = Nothing
eventParentRule (princ, index) = Just $
  Rule [nodeAtom (princ, index) term_n]
  (Exists [var_m] [parentAtom term_m term_n,
                   nodeAtom (princ, index - 1) term_m])

-- "init1(n) -> Exists a, n1: a_init(n, a) & n1_init(n, n1)"
eventParamRule :: Node -> Event -> Rule
eventParamRule (princ, index) (Event pred mesg) =
  let vars = freeVars (translateMesg id mesg) in
  Rule [nodeAtom (princ, index) term_n]
  (Exists vars (paramConj princ term_n
                (zip vars (map Variable vars))))

-- "a_init(n, x) -> Node(n) & name(x)"
paramSortRule :: Princ -> Param -> Sort -> Rule
paramSortRule princ param sort =
  Rule [paramAtom princ param [term_n, term_x]]
  (Exists [] [sortAtom "Node" term_n,
              sortAtom sort term_x])

-- "a_init(n, x) & a_init(n, y) -> x = y"
paramFuncRule :: Princ -> Param -> Rule
paramFuncRule princ param =
  Rule [paramAtom princ param [term_n, term_x],
        paramAtom princ param [term_n, term_y]]
  (Exists [] [Equation term_x term_y])

-- "a_init(n, x) & Parent(n, m) -> a_init(m, x)"
paramParentRule :: Princ -> Var -> Rule
paramParentRule princ var =
  Rule [paramAtom princ var [term_n, term_x],
        parentAtom term_n term_m]
  (Exists [] [paramAtom princ var [term_m, term_x]])


translateSkeleton :: Form -> Rule
translateSkeleton (Skeleton vars strands assums precs) =
  let nodeVars = take (length strands) var_ns
      precVars = take (length precs) var_ms
      princs = map (\(Strand (princ, _) _) -> princ) strands in
  Rule []
  (Exists (nodeVars ++ Map.keys vars ++
          map fst precVars ++ map snd precVars)
   (sortConj (Map.toList vars) ++
    assumConj assums ++
    concat (zipWith translateStrand strands nodeVars) ++
    concat (zipWith (precConj princs) precs precVars)))

translateStrand :: Strand -> Var -> Conjunction
translateStrand (Strand (princ, index) maplets) var =
  let params = map (\(v, m) -> (v, translateMesg id m)) maplets in
  nodeAtom (princ, index) (Variable var) :
  paramConj princ (Variable var) params


{- Conjunctions -}

-- "b_init(n, b) & n2_init(n, n2) & ..."
paramConj :: Princ -> Term -> [(Param, Term)] -> Conjunction
paramConj princ node params =
  map (\(name, term) -> paramAtom princ name [node, term]) params

-- "Name(a) & Text(x)"
sortConj :: [(Var, Sort)] -> Conjunction
sortConj = map (\(v, s) -> sortAtom s (Variable v))

-- "Nonorig(Privk<b>) & Orig(x)
assumConj :: [SkelAssum] -> Conjunction
assumConj =
  map (\(SkelAssum pred mesg) -> assumAtom pred (translateMesg id mesg))

-- resp2(m0) & init3(m1) & lprec(m0, n2) & lprec(m1, n3) & prec(m0, m1)"
precConj :: [Princ] -> Precedes -> (Var, Var) -> Conjunction
precConj princs (Precedes (s, n) (s', n')) (v, v') =
  let t = Variable v
      t' = Variable v' in
  [nodeAtom (princs !! (s - 1), n) t,
   nodeAtom (princs !! (s' - 1), n') t',
   lprecAtom t (Variable $ var_ns !! (s - 1)),
   lprecAtom t' (Variable $ var_ns !! (s' - 1)),
   precAtom t t']


{- *Atoms* -}

-- "prec(m, n)"
precAtom :: Term -> Term -> Atom
precAtom n1 n2 = Relation "precedes" [n1, n2]

-- "lprec(m, n)"
lprecAtom :: Term -> Term -> Atom
lprecAtom n1 n2 = Relation "lprec" [n1, n2]
  
-- "n1_init(n, x)"
paramAtom :: Princ -> Param -> [Term] -> Atom
paramAtom princ param args =
  Relation (param ++ "_" ++ princ) args

-- "init1(n)"
nodeAtom :: Node -> Term -> Atom
nodeAtom (princ, index) node =
  Relation (princ ++ show index) [node]

-- "Node(n)"
sortAtom :: Sort -> Term -> Atom
sortAtom sort param = Relation sort [param]

-- "Send(n, Enc<n1, Pubk<b>>)"
eventAtom :: Direction -> Term -> Term -> Atom
eventAtom dir node mesg = Relation dir [node, mesg]

-- "Nonorig(Privk<b>)"
assumAtom :: Predicate -> Term -> Atom
assumAtom pred term = Relation pred [term]

-- "Parent(n, m)"
parentAtom :: Term -> Term -> Atom
parentAtom parent child = Relation "Parent" [child, parent]

-- "Listen(Enc<x, y>)"
listenAtom :: Term -> Atom
listenAtom term = Relation "Listen" [term]


{- *Naming* -}

var_n = "_n"
var_m = "_m"
var_x = "_x"
var_y = "_y"
term_n = Variable var_n
term_m = Variable var_m
term_x = Variable var_x
term_y = Variable var_y
var_ns = map (\x -> "_n" ++ show x) [0..]
var_ms = map (\x -> ("_m" ++ show x, "_m" ++ show (x + 1))) [0, 2..]

compileSort :: SExpr a -> Sort
compileSort (S _ sort) = case sort of
  "text" -> "Text"
  "data" -> "Data"
  "name" -> "Name"
  "skey" -> "SKey"
  "akey" -> "AKey"
  "mesg" -> "Basic"


{- *Transcription* -}

instance Show Theory where
  showsPrec _ (Theory name rules) =
    formatTheory $ map (showsPrec 0) rules

instance Show Rule where
  showsPrec _ (Rule [] rhs) =
    formatRule (showString "True") (showsPrec 0 rhs)
  showsPrec _ (Rule lhs rhs) =
    formatRule (formatConj (map (showsPrec 0) lhs)) (showsPrec 0 rhs)

instance Show Existential where
  showsPrec _ (Exists vars conj) =
    formatExistential
    (map showString vars)
    (formatConj (map (showsPrec 0) conj))

instance Show Atom where
  showsPrec _ (Relation rel terms) =
    formatAtom (showString rel) (map (showsPrec 0) terms)
  showsPrec _ (Equation left right) =
    formatEquation (showsPrec 0 left) (showsPrec 0 right)

instance Show Term where
  showsPrec _ (Variable var) = showString var
  showsPrec _ (Function fun subterms) =
    formatTerm (showString fun) (map (showsPrec 0) subterms)


{- *Formatting* -}

-- from Haskell wiki
compose :: [a -> a] -> a -> a
compose fs v = foldl (.) id fs $ v

joinShows :: String -> [ShowS] -> ShowS
joinShows delim = compose . intersperse (showString delim)

formatTheory :: [ShowS] -> ShowS
formatTheory rules = joinShows "" rules

formatRule lhs rhs =
  lhs . showString " -> " . rhs . showString "\n"

formatExistential [] body = body
formatExistential vars body =
  showString "Exists " . joinShows ", " vars . showString ": " . body

formatConj conjuncts = joinShows " & " conjuncts

formatEquation lhs rhs = lhs . showString " = " . rhs

formatAtom rel vars =
  rel . showString "(" . joinShows ", " vars . showString ")"

formatTerm fun subterms =
  fun . showString "<" . joinShows ", " subterms . showString ">"