module Obs.Typing

import Control.Monad.Identity

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

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


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

check' : (tyCtx : Context ctx)
  -> (sort : Universe)
  -> (type : TypeNormalForm ctx)
  -> (term : Term (length ctx))
  -> Logging ann (NormalForm (relevance sort) ctx)

infer' : (tyCtx : Context ctx)
  -> (term : Term (length ctx))
  -> Logging ann (sort : Universe ** (TypeNormalForm ctx, NormalForm (relevance sort) ctx))

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

check : (tyCtx : Context ctx)
  -> (sort : Universe)
  -> (type : TypeNormalForm ctx)
  -> (term : WithBounds (Term (length ctx)))
  -> Logging ann (NormalForm (relevance sort) ctx)
check ctx sort type (MkBounded term irrel bounds) =
  inBounds (MkBounded (check' ctx sort type term) irrel bounds)

checkType : (tyCtx : Context ctx)
  -> (sort : Universe)
  -> (term : WithBounds (Term (length ctx)))
  -> Logging ann (TypeNormalForm ctx)
checkType ctx sort term = check ctx (succ sort) (cast sort) term

infer : (tyCtx : Context ctx)
  -> (term : WithBounds (Term (length ctx)))
  -> Logging ann (sort : Universe ** (TypeNormalForm ctx, NormalForm (relevance sort) ctx))
infer ctx (MkBounded term irrel bounds) = inBounds (MkBounded (infer' ctx term) irrel bounds)

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

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

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

  let bodyCtx = ctx ::< MkParameter {sort = domainSort, ty = domain.type, name = var}

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

  pure $ MkLambda {var = Just var.val, body}

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

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

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

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

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

  pure $ MkPair {first, second}

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

check' ctx sort type (MkContainer {shape, position, next}) = do
  container <- matchContainer type

  let Yes Refl = decEq sort $
        container.inputSort ~>
        max (succ container.shapeSort)
            (container.shapeSort ~>
             max (succ container.positionSort) (container.positionSort ~> container.outputSort))
    | No _ => fatal "internal universe mismatch"

  info "encountered container constructor"

  let shapeType = containerShapeType container
  trace $ pretty {ann} "checking shape has type" <++> pretty shapeType
  shape <- check ctx (container.inputSort ~> succ container.shapeSort) shapeType shape

  shapeInstance <- doApp (weaken [relevance container.inputSort] shape) (point Nothing Here)

  let positionType = containerPositionType container shapeInstance
  trace $ pretty {ann} "checking position has type" <++> pretty positionType
  position <- check
    { tyCtx = ctx
    , sort = container.inputSort ~> container.shapeSort ~> succ container.positionSort
    , type = positionType
    , term = position
    }

  let positionInstance = weaken [relevance container.shapeSort, relevance container.inputSort] position
  positionInstance <- doApp positionInstance (point Nothing (There Here))
  positionInstance <- doApp positionInstance (point Nothing Here)

  let nextType = containerNextType container shapeInstance positionInstance
  trace $ pretty {ann} "checking next has type" <++> pretty nextType
  next <- check
    { tyCtx = ctx
    , sort = container.inputSort ~> container.shapeSort ~> container.positionSort ~> container.outputSort
    , type = nextType
    , term = next
    }

  MkContainer
    { inputRel = relevance container.inputSort
    , shapeRel = relevance container.shapeSort
    , shape
    , position
    , next
    }

check' ctx (Set 0) (Cnstr (Box {prop})) (MkBox {arg}) = do
  info "encountered box"

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

  pure (Cnstr MkBox)

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)

  let True = runIdentity $ convert Relevant (cast sort) expandedType expandedInferredType
    | False => 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 {domain = MkDecl var domain, codomain}) = do
    info "encountered pi"

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

    let domainVar = MkParameter {sort = domainSort, ty = domain, name = var}
    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
          { domainSort = domainVar.sort
          , codomainSort
          , domain = MkDecl (Just var.val) domain
          , codomain
          }
    pure (succ piSort ** (cast piSort, term))

infer' ctx (Lambda {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.type
    arg <- check ctx domainSort domain.type 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 {index = MkDecl var index, element}) = do
    info "encountered sigma"

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

    let indexVar = MkParameter {sort = indexSort, ty = index, name = var}
    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
          { indexSort = indexVar.sort
          , elementSort
          , index = MkDecl (Just 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.type
    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 (Container {input, output, shapeSort, positionSort}) = do
  info "encountered container"

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

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

  let shapeSort = shapeSort.val
  let positionSort = positionSort.val

  let containerSort = foldr1 max
        [ inputSort ~> succ shapeSort
        , inputSort ~> shapeSort ~> succ positionSort
        , inputSort ~> shapeSort ~> positionSort ~> outputSort
        ]
  let container = MkContainerTy {inputSort, shapeSort, positionSort, outputSort, input, output}
  let container = doContainer container

  pure (succ containerSort ** (cast containerSort, container))

infer' ctx (MkContainer {shape, position, next}) = fatal "cannot infer type of container"

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

  (sort ** (type, arg)) <- infer ctx arg
  trace $ pretty {ann} "argument type is" <++> pretty type

  container <- matchContainer type

  let Yes Refl = decEq sort (containerSort container)
    | No _ => fatal "internal universe mismatch"

  let shapeType = containerShapeType container
  shape <- doShape {inputRel = relevance container.inputSort, arg}

  pure ((container.inputSort ~> succ container.shapeSort) ** (shapeType, shape))

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

  (sort ** (type, arg)) <- infer ctx arg
  trace $ pretty {ann} "argument type is" <++> pretty type

  container <- matchContainer type

  let Yes Refl = decEq sort (containerSort container)
    | No _ => fatal "internal universe mismatch"

  shape <- doShape {inputRel = relevance container.inputSort, arg}
  shape <- doApp (weaken [_] shape) (point Nothing Here)

  let positionType = containerPositionType container shape
  position <- doPosition
    { inputRel = relevance container.inputSort
    , shapeRel = relevance container.shapeSort
    , outputRel = relevance container.outputSort
    , arg
    }

  pure ((container.inputSort ~> container.shapeSort ~> succ container.positionSort) **
        (positionType, position))

infer' ctx (Next {arg}) =
  let eq : (a, b, c, d : Universe) -> relevance (a ~> b ~> c ~> d) = relevance d
      eq a b c d =
        rewrite sym $ functionRelevance c d in
        rewrite sym $ functionRelevance b (c ~> d) in
        rewrite sym $ functionRelevance a (b ~> c ~> d) in
        Refl
  in do
  info "encountered nextIndex"

  (sort ** (type, arg)) <- infer ctx arg
  trace $ pretty {ann} "argument type is" <++> pretty type

  container <- matchContainer type

  let Yes Refl = decEq sort (containerSort container)
    | No _ => fatal "internal universe mismatch"

  shape <- doShape {inputRel = relevance container.inputSort, arg}
  shape <- doApp (weaken [_] shape) (point Nothing Here)

  position <- doPosition
    { inputRel = relevance container.inputSort
    , shapeRel = relevance container.shapeSort
    , outputRel = relevance container.outputSort
    , arg
    }
  position <- doApp (weaken [_, _] position) (point Nothing (There Here))
  position <- doApp position (point Nothing Here)

  let nextType = containerNextType container shape position
  nextIndex <- doNext
    { inputRel = relevance container.inputSort
    , shapeRel = relevance container.shapeSort
    , outputRel = relevance container.outputSort
    , arg
    }

  let next' = rewrite eq container.inputSort container.shapeSort container.positionSort container.outputSort in nextIndex

  pure ((container.inputSort ~> container.shapeSort ~> container.positionSort ~> container.outputSort) **
        (nextType, next'))

infer' ctx (Sem {pred = MkLambda var pred, arg}) = do
  info "encountered extension"

  (sort ** (type, arg)) <- infer ctx arg
  trace $ pretty {ann} "argument type is" <++> pretty type

  container <- matchContainer type

  let Yes Refl = decEq sort (containerSort container)
    | No _ => fatal "internal universe mismatch"

  let predCtx = ctx ::< MkParameter
        { sort = container.outputSort
        , ty = container.output
        , name = var
        }
  (predSort, pred) <- inferType predCtx pred

  let semType = containerSemType container predSort
  sem <- doSem container predSort pred arg

  pure ((container.inputSort ~>
         succ (max container.shapeSort (container.positionSort ~> predSort))) **
        (semType, sem))

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 {returnType = MkLambda 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 ::< MkParameter {sort = Set 0, ty = Cnstr Bool, name = var}

    (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 (Box {prop}) = do
  info "encountered Box"

  trace "checking prop is a proposition"
  prop <- check ctx (Set 0) (Cnstr (Universe Prop)) prop
  trace "prop is well typed"

  pure (Set 1 ** (cast (Set 0), Cnstr $ Box {prop}))

infer' ctx (MkBox {arg}) = do
  info "encountered box"

  (Prop ** (prop, Irrel)) <- infer ctx arg
    | (sort ** (type, _)) => universeMismatch (pretty "Prop") (pretty sort)

  pure (Set 0 ** (Cnstr $ Box {prop}, Cnstr MkBox))

infer' ctx (Unbox {arg}) = do
  info "encountered unbox"

  (sort ** (type, arg)) <- infer ctx arg
  trace $ pretty {ann} "argument type is" <++> pretty type

  let Cnstr (Box {prop}) = type
    | _ => typeMismatch (pretty "box") (pretty type)

  pure (Prop ** (prop, Irrel))

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
          { domainSort = indexSort
          , codomainSort = Set 0
          , domain = MkDecl Nothing indexType
          , codomain = 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 ---------------------------------------------

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

    pure $ MkParameter {name = param.name, ty, sort}

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

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}
        ctx def.ty
    debug $ pretty {ann} "\{def.name.val} has type" <++> pretty ty

    tm <- check 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))
checkBlock (blk ::< param) = do
  (_ ** (ctx, Refl)) <- checkBlock blk
  param <- checkParam ctx param
  pure (_ ** (ctx ::< param, Refl))