module Obs.Typing.Conversion

import Control.Monad.Maybe
import Control.Monad.Trans

import Data.Bool
import Data.So

import Decidable.Equality

import Obs.Logging
import Obs.NormalForm
import Obs.NormalForm.Normalise
import Obs.Sort
import Obs.Substitution

%default total

-- Conversion ------------------------------------------------------------------

export
convert : (s : Sort)
  -> (a : NormalForm True ctx)
  -> (t, u : NormalForm (isSet s) ctx)
  -> MaybeT (Logging ann) (NormalForm (isSet s) ctx)

untypedConvert : (t, u : NormalForm True ctx) -> MaybeT (Logging ann) (NormalForm True ctx)
-- Diagonals
untypedConvert (Cnstr (Sort s)) (Cnstr (Sort s')) =
  if s == s' then pure (Cnstr (Sort s)) else nothing
untypedConvert (Cnstr (Pi s s' var a b)) (Cnstr (Pi l l' _ a' b')) = do
  let Yes Refl = decEq s l
    | No _ => nothing
  lift $ trace "going back to type-directed conversion"
  a <- assert_total (convert (suc s) (cast s) a a')
  b <- assert_total (convert (suc s') (cast s') b b')
  pure (Cnstr $ Pi s s' var a b)
untypedConvert (Cnstr (Lambda s var t)) (Cnstr (Lambda s' _ u)) = do
  let Yes Refl = decEq s s'
    | No _ => nothing
  lift $ trace "converting under lambda"
  t <- untypedConvert t u
  pure (Cnstr $ Lambda s var t)
untypedConvert (Cnstr (Sigma s s' var a b)) (Cnstr (Sigma l l' _ a' b')) = do
  let Yes Refl = decEq s l
    | No _ => nothing
  lift $ trace "going back to type-directed conversion"
  a <- assert_total (convert (suc s) (cast s) a a')
  b <- assert_total (convert (suc s') (cast s') b b')
  pure (Cnstr $ Sigma s s' var a b)
untypedConvert (Cnstr (Pair s@(Set _) s'@(Set _) prf t u)) (Cnstr (Pair l l' prf' t' u')) = do
  let Yes Refl = decEq (s, s') (l, l')
    | No _ => nothing
  lift $ trace "converting pair pointwise"
  t <- untypedConvert t t'
  u <- untypedConvert u u'
  pure (Cnstr $ Pair l l' Oh t u)
untypedConvert (Cnstr (Pair s@(Set _) Prop prf t u)) (Cnstr (Pair l Prop prf' t' u')) = do
  let Yes Refl = decEq s l
    | No _ => nothing
  lift $ trace "converting pair pointwise"
  t <- untypedConvert t t'
  pure (Cnstr $ Pair l Prop Oh t Irrel)
untypedConvert (Cnstr (Pair Prop s'@(Set _) prf t u)) (Cnstr (Pair Prop l' prf' t' u')) = do
  let Yes Refl = decEq s' l'
    | No _ => nothing
  lift $ trace "converting pair pointwise"
  u <- untypedConvert u u'
  pure (Cnstr $ Pair Prop l' Oh Irrel u)
untypedConvert (Cnstr Bool) (Cnstr Bool) = pure (Cnstr Bool)
untypedConvert (Cnstr True) (Cnstr True) = pure (Cnstr True)
untypedConvert (Cnstr False) (Cnstr False) = pure (Cnstr False)
untypedConvert (Cnstr Top) (Cnstr Top) = pure (Cnstr Top)
untypedConvert (Cnstr Bottom) (Cnstr Bottom) = pure (Cnstr Bottom)
untypedConvert t@(Ntrl (Var _ _ _ i)) (Ntrl (Var _ _ _ j)) =
  if elemToNat i == elemToNat j then pure t else nothing
untypedConvert lhs@(Ntrl (App b t u)) rhs@(Ntrl (App b' t' u')) = do
  let Yes Refl = decEq b b'
    | No _ => nothing
  lift $ trace "checking parts of a spine"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  let True = b'
    | False => lift $ doApp t Irrel
  u <- untypedConvert u u'
  lift $ doApp t u
untypedConvert lhs@(Ntrl (Fst b t)) rhs@(Ntrl (Fst b' t')) = do
  lift $ trace "checking full pair for fst"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  lift $ doFst True b t
untypedConvert lhs@(Ntrl (Snd b t)) rhs@(Ntrl (Snd b' t')) = do
  lift $ trace "checking full pair for snd"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  lift $ doSnd b True (rewrite orTrueTrue b in t)
untypedConvert lhs@(Ntrl (If t f g)) rhs@(Ntrl (If t' f' g')) = do
  lift $ trace "checking all branches of if"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  f <- untypedConvert f f'
  g <- untypedConvert g g'
  lift $ doIf t f g
untypedConvert (Ntrl Absurd) (Ntrl Absurd) = pure (Ntrl Absurd)
untypedConvert lhs@(Ntrl (Equal a t u)) rhs@(Ntrl (Equal a' t' u')) = do
  lift $ trace "checking equal (stuck on type)"
  a <- untypedConvert (assert_smaller lhs (Ntrl a)) (assert_smaller rhs (Ntrl a'))
  t <- untypedConvert t t'
  u <- untypedConvert u u'
  lift $ doEqual _ a t u
untypedConvert lhs@(Ntrl (EqualL t u)) rhs@(Ntrl (EqualL t' u')) = do
  lift $ trace "checking equal (stuck on left)"
  t <- untypedConvert (assert_smaller lhs (Ntrl t)) (assert_smaller rhs (Ntrl t'))
  u <- untypedConvert u u'
  lift $ doEqualSet t u
untypedConvert lhs@(Ntrl (EqualR t u)) rhs@(Ntrl (EqualR t' u')) = do
  lift $ trace "checking equal (stuck on right)"
  t <- untypedConvert (assert_smaller lhs (Cnstr t)) (assert_smaller rhs (Cnstr t'))
  u <- untypedConvert (assert_smaller lhs (Ntrl u)) (assert_smaller rhs (Ntrl u'))
  lift $ doEqualSet t u
untypedConvert lhs@(Ntrl (EqualStuck t u)) rhs@(Ntrl (EqualStuck t' u')) = do
  lift $ trace "checking equal (head mismatch)"
  t <- untypedConvert (assert_smaller lhs (Cnstr t)) (assert_smaller rhs (Cnstr t'))
  u <- untypedConvert (assert_smaller lhs (Cnstr u)) (assert_smaller rhs (Cnstr u'))
  lift $ doEqualSet t u
untypedConvert lhs@(Ntrl (CastL a b t)) rhs@(Ntrl (CastL a' b' t')) = do
  lift $ trace "checking cast (stuck on left)"
  a <- untypedConvert (assert_smaller lhs (Ntrl a)) (assert_smaller rhs (Ntrl a'))
  b <- untypedConvert b b'
  t <- untypedConvert t t'
  lift $ doCast _ a b t
untypedConvert lhs@(Ntrl (CastR a b t)) rhs@(Ntrl (CastR a' b' t')) = do
  lift $ trace "checking cast (stuck on right)"
  a <- untypedConvert (assert_smaller lhs (Cnstr a)) (assert_smaller rhs (Cnstr a'))
  b <- untypedConvert (assert_smaller lhs (Ntrl b)) (assert_smaller rhs (Ntrl b'))
  t <- untypedConvert t t'
  lift $ doCast _ a b t
untypedConvert lhs@(Ntrl (CastStuck a b t)) rhs@(Ntrl (CastStuck a' b' t')) = do
  lift $ trace "checking cast (stuck on left)"
  a <- untypedConvert (assert_smaller lhs (Cnstr a)) (assert_smaller rhs (Cnstr a'))
  b <- untypedConvert (assert_smaller lhs (Cnstr b)) (assert_smaller rhs (Cnstr b'))
  t <- untypedConvert t t'
  lift $ doCast _ a b t
-- Pi types
untypedConvert t@(Cnstr (Lambda s var _)) (Ntrl u) = do
  u <- lift $ doApp (Ntrl $ weaken [isSet s] u) (point (var, (s ** Refl)) Here)
  assert_total (untypedConvert t (Cnstr (Lambda s var u)))
untypedConvert (Ntrl t) u@(Cnstr (Lambda s var _)) = do
  t <- lift $ doApp (Ntrl $ weaken [isSet s] t) (point (var, (s ** Refl)) Here)
  assert_total (untypedConvert (Cnstr (Lambda s var t)) u)
-- Sigma types
untypedConvert t@(Cnstr (Pair s s' prf _ _)) t'@(Ntrl _) =
  let t'' : NormalForm (isSet s || isSet s') ctx
      t'' = rewrite trans (sym $ maxIsSet s s') (soToEq prf) in t'
  in do
  t' <- lift $ doFst (isSet s) (isSet s') t''
  u' <- lift $ doSnd (isSet s) (isSet s') t''
  assert_total (untypedConvert t (Cnstr (Pair s s' prf t' u')))
untypedConvert t@(Ntrl _) t'@(Cnstr (Pair s s' prf _ _)) =
  let t'' : NormalForm (isSet s || isSet s') ctx
      t'' = rewrite trans (sym $ maxIsSet s s') (soToEq prf) in t
  in do
  t <- lift $ doFst (isSet s) (isSet s') t''
  u <- lift $ doSnd (isSet s) (isSet s') t''
  assert_total (untypedConvert (Cnstr (Pair s s' prf t u)) t')
-- Bools
-- no cases because diagonals complete
-- Fallback
untypedConvert _ _ = nothing

convert Prop a Irrel Irrel = pure Irrel
convert (Set k) (Ntrl a) t u = do
  lift $ trace $ pretty {ann} "cannot do type-directed conversion on" <++> pretty a
  untypedConvert t u
convert (Set k) (Cnstr (Sort s)) t u = do
  lift $ trace $ pretty {ann} "converting in type" <++> pretty s
  untypedConvert t u
convert (Set k) (Cnstr (Pi s s'@(Set _) var a b)) t u = do
  -- s' must be Set, otherwise ty is in Prop
  lift $ trace "converting pi type"
  t' <- lift $ doApp (Sorted.weaken [isSet s] t) (point (var, (s ** Refl)) Here)
  u' <- lift $ doApp (Sorted.weaken [isSet s] u) (point (var, (s ** Refl)) Here)
  t <- convert s' b t' u'
  pure (Cnstr $ Lambda s var t)
convert (Set k) ty@(Cnstr (Sigma s@(Set j) s' var a b)) t u = do
  lift $ trace "converting sigma type"
  t1 <- lift $ doFst True (isSet s') t
  u1 <- lift $ doFst True (isSet s') u
  t1 <- convert s a t1 u1
  t2 <- lift $ doSnd True (isSet s') t
  u2 <- lift $ doSnd True (isSet s') u
  b <- lift $ subst1 t1 b
  t2 <- convert s' (assert_smaller ty b) t2 u2
  pure (Cnstr $ Pair s s' Oh t1 t2)
convert (Set k) ty@(Cnstr (Sigma Prop s'@(Set _) var a b)) t u = do
  lift $ trace "converting sigma type (snd only)"
  t2 <- lift $ doSnd False True t
  u2 <- lift $ doSnd False True u
  b <- lift $ subst1 Irrel b
  t2 <- convert s' (assert_smaller ty b) t2 u2
  pure (Cnstr $ Pair Prop s' Oh Irrel t2)
convert (Set 0) (Cnstr Bool) (Cnstr True) (Cnstr True) = do
  lift $ trace "converting bool (true)"
  pure (Cnstr True)
convert (Set 0) (Cnstr Bool) (Cnstr False) (Cnstr False) = do
  lift $ trace "converting bool (false)"
  pure (Cnstr False)
convert (Set k) (Cnstr Bool) t@(Ntrl _) u@(Ntrl _) = do
  lift $ trace "converting bool (neutral)"
  untypedConvert t u
convert (Set k) (Cnstr a) t u = lift $ inScope "bad type constructor" $ fatal a