module Bergamot.Syntax exposing (Term(..))

import Set exposing (Set)
import Dict exposing (Dict)
import Maybe exposing (Maybe)
import Tuple
import Debug

type alias Name = String
type alias Metavariable = String

type Term a
    = IntLit Int
    | StringLit String
    | Call Name (List (Term a))
    | Var a

type alias UnificationInfo =
    { equivalence : Set Metavariable
    , term : Maybe (Term Metavariable)
    }

type alias UnificationState = Dict Metavariable UnificationInfo

reconcile : Set Metavariable -> Maybe (Term Metavariable) -> UnificationState -> UnificationState
reconcile eq mt us =
    let newValue = { equivalence = eq, term = mt }
    in Set.foldl (\v -> Dict.insert v newValue) us eq

merge : Metavariable -> Metavariable -> UnificationState -> Maybe UnificationState
merge v1 v2 us =
    case (Dict.get v1 us, Dict.get v2 us) of
        (Just ui1, Just ui2) ->
            let
                newEq = Set.union ui1.equivalence ui2.equivalence
            in
                case (ui1.term, ui2.term) of
                    (Just t1, Just t2) ->
                        unify t1 t2 us
                        |> Maybe.map (\(t, usp) -> reconcile newEq (Just t) usp)
                    (Just t1, Nothing) ->
                        Just (reconcile newEq (Just t1) us)
                    (Nothing, Just t2) ->
                        Just (reconcile newEq (Just t2) us)
                    (Nothing, Nothing) ->
                        Just (reconcile newEq Nothing us)
        (Just ui1, Nothing) -> Just (reconcile (Set.insert v2 ui1.equivalence) ui1.term us)
        (Nothing, Just ui2) -> Just (reconcile (Set.insert v1 ui2.equivalence) ui2.term us)
        (Nothing, Nothing) -> Just (reconcile (Set.fromList [v1,v2]) Nothing us)

set : Metavariable -> Term Metavariable -> UnificationState -> Maybe (Term Metavariable, UnificationState)
set v t us =
    case Dict.get v us of
        Just ui ->
            case ui.term of
                Just tp ->
                    unify t tp us
                    |> Maybe.map (\(tpp, usp) -> (tpp, reconcile ui.equivalence (Just tpp) usp))
                Nothing ->
                    Just (t, reconcile ui.equivalence (Just t) us)
        Nothing -> Just (t, Dict.insert v { equivalence = Set.singleton v, term = Just t } us)

unifyList : List (Term Metavariable) -> List (Term Metavariable) -> UnificationState -> Maybe (List (Term Metavariable), UnificationState)
unifyList l1 l2 us =
    case (l1, l2) of
        (t1 :: ts1, t2 :: ts2) ->
            unify t1 t2 us
            |> Maybe.andThen (\(t, usp) -> unifyList ts1 ts2 usp)
        ([], []) -> Just ([], us)
        _ -> Nothing

unify : Term Metavariable -> Term Metavariable -> UnificationState -> Maybe (Term Metavariable, UnificationState)
unify t1 t2 us =
    case (t1, t2) of
        (IntLit i1, IntLit i2) -> if i1 == i2 then Just (t1, us) else Nothing
        (StringLit s1, StringLit s2) -> if s1 == s2 then Just (t1, us) else Nothing
        (Call n1 ts1, Call n2 ts2) ->
            if n1 == n2
            then Maybe.map (Tuple.mapFirst (Call n1)) (unifyList ts1 ts2 us)
            else Nothing
        (Var v1, Var v2) ->
            merge v1 v2 us
            |> Maybe.map (\usp -> (Var v1, usp))
        (Var v1, _) -> set v1 t2 us
        (_, Var v2) -> set v2 t1 us
        _ -> Nothing