module Obs.Typing

import Control.Monad.Maybe

import Data.Nat

import Decidable.Equality

import Obs.Logging
import Obs.NormalForm
import Obs.NormalForm.Normalise
import Obs.Substitution
import Obs.Term
import Obs.Typing.Context
import Obs.Typing.Conversion
import Obs.Universe

import System

import Text.Bounded
import Text.PrettyPrint.Prettyprinter
import Text.PrettyPrint.Prettyprinter.Render.Terminal

%default total

-- Loggings ----------------------------------------------------------------------

mismatch : (expected, got : Doc ann) -> Logging ann a
mismatch lhs rhs = fatal $
  hang 2 (pretty "expected" <+> line <+> lhs) <+> comma <+> line <+>
  hang 2 (pretty "got" <+> line <+> align rhs)

typeMismatch : Doc ann -> Doc ann -> Logging ann a
typeMismatch lhs rhs = inScope "type mismatch" $ mismatch lhs rhs

universeMismatch : Doc ann -> Doc ann -> Logging ann a
universeMismatch lhs rhs = inScope "universe mismatch" $ mismatch lhs rhs

-- Utilities -------------------------------------------------------------------

MkPair : (s, s' : Universe)
  -> NormalForm (relevance s) ctx
  -> NormalForm (relevance s') ctx
  -> NormalForm (relevance (max s s')) ctx
MkPair Prop Prop t u = u
MkPair Prop s'@(Set _) t u = Cnstr $ Pair Prop s' Set t u
MkPair s@(Set _) s' t u = Cnstr $ Pair s s' Set t u

-- Checking and Inference ------------------------------------------------------

check' : (tyCtx : TyContext ctx)
  -> (s : Universe)
  -> TypeNormalForm ctx
  -> Term (length ctx)
  -> Logging ann (NormalForm (relevance s) ctx)

infer' : (tyCtx : TyContext ctx)
  -> Term (length ctx)
  -> Logging ann (s ** (TypeNormalForm ctx, NormalForm (relevance s) ctx))

inferType' : {default typeMismatch tag : forall a . Doc ann -> Doc ann -> Logging ann a}
  -> (tyCtx : TyContext ctx)
  -> Term (length ctx)
  -> Logging ann (Universe, TypeNormalForm ctx)
inferType' ctx t = do
  (Set _ ** (Cnstr (Universe s), a)) <- infer' ctx t
    | (_ ** (type, a)) => tag (pretty "sort") (pretty type)
  pure (s, a)

check : (tyCtx : TyContext ctx)
  -> (s : Universe)
  -> TypeNormalForm ctx
  -> WithBounds (Term (length ctx))
  -> Logging ann (NormalForm (relevance s) ctx)
check ctx s a (MkBounded t irrel bounds) = inBounds (MkBounded (check' ctx s a t) irrel bounds)

checkType : (tyCtx : TyContext ctx)
  -> (s : Universe)
  -> WithBounds (Term (length ctx))
  -> Logging ann (TypeNormalForm ctx)
checkType ctx s tm = check ctx (succ s) (cast s) tm

infer : (tyCtx : TyContext ctx)
  -> WithBounds (Term (length ctx))
  -> Logging ann (s ** (TypeNormalForm ctx, NormalForm (relevance s) ctx))
infer ctx (MkBounded t irrel bounds) = inBounds (MkBounded (infer' ctx t) irrel bounds)

inferType : {default typeMismatch tag : forall a . Doc ann -> Doc ann -> Logging ann a}
  -> (tyCtx : TyContext ctx)
  -> WithBounds (Term (length ctx))
  -> Logging ann (Universe, TypeNormalForm ctx)
inferType ctx (MkBounded t irrel bounds) = inBounds (MkBounded (inferType' ctx t) irrel bounds)

check' ctx sort (Cnstr (Pi domainSort codomainSort _ domain codomain)) (Lambda {var, body}) = do
  info "encountered lambda"

  let Yes Refl = decEq sort (domainSort ~> codomainSort)
    | No _ => fatal "internal universe mismatch"

  let bodyCtx = ctx ::< MkFreeVar domainSort domain var.val

  trace $ pretty {ann} "checking body has type" <++> pretty codomainSort
  body <- check bodyCtx codomainSort codomain body

  pure $ case codomainSort of
    Prop => Irrel
    Set _ => Cnstr (Lambda domainSort var.val body)

check' ctx sort type (Lambda {var, body}) = typeMismatch (pretty "function") (pretty type)

check' ctx sort (Cnstr (Sigma indexSort elementSort var index element)) (Pair {first, second}) = do
  info "encountered pair"

  let Yes Refl = decEq sort (max indexSort elementSort)
    | No _ => fatal "internal universe mismatch"

  trace $ pretty {ann} "checking first has type" <++> pretty index
  first <- check ctx indexSort index first

  element <- subst1 first element
  trace $ pretty {ann} "checking second has type" <++> pretty element
  second <- check ctx elementSort element second

  pure $ MkPair indexSort elementSort first second

check' ctx sort type (Pair {first, second}) = typeMismatch (pretty "pair type") (pretty type)

check' ctx sort type (Absurd {contradiction}) = do
  info "encountered absurd"

  trace "checking proof of contradiction"
  Irrel <- check ctx Prop (Cnstr Bottom) contradiction

  pure $ doAbsurd _

check' ctx sort type t = do
  info "falling through to inference"

  (inferredSort ** (inferredType, value)) <- infer' ctx t
  trace $ pretty {ann} "inferred type is" <++> pretty inferredType

  let Yes Refl = decEq sort inferredSort
    | No _ => typeMismatch (pretty type) (pretty inferredType)

  expandedType <- subst type (reduce ctx)
  expandedInferredType <- subst inferredType (reduce ctx)

  Just _ <- inScope "convert" $ runMaybeT $ convert (succ sort) (cast sort) expandedType expandedInferredType
  | Nothing => inScope "conversion failed" $ mismatch (pretty type) (pretty inferredType)

  pure value


infer' ctx tm@(Var {var, i}) = do
    info $ pretty {ann} "encountered variable" <++> pretty tm

    let def = index' ctx i
    pure (def.sort ** (def.ty, def.tm))

infer' ctx (Universe {s}) = do
    info $ pretty {ann} "encountered universe" <++> pretty s

    pure (succ (succ s) ** (cast (succ s), cast s))

infer' ctx (Pi {var, domain, codomain}) = do
    info "encountered pi"

    (domainSort, domain) <- inferType ctx domain
    trace $ pretty {ann} "domain type is" <++> pretty domain

    let domainVar = MkFreeVar domainSort domain var.val
    let codomainCtx = ctx ::< domainVar

    (codomainSort, codomain) <- inferType codomainCtx codomain
    trace $ pretty {ann} "codomain type is" <++> pretty codomain

    let piSort = domainVar.sort ~> codomainSort
    let term = Cnstr (Pi domainVar.sort codomainSort var.val domain codomain)
    pure (succ piSort ** (cast piSort, term))

infer' ctx (Lambda {var, body}) = fatal "cannot infer type of lambda"

infer' ctx (App {fun, arg}) = do
    info "encountered application"

    (piSort ** (funType, funValue)) <- infer ctx fun
    trace $ pretty {ann} "function has type" <++> pretty funType

    let Cnstr (Pi domainSort codomainSort _ domain codomain) = funType
      | _ => inBounds $ map (const $ typeMismatch (pretty "function") (pretty funType)) fun

    let Yes Refl = decEq piSort (domainSort ~> codomainSort)
      | No _ => fatal "internal universe mismatch"

    let fun = funValue

    trace $ pretty {ann} "checking argument has type" <++> pretty domain
    arg <- check ctx domainSort domain arg
    trace $ pretty {ann} "argument is well typed"

    returnType <- subst1 arg codomain
    returnValue <- doApp fun arg

    let returnValue = rewrite sym $ functionRelevance domainSort codomainSort in returnValue

    pure (codomainSort ** (returnType, returnValue))

infer' ctx (Sigma {var, index, element}) = do
    info "encountered sigma"

    (indexSort, index) <- inferType ctx index
    trace $ pretty {ann} "index type is" <++> pretty index

    let indexVar = MkFreeVar indexSort index var.val
    let elementCtx = ctx ::< indexVar

    (elementSort, element) <- inferType elementCtx element
    trace $ pretty {ann} "element type is" <++> pretty element

    let sigmaSort = max indexVar.sort elementSort
    let term = Cnstr (Sigma indexVar.sort elementSort var.val index element)
    pure (succ sigmaSort ** (cast sigmaSort, term))

infer' ctx (Pair {first, second}) = fatal "cannot infer type of pair"

infer' ctx (First {arg}) = do
    info "encountered first"

    (sigmaSort ** (argType, argValue)) <- infer ctx arg
    trace $ pretty {ann} "pair has type" <++> pretty argType

    let Cnstr (Sigma indexSort elementSort _ index element) = argType
      | _ => inBounds $ map (const $ typeMismatch (pretty "pair type") (pretty argType)) arg

    let Yes Refl = decEq sigmaSort (max indexSort elementSort)
      | No _ => fatal "internal universe mismatch"

    let argValue = rewrite sym $ pairRelevance indexSort elementSort in argValue

    let returnType = index
    returnValue <- doFst (relevance indexSort) (relevance elementSort) argValue

    pure (indexSort ** (returnType, returnValue))

infer' ctx (Second {arg}) = do
    info "encountered second"

    (sigmaSort ** (argType, argValue)) <- infer ctx arg
    trace $ pretty {ann} "pair has type" <++> pretty argType

    let Cnstr (Sigma indexSort elementSort _ index element) = argType
      | _ => inBounds $ map (const $ typeMismatch (pretty "pair type") (pretty argType)) arg

    let Yes Refl = decEq sigmaSort (max indexSort elementSort)
      | No _ => fatal "internal universe mismatch"

    let argValue = rewrite sym $ pairRelevance indexSort elementSort in argValue

    indexValue <- doFst (relevance indexSort) (relevance elementSort) argValue
    returnType <- subst1 indexValue element
    returnValue <- doSnd (relevance indexSort) (relevance elementSort) argValue

    pure (elementSort ** (returnType, returnValue))

infer' ctx Bool = do
    info "encountered bool"
    pure (Set 1 ** (cast (Set 0), Cnstr Bool))

infer' ctx True = do
    info "encountered true"
    pure (Set 0 ** (Cnstr Bool, Cnstr True))

infer' ctx False = do
    info "encountered false"
    pure (Set 0 ** (Cnstr Bool, Cnstr False))

infer' ctx (If {var, returnType, discriminant, true, false}) = do
    info "encountered if"
    trace "checking discriminant is bool"

    discriminant <- check ctx (Set 0) (Cnstr Bool) discriminant
    trace "discriminant is well typed"

    let returnTypeCtx = ctx ::< MkFreeVar (Set 0) (Cnstr Bool) var.val

    (returnSort, returnType) <- inferType returnTypeCtx returnType
    trace $ pretty {ann} "result is an index of type" <++> pretty returnType

    trueType <- subst1 (Cnstr True) returnType
    trace $ pretty {ann} "checking true branch has type" <++> pretty trueType
    true <- check ctx returnSort trueType true
    trace "true branch is well typed"

    falseType <- subst1 (Cnstr False) returnType
    trace $ pretty {ann} "checking false branch has type" <++> pretty falseType
    false <- check ctx returnSort falseType false
    trace "false branch is well typed"

    returnType <- subst1 discriminant returnType
    returnValue <- doIf discriminant true false

    pure (returnSort ** (returnType, returnValue))

infer' ctx Top = do
    info "encountered top"
    pure (Set 0 ** (Cnstr (Universe Prop), Cnstr Top))

infer' ctx Point = do
    info "encountered point"
    pure (Prop ** (Cnstr Top, Irrel))

infer' ctx Bottom = do
    info "encountered bottom"
    pure (Set 0 ** (Cnstr (Universe Prop), Cnstr Bottom))

infer' ctx (Absurd {contradiction}) = fatal "cannot infer type of absurd"

infer' ctx (Equal {left, right}) = do
    info "encountered equal"

    (sort ** (type, left)) <- infer ctx left
    trace $ pretty {ann} "left side has type" <++> pretty type

    right <- check ctx sort type right
    trace "right side is well typed"

    let equalitySort = Prop
    equalityType <- doEqual (relevance sort) type left right

    pure (succ equalitySort ** (cast equalitySort, equalityType))

infer' ctx (Refl {value}) = do
    info "encountered refl"

    (sort ** (type, value)) <- infer ctx value
    trace $ pretty {ann} "value has type" <++> pretty type

    equalityType <- doEqual (relevance sort) type value value

    pure (Prop ** (equalityType, Irrel))

infer' ctx (Transp {indexedType, oldIndex, value, newIndex, equality}) = do
    info "encountered transp"

    (indexSort ** (indexType, oldIndex)) <- infer ctx oldIndex
    trace $ pretty {ann} "index type is" <++> pretty indexType

    newIndex <- check ctx indexSort indexType newIndex
    trace "new index is well typed"

    let indexedTypeShape = Cnstr $ Pi indexSort (Set 0) "" indexType (Cnstr (Universe Prop))
    indexedType <- check ctx (indexSort ~> Set 0) indexedTypeShape indexedType
    trace $ pretty {ann} "result is an index of type" <++> pretty indexedType

    oldType <- doApp indexedType oldIndex
    trace $ pretty {ann} "checking value has type" <++> pretty oldType
    Irrel <- check ctx Prop oldType value
    trace "value is well typed"

    equalityType <- doEqual (relevance indexSort) indexType oldIndex newIndex
    trace $ pretty {ann} "checking equality has type" <++> pretty equalityType
    equality <- check ctx Prop equalityType equality
    trace "equality is well typed"

    returnType <- doApp indexedType newIndex

    pure (Prop ** (returnType, Irrel))

infer' ctx (Cast {oldType, newType, equality, value}) = do
    info "encountered cast"

    (sort, oldType) <- inferType ctx oldType
    trace $ pretty {ann} "old type is" <++> pretty oldType

    newType <- checkType ctx sort newType
    trace $ pretty {ann} "new type is" <++> pretty newType

    equalityType <- doEqual Relevant (cast sort) oldType newType
    trace $ pretty {ann} "checking equality has type" <++> pretty equalityType
    equality <- check ctx Prop equalityType equality
    trace "equality is well typed"

    trace $ pretty {ann} "checking value has type" <++> pretty oldType
    value <- check ctx sort oldType value
    trace "value is well typed"

    returnValue <- doCast (relevance sort) oldType newType value

    pure (sort ** (newType, returnValue))

infer' ctx (CastId {value}) = do
    info "encountered castRefl"

    (sort ** (type, value)) <- infer ctx value
    trace $ pretty {ann} "value has type" <++> pretty value

    castRefl <- doCast (relevance sort) type type value

    returnType <- doEqual (relevance sort) type castRefl value

    pure (Prop ** (returnType, Irrel))

-- Checking Definitions and Blocks ---------------------------------------------

checkDef : Context ctx -> Term.Definition (length ctx) -> Logging ann (NormalForm.Definition ctx)
checkDef ctx def =
  let block = do
    debug $ "inferring type of \{def.name.val}"
    (sort, ty) <-
      inferType {tag = \lhs, rhs => inScope "invalid declaration" $ mismatch lhs rhs}
        (fromContext ctx) def.ty
    debug $ pretty {ann} "\{def.name.val} has type" <++> pretty ty

    tm <- check (fromContext ctx) sort ty def.tm
    debug $ pretty {ann} "\{def.name.val} is well typed with value" <++> pretty tm

    pure $ MkDefinition {name = def.name, ty, tm, sort}

  in inBounds $ map (const $ inScope "check" block) def.name

export
checkBlock : Block n -> Logging ann (ctx ** (Context ctx, length ctx = n))
checkBlock [] = pure ([] ** ([], Refl))
checkBlock (blk :< def) = do
  (_ ** (ctx, Refl)) <- checkBlock blk
  def <- checkDef ctx def
  pure (_ ** (ctx :< def, Refl))