port module Main exposing (main)

import Browser
import Html exposing (Html)
import Html.Events exposing (onInput, onClick)
import Html.Attributes exposing (type_, class, classList, value)
import Html.Lazy
import Bergamot.Syntax exposing (..)
import Bergamot.Search exposing (..)
import Bergamot.Rules exposing (..)
import Bergamot.Parser exposing (..)
import Bergamot.Latex exposing (..)
import Bergamot.ObjectLanguage exposing (topLevelExpr, exprToTerm)
import Bergamot.Utils exposing (..)
import Json.Encode
import Json.Decode exposing (string, field, list, oneOf, succeed, fail)
import Maybe
import Tuple

type Tab
    = Editor
    | MetaEditor
    | Rendered

tabEq : Tab -> Tab -> Bool
tabEq t1 t2 =
    case (t1, t2) of
        (Editor, Editor) -> True
        (MetaEditor, MetaEditor) -> True
        (Rendered, Rendered) -> True
        _ -> False

type EditMode
    = Query
    | Syntax
    | Custom String

type ResultMode
    = Conclusion
    | Tree

resultModeEq : ResultMode -> ResultMode -> Bool
resultModeEq rm1 rm2 =
    case (rm1, rm2) of
        (Conclusion, Conclusion) -> True
        (Tree, Tree) -> True
        _ -> False

type alias Model =
    { program : String
      -- ^ The Bergamot rules to execute a search against
    , renderProgram: String
      -- ^ The Bergamot render rules to apply when generating LaTeX
    , inputProgram : String
      -- ^ The Bergamot rules to apply to convert a Syntax term into a query
    , tab : Tab
      -- ^ The current tab (editor, redner rule editor, rendered)
    , input : String
      -- ^ The string the user has in the input box
    , desugaredQuery : Maybe String
      -- ^ The Bergamot query corresponding to input. May be Nothing while waiting for
      --   Custom mode, which has user JS convert input to a query.
    , inputModes : List (String, EditMode)
      -- ^ A List of input modes that can be used with the input box.
    , inputMode : Int
      -- ^ The index of the currently selected input mode.
    , resultMode : ResultMode
      -- ^ How the result should be displayed (judgement or proof tree)
    }
type alias Flags =
    { renderRules : String
    , inputRules : String
    , rules : String
    , input : String
    , inputModes : Json.Decode.Value
    }
type Msg
    = SetProgram String
    | SetRenderProgram String
    | SetInput String
    | SetTab Tab
    | SetInputMode Int
    | SetResultMode ResultMode
    | SetDesugaredQuery { input: String, query: String }

decodeInputModes: Json.Decode.Value -> List (String, EditMode)
decodeInputModes val =
    let
        exactString s1 v =
            string
            |> Json.Decode.andThen (\s2 ->
                if s1 == s2
                then succeed v
                else Json.Decode.fail "did not match expected string")
        editModeDecoder = oneOf
            [ exactString "query" Query
            , exactString "syntax" Syntax
            , field "custom" string |> Json.Decode.map Custom
            ]
        decoder = Json.Decode.keyValuePairs editModeDecoder
    in
        case Json.Decode.decodeValue decoder val of
            Ok l -> l
            Err _ -> [("Language Term", Syntax), ("Query", Query)]

-- Convert the user-entered string 'input' using custom query mode 'mode'
port convertInput : { mode : String, input : String } -> Cmd msg

-- Invoked when user code finishes converting 'input' into a Bergamot query
port receiveConverted : ({ input : String, query : String } -> msg) -> Sub msg

convertInputCmd : Int -> List (String, EditMode) -> String -> Cmd Msg
convertInputCmd inputMode inputModes input =
    case getEditMode inputMode inputModes of
        Ok (Custom modeType) ->
            convertInput { mode = modeType, input = input }
        _ -> Cmd.none

init : Flags -> (Model, Cmd Msg)
init fs =
    let
        inputModes = decodeInputModes fs.inputModes
        inputMode = 0
    in
        ( { program = fs.rules
          , renderProgram = fs.renderRules
          , inputProgram = fs.inputRules
          , input = fs.input
          , desugaredQuery = Nothing
          , tab = Rendered
          , inputModes = inputModes
          , inputMode = 0
          , resultMode = Conclusion
          }
        , convertInputCmd inputMode inputModes fs.input
        )

viewSection : String -> Html Msg -> Html Msg
viewSection name content =
    Html.div [ class "bergamot-section" ] [ Html.em [ class "bergamot-section-heading" ] [ Html.text name ], content ]

viewTab : Tab -> Html Msg -> Html Msg -> Html Msg -> Html Msg
viewTab tab editor metaEditor rendered =
    case tab of
        Editor -> editor
        MetaEditor -> metaEditor
        Rendered -> rendered

viewSelector : (Int -> a -> b -> Bool) -> (Int -> a -> Msg) -> b -> List (a, String) -> Html Msg
viewSelector eq mkMsg mode modeNames =
    Html.div [ class "bergamot-selector" ] <|
        List.indexedMap (\idx (modep, name) ->
            Html.button
                [ classList [("active", eq idx modep mode)]
                , onClick (mkMsg idx modep)
                ] [ Html.text name ]) modeNames

viewTabSelector : Tab -> List (Tab, String) -> Html Msg
viewTabSelector = viewSelector (\_ -> tabEq) (\_ -> SetTab)

viewInputModeSelector : Int -> List (EditMode, String) -> Html Msg
viewInputModeSelector = viewSelector (\i _ idx -> i == idx) (\idx _ -> SetInputMode idx)

viewResultModeSelector : ResultMode -> List (ResultMode, String) -> Html Msg
viewResultModeSelector = viewSelector (\_ -> resultModeEq) (\_ -> SetResultMode)

viewRule : RuleEnv -> Rule -> Maybe (Html Msg)
viewRule env r = renderRuleViaRules env r
    |> Maybe.map latex

viewRuleSection : RuleEnv -> Section -> Maybe (Html Msg)
viewRuleSection env sec =
    List.map (viewRule env) sec.rules
    |> sequenceMaybes
    |> Maybe.map (\rs ->
        Html.div [ class "bergamot-rule-section" ]
        [ Html.div [ class "bergamot-rule-list" ] rs
        , Html.p [class "bergamot-rule-section-name"] [Html.text (sec.name)]
        ])

viewRules : String -> String -> Html Msg
viewRules renderProgs progs =
    Html.div [ class "bergamot-section-list" ] <|
        case progAndRenderProg progs renderProgs of
            Ok (prog, renderProg) -> List.filterMap (viewRuleSection renderProg) prog.sections
            _ -> []

type Error
    = BadQuery
    | BadObjectTerm
    | NoConversionResults
    | BadInputProg
    | BadProg
    | BadRenderProg
    | FailedRender
    | InvalidInputMode
    | Silent

combineTwoResults : (a -> b -> Result Error c) -> Result Error a -> Result Error b -> Result Error c
combineTwoResults f ra rb =
    case ra of
        Err _ -> Err Silent
        Ok a ->
            case rb of
                Err _ -> Err Silent
                Ok b -> f a b

errorToString : Error -> String
errorToString err =
    case err of
        BadQuery -> "Unable to parse input query"
        BadObjectTerm -> "Unable to parse input object language term"
        NoConversionResults -> "Failed to convert object language term to proof goal"
        BadInputProg -> "Unable to parse conversion rules from object language to proof goals"
        BadProg -> "Unable to parse rules"
        BadRenderProg -> "Unable to parse rendering rules"
        FailedRender -> "Unable to render terms using provided rendering rules"
        InvalidInputMode -> "Invalid edit mode specified"
        Silent -> ""

viewError : Error -> Html Msg
viewError e = Html.div [ class "bergamot-error" ] [ Html.text <| errorToString e ]

viewIfError : Result Error a -> List (Html Msg)
viewIfError r =
    case r of
        Err Silent -> []
        Err e -> [ viewError e ]
        _ -> []

viewOrError : Result Error (Html Msg) -> Html Msg
viewOrError r =
    case r of
        Err Silent -> Html.div [] []
        Err e -> Html.div [] [ viewError e ] -- The div wrapper is needed because Elm has a bug?
        Ok html -> html

getEditMode : Int -> List (String, EditMode) -> Result Error EditMode
getEditMode i l =
    case l of
        [] -> Err InvalidInputMode
        ((_, editMode) :: xs) ->
            if i == 0 then Ok editMode else getEditMode (i - 1) xs

proofGoal : EditMode -> String -> String -> Maybe String -> Result Error (Term Metavariable)
proofGoal editMode inputProgs input desugaredQuery =
    if input == ""
    then Err Silent
    else
        case editMode of
            Query ->
                case run topLevelTerm input of
                    Nothing -> Err BadQuery
                    Just query -> Ok query
            Syntax ->
                case (run program inputProgs, run topLevelExpr input) of
                    (Just inputProg, Just e) ->
                        let
                            inputRule = { name = "UserInput", premises = [], conclusion = Call "input" [exprToTerm e] }
                            fullProgram = { sections = { name = "", rules = [inputRule] } :: inputProg.sections }
                            query = Call "prompt" [Var (MkMetavariable "?p")]
                        in
                            case single fullProgram (prove query |> Bergamot.Rules.andThen reifyProofTree) of
                                Just (MkProofTree node) ->
                                    case node.conclusion of
                                        Call "prompt" [t] -> Ok <| Bergamot.Syntax.map (MkMetavariable << unUnificationVar) t
                                        _ -> Err NoConversionResults
                                _ -> Err NoConversionResults
                    (_, Nothing) -> Err BadObjectTerm
                    (Nothing, _) -> Err BadInputProg
            Custom _ ->
                case desugaredQuery of
                    Just querys ->
                        case run topLevelTerm querys of
                            Nothing -> Err BadQuery
                            Just query -> Ok query
                    Nothing -> Err Silent

progAndRenderProg : String -> String -> Result Error (RuleEnv, RuleEnv)
progAndRenderProg progs renderProgs =
    case (run program progs, run program renderProgs) of
        (Just prog, Just renderProg) -> Ok (prog, renderProg)
        (Nothing, _) -> Err BadProg
        (_, Nothing) -> Err BadRenderProg

renderProofTree : ResultMode -> ProofTree -> RuleEnv -> Result Error (Html Msg)
renderProofTree resultMode (MkProofTree node) renderProg =
    let
        maybeLatexString =
            case resultMode of
                Conclusion -> renderTermViaRules renderProg (quoteVariables node.conclusion)
                Tree -> renderTreeViaRules renderProg (MkProofTree node)
    in
        List.filterMap (Maybe.map latex) [maybeLatexString]
        |> Html.div [ class "bergamot-proof-tree" ]
        |> Ok


viewProofTree : ResultMode -> Term Metavariable -> (RuleEnv, RuleEnv) -> Result Error (Html Msg)
viewProofTree resultMode query (prog, renderProg) =
    case single prog (prove query |> Bergamot.Rules.andThen reifyProofTree) of
        Just proofTree -> renderProofTree resultMode proofTree renderProg
        Nothing -> Ok <| Html.div [ class "bergamot-no-proofs" ]
            [ Html.text "No applicable rules to prove input" ]

view : Model -> Html Msg
view m =
    let
        mode = getEditMode m.inputMode m.inputModes
        termOrErr = mode |> Result.andThen (\editMode -> proofGoal editMode m.inputProgram m.input m.desugaredQuery)
        progsOrErr = progAndRenderProg m.program m.renderProgram
        proofTreeOrErr = combineTwoResults (viewProofTree m.resultMode) termOrErr progsOrErr
    in
        Html.div [ class "bergamot-root" ]
        [ viewSection "Input" <| Html.div [] <|
            [ viewInputModeSelector m.inputMode <|
                List.map (\(a, b) -> (b, a)) m.inputModes
            , Html.input [ type_ "text", onInput SetInput, value m.input ] []
            ] ++
            viewIfError termOrErr
        , viewSection "Result" <| Html.div[]
            [ viewResultModeSelector m.resultMode [(Conclusion, "Conclusion Only"), (Tree, "Full Proof Tree")]
            , viewOrError proofTreeOrErr
            ]
        , viewSection "Rules" <| Html.div [] <|
            [ viewTabSelector m.tab [(Rendered, "Rendered"), (Editor, "Editor"), (MetaEditor, "Presentation Rules")]
            , viewTab m.tab
                (Html.textarea [ onInput SetProgram, value m.program ] [ ])
                (Html.textarea [ onInput SetRenderProgram, value m.renderProgram ] [ ])
                (Html.Lazy.lazy2 viewRules m.renderProgram m.program)
            ] ++
            viewIfError progsOrErr
        ]

update : Msg -> Model -> (Model, Cmd Msg)
update msg m =
    case msg of
        SetProgram prog -> ({ m | program = prog }, Cmd.none)
        SetRenderProgram prog -> ({ m | renderProgram = prog }, Cmd.none)
        SetInput input ->
            ( { m | input = input }
            , convertInputCmd m.inputMode m.inputModes input
            )
        SetTab tab -> ({ m | tab = tab }, Cmd.none)
        SetInputMode mode ->
            ( { m | inputMode = mode }
            , convertInputCmd mode m.inputModes m.input
            )
        SetResultMode mode -> ({ m | resultMode = mode }, Cmd.none)
        SetDesugaredQuery data ->
            ({ m | desugaredQuery =
                if m.input == data.input
                then Just data.query
                else m.desugaredQuery
             }
            , Cmd.none
            )

subscriptions : Model -> Sub Msg
subscriptions _ = receiveConverted SetDesugaredQuery

main =
    Browser.element
        { init = init
        , view = view
        , update = update
        , subscriptions = subscriptions
        }

-- Latex support:
-- Based on https://stackoverflow.com/questions/75492820/embedding-mathematical-equations-in-an-elm-spa

latex : String -> Html Msg
latex expr =
  Html.node "katex-expression"
    [ Html.Attributes.attribute "expression" expr
    , Html.Attributes.attribute "katex-options" (Json.Encode.encode 0 options)
    ]
    []


options : Json.Encode.Value
options =
    Json.Encode.object
        [ ( "displayMode", Json.Encode.bool True )
        ]