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

-- 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 : (indexRel, elementRel : Relevance)
  -> (first : NormalForm indexRel ctx)
  -> (second : NormalForm elementRel ctx)
  -> NormalForm (pair indexRel elementRel) ctx
MkPair Irrelevant Irrelevant first second = Irrel
MkPair Irrelevant Relevant first second =
  Cnstr $ Pair {indexRel = Irrelevant, elementRel = Relevant, prf = Relevant, first, second}
MkPair Relevant elementRel first second =
  Cnstr $ Pair {indexRel = Relevant, elementRel, prf = Relevant, first, second}

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

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

infer' : (tyCtx : TyContext 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 : TyContext 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 : TyContext 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 : TyContext 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 : TyContext 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 : TyContext 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"

  let bodyCtx = ctx ::< MkFreeVar domainSort domain.type 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 {domainRel = _, 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 {indexRel = _, elementRel = _, 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)

  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 = 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
          { 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 = 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
          { 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 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 ::< 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
          { 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 ---------------------------------------------

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))