data ExprType
  = IntType
  | BoolType
  | StringType
  | PairType ExprType ExprType

repr : ExprType -> Type
repr IntType = Int
repr BoolType = Bool
repr StringType = String
repr (PairType t1 t2) = Pair (repr t1) (repr t2)

decEq : (a : ExprType) -> (b : ExprType) -> Maybe (a = b)
decEq IntType IntType = Just Refl
decEq BoolType BoolType = Just Refl
decEq StringType StringType = Just Refl
decEq (PairType lt1 lt2) (PairType rt1 rt2) = do
  subEq1 <- decEq lt1 rt1
  subEq2 <- decEq lt2 rt2
  let firstEqual = replace {P = \t1 => PairType lt1 lt2 = PairType t1 lt2} subEq1 Refl
  let secondEqual = replace {P = \t2 => PairType lt1 lt2 = PairType rt1 t2} subEq2 firstEqual
  pure secondEqual
decEq _ _ = Nothing

data Op
  = Add
  | Subtract
  | Multiply
  | Divide

data Expr
  = IntLit Int
  | BoolLit Bool
  | StringLit String
  | BinOp Op Expr Expr
  | IfElse Expr Expr Expr
  | Pair Expr Expr
  | Fst Expr
  | Snd Expr

data SafeExpr : ExprType -> Type where
  IntLiteral : Int -> SafeExpr IntType
  BoolLiteral : Bool -> SafeExpr BoolType
  StringLiteral : String -> SafeExpr StringType
  BinOperation : (repr a -> repr b -> repr c) -> SafeExpr a -> SafeExpr b -> SafeExpr c
  IfThenElse : SafeExpr BoolType -> SafeExpr t -> SafeExpr t -> SafeExpr t
  NewPair : SafeExpr t1 -> SafeExpr t2 -> SafeExpr (PairType t1 t2)
  First : SafeExpr (PairType t1 t2) -> SafeExpr t1
  Second : SafeExpr (PairType t1 t2) -> SafeExpr t2

typecheckOp : Op -> (a : ExprType) -> (b : ExprType) -> Either String (c : ExprType ** repr a -> repr b -> repr c) 
typecheckOp Add IntType IntType = Right (IntType ** (+))
typecheckOp Subtract IntType IntType = Right (IntType ** (-))
typecheckOp Multiply IntType IntType = Right (IntType ** (*))
typecheckOp Divide IntType IntType = Right (IntType ** div)
typecheckOp _ _ _ = Left "Invalid binary operator application"

requireBool : (n : ExprType ** SafeExpr n) -> Either String (SafeExpr BoolType)
requireBool (BoolType ** e) = Right e
requireBool _ = Left "Not a boolean."

typecheck : Expr -> Either String (n : ExprType ** SafeExpr n)
typecheck (IntLit i) = Right (_ ** IntLiteral i)
typecheck (BoolLit b) = Right (_ ** BoolLiteral b)
typecheck (StringLit s) = Right (_ ** StringLiteral s)
typecheck (BinOp o l r) = do
  (lt ** le) <- typecheck l
  (rt ** re) <- typecheck r
  (ot ** f) <- typecheckOp o lt rt
  pure (_ ** BinOperation f le re)
typecheck (IfElse c t e) =
  do
    ce <- typecheck c >>= requireBool
    (tt ** te) <- typecheck t
    (et ** ee) <- typecheck e
    case decEq tt et of
      Just p => pure (_ ** IfThenElse ce (replace p te) ee)
      Nothing => Left "Incompatible branch types."
typecheck (Pair l r) =
  do
    (lt ** le) <- typecheck l
    (rt ** re) <- typecheck r
    pure (_ ** NewPair le re)
typecheck (Fst p) =
  do
    (pt ** pe) <- typecheck p
    case pt of
      PairType _ _ => pure $ (_ ** First pe)
      _ => Left "Applying fst to non-pair."
typecheck (Snd p) =
  do
    (pt ** pe) <- typecheck p
    case pt of
      PairType _ _ => pure $ (_ ** Second pe)
      _ => Left "Applying snd to non-pair."

eval : SafeExpr t -> repr t
eval (IntLiteral i) = i
eval (BoolLiteral b) = b
eval (StringLiteral s) = s
eval (BinOperation f l r) = f (eval l) (eval r)
eval (IfThenElse c t e) = if (eval c) then (eval t) else (eval e)
eval (NewPair l r) = (eval l, eval r)
eval (First p) = fst (eval p)
eval (Second p) = snd (eval p)

resultStr : {t : ExprType} -> repr t -> String
resultStr {t=IntType} i = show i
resultStr {t=BoolType} b = show b
resultStr {t=StringType} s = show s
resultStr {t=PairType t1 t2} (l,r) = "(" ++ resultStr l ++ ", " ++ resultStr r ++ ")"

tryEval : Expr -> String
tryEval ex =
  case typecheck ex of
    Left err => "Type error: " ++ err
    Right (t ** e) => resultStr $ eval {t} e

main : IO ()
main = putStrLn $ tryEval $ BinOp Add (Fst (IfElse (BoolLit True) (Pair (IntLit 6) (BoolLit True)) (Pair (IntLit 7) (BoolLit False)))) (BinOp Multiply (IntLit 160) (IntLit 2))