:- module(types,
    [ infer_type/3, % infer_type(+AST, +Env, -Type)
      unify_types/3, % unify_types(+Type1, +Type2, -UnifiedType)
      refine_env/4,  % refine_env(+Var, +Type, +EnvIn, -EnvOut)
      get_type/3,    % get_type(+Var, +Env, -Type)
      % Type representations (examples)
      type_number/0, type_string/0, type_boolean/0, type_list_nil/0,
      type_list/1, type_tuple/1, type_union/2, type_intersection/2, type_negation/1,
      type_any/0, type_never/0,
      initial_env/1
    ]).

:- use_module(log).
:- discontiguous unify_types/3.
:- discontiguous infer_type/3. % Added to handle infer_type_arg/3 in between

% --- Type Representations (as atoms/compound terms) ---
type_number :- _ = number.
type_string :- _ = string.
type_boolean :- _ = boolean. % Represents the type 'boolean'
type_list_nil :- _ = list_nil. % AST node for empty list literal '()'
type_list(_T) :- _ = list(_). % _T is intentionally a singleton, structure check: list(Anything)
type_tuple(Ts) :- _ = tuple(Ts), is_list(Ts). % Ts is used, not singleton
type_union(T1, T2) :- _ = union(T1, T2). % T1, T2 are used, not singletons
type_intersection(T1, T2) :- _ = intersection(T1, T2). % T1, T2 are used, not singletons
type_negation(T) :- _ = negation(T). % T is used, not singleton
type_any :- _ = any. % Top type
type_never :- _ = never. % Bottom type, result of failed branches or contradictions

% --- Environment ---
% Env is a list of Var:Type pairs.
initial_env([]).

get_type(Var, [Var:Type | _], Type) :- !.
get_type(Var, [_ | RestEnv], Type) :- get_type(Var, RestEnv, Type).
get_type(Var, [], _) :-
    log(error, unbound_variable(Var, 'unknown_location')), % Location should be passed
    fail.

refine_env(Var, Type, EnvIn, [Var:Type | EnvSansOldBinding]) :-
    delete(EnvIn, Var:_, EnvSansOldBinding), % EnvSansOldBinding is EnvIn with all Var:_ bindings removed.
    log(type_refinement, env_refined(Var, Type)).


% --- Type Inference ---
infer_type(int(_), _Env, number) :- log(type_inference, 'Integer literal -> number').
infer_type(string_val(_), _Env, string) :- log(type_inference, 'String literal -> string').
infer_type(bool(_), _Env, boolean) :- log(type_inference, 'Boolean literal -> boolean').
infer_type(list_nil, _Env, list(never)) :- log(type_inference, 'Empty list literal -> list(never)'). % Or a polymorphic list type

infer_type(id(Var), Env, Type) :-
    ( get_type(Var, Env, Type) ->
        log(type_inference, id(Var) -> Type)
    ; log(error, type_error(unbound_variable(Var), id(Var))),
      Type = never, % Or fail, depending on error handling strategy
      explain_error(unbound_variable(Var, id(Var)), _Msg)
    ).

infer_type(let(Var, ValueAst, BodyAst), EnvIn, BodyType) :-
    log(type_inference, 'Inferring type for let expression'),
    infer_type(ValueAst, EnvIn, ValueType),
    log(type_inference, let_value(Var, ValueType)),
    refine_env(Var, ValueType, EnvIn, EnvMid),
    infer_type(BodyAst, EnvMid, BodyType),
    log(type_inference, let_body(BodyType)).

infer_type(if(CondAst, ThenAst, ElseAst), EnvIn, IfType) :-
    log(type_inference, 'Inferring type for if expression'),
    infer_type(CondAst, EnvIn, CondType),
    ( CondType == boolean -> true
    ; log(error, type_error(expected_boolean_condition, CondAst)), IfType = never, fail
    ),
    % Flow-sensitive refinement example:
    ( CondAst = is_number(id(X)) -> % If condition is is_number(X)
        log(type_inference, flow_refinement_condition(is_number(id(X)))),
        refine_env(X, number, EnvIn, EnvThen) % X is number in Then branch
    ; EnvThen = EnvIn % No specific refinement from condition structure
    ),
    infer_type(ThenAst, EnvThen, ThenType),
    % For Else branch, if CondAst was `is_number(X)`, then X is `not(number)`
    ( CondAst = is_number(id(X)) ->
        get_type(X, EnvIn, OriginalXType), % Get original type of X before refinement
        refine_env(X, intersection(OriginalXType, negation(number)), EnvIn, EnvElse)
    ; EnvElse = EnvIn
    ),
    infer_type(ElseAst, EnvElse, ElseType),
    unify_types(ThenType, ElseType, IfType), % Branches must have compatible types
    log(type_inference, if_expression(CondType, ThenType, ElseType) -> IfType).

% Example: is_number/1 predicate (built-in)
infer_type(is_number(ArgAst), Env, boolean) :-
    log(type_inference, 'Inferring type for is_number/1 call'),
    infer_type(ArgAst, Env, _ArgType). % ArgType can be anything, is_number checks it.

% Lambda expressions (placeholder - full function type inference is complex)
infer_type(lambda(_Params, _BodyAst), _Env, any) :- % For (lambda (params...) body)
    % A proper implementation would construct a function type: fun_type(ParamTypes, ReturnType)
    % This requires inferring types for params (possibly from annotations) and body.
    log(type_inference, 'Lambda expression -> any (placeholder)').

% General function application (placeholder - requires function type for FunctorSExpr)
infer_type(apply(FunctorSExpr, ArgsSExprs), Env, any) :- % For ((lambda ...) arg) or (f arg) where f is complex
    log(type_inference, 'General application (apply/2) -> any (placeholder)'),
    infer_type(FunctorSExpr, Env, FunctorType),
    % Infer types of ArgsSExprs
    maplist(infer_type_arg(Env), ArgsSExprs, ArgTypes),
    log(type_inference, apply_functor_type(FunctorType)),
    log(type_inference, apply_arg_types(ArgTypes)).
    % A proper implementation would:
    % 1. Ensure FunctorType is a function type, e.g., fun_type(ExpectedParamTypes, ReturnType).
    % 2. Check Arity and if ArgTypes are subtypes of ExpectedParamTypes.
    % 3. Return ReturnType.
    % For now, it's 'any'.

infer_type_arg(Env, ArgSExpr, ArgType) :- infer_type(ArgSExpr, Env, ArgType).


% Example: validate_user/1 (hypothetical predicate that narrows type)
% Assume validate_user/1 takes 'any' and if it succeeds, the arg is a 'user_record'.
% This would typically be declared elsewhere (e.g. function signatures)
% For now, we simulate its effect.
infer_type(validate_user(ArgAst), Env, boolean) :- % validate_user returns boolean
    log(type_inference, 'Inferring type for validate_user/1 call'),
    infer_type(ArgAst, Env, _ArgType).
    % The actual refinement happens in the 'then' branch of an 'if' or similar construct
    % e.g., if validate_user(x) then ... (x is now user_record)

% Pattern Matching
infer_type(match(ExprAst, Clauses), EnvIn, MatchType) :-
    log(type_inference, 'Inferring type for match expression'),
    infer_type(ExprAst, EnvIn, ExprType),
    infer_clause_types(Clauses, ExprType, EnvIn, ClauseTypes),
    ( ClauseTypes = [] -> MatchType = never % Or error: non-exhaustive match if not desired
    ; reduce_types(ClauseTypes, MatchType) % Unify all clause body types
    ),
    log(type_inference, match_result_type(MatchType)).

infer_clause_types([], _ExprType, _EnvIn, []).
infer_clause_types([clause(Pattern, _Guard, BodyAst) | RestClauses], ExprType, EnvIn, [BodyType | RestBodyTypes]) :-
    log(type_inference, inferring_clause_pattern(Pattern)),
    refine_env_from_pattern(Pattern, ExprType, EnvIn, EnvPattern),
    ( EnvPattern == fail -> % Pattern doesn't match ExprType or is contradictory
        log(type_warning, pattern_will_not_match(Pattern, ExprType)),
        BodyType = never % This branch is effectively dead
    ; infer_type(BodyAst, EnvPattern, BodyType)
    ),
    infer_clause_types(RestClauses, ExprType, EnvIn, RestBodyTypes).

% refine_env_from_pattern/4: (+Pattern, +MatchedExprType, +EnvIn, -EnvOutOrFail)
refine_env_from_pattern(pvar(Name), MatchedExprType, EnvIn, EnvOut) :-
    !, refine_env(Name, MatchedExprType, EnvIn, EnvOut).
refine_env_from_pattern(pwild, _MatchedExprType, EnvIn, EnvIn) :- !.
refine_env_from_pattern(pint(_), MatchedExprType, EnvIn, EnvIn) :-
    ( unify_types(MatchedExprType, number, number) -> true % Check if MatchedExprType is compatible with number
    ; log(type_error, pattern_type_mismatch(pint, MatchedExprType)), fail
    ).
refine_env_from_pattern(pstring(_), MatchedExprType, EnvIn, EnvIn) :-
    ( unify_types(MatchedExprType, string, string) -> true % Check if MatchedExprType is compatible with string
    ; log(type_error, pattern_type_mismatch(pstring, MatchedExprType)), fail
    ).
refine_env_from_pattern(pbool(_), MatchedExprType, EnvIn, EnvIn) :- % Added for boolean patterns
    ( unify_types(MatchedExprType, boolean, boolean) -> true % Check if MatchedExprType is compatible with boolean
    ; log(type_error, pattern_type_mismatch(pbool, MatchedExprType)), fail
    ).
refine_env_from_pattern(ptuple(Patterns), MatchedExprType, EnvIn, EnvOut) :-
    ( MatchedExprType = tuple(ElementTypes) ; MatchedExprType = any ), % Allow matching 'any' as a tuple
    ( var(ElementTypes) -> % MatchedExprType was 'any' or tuple(_)
        length(Patterns, L), length(ElementTypes, L), % Infer arity
        maplist(=(any), ElementTypes) % Assume elements are 'any' if not specified
    ; length(Patterns, L1), length(ElementTypes, L2), L1 == L2 % Check arity
    ), !,
    refine_env_from_patterns(Patterns, ElementTypes, EnvIn, EnvOut).
refine_env_from_pattern(ptuple(_Patterns), MatchedExprType, _EnvIn, fail) :-
    log(type_error, pattern_type_mismatch(ptuple, MatchedExprType)), fail.

refine_env_from_pattern(plist(Patterns), MatchedExprType, EnvIn, EnvOut) :-
    ( MatchedExprType = list(ElementType) ; MatchedExprType = any ),
    ( var(ElementType) -> ElementType = any ), % If MatchedExprType was 'any' or list(_), treat element type as 'any'
    !,
    length(Patterns, Len),
    length(TypesForPatterns, Len),          % Create a list of unbound variables of the same length as Patterns
    maplist(=(ElementType), TypesForPatterns), % Unify each variable in TypesForPatterns with ElementType
    refine_env_from_patterns(Patterns, TypesForPatterns, EnvIn, EnvOut).
refine_env_from_pattern(plist(_Patterns), MatchedExprType, _EnvIn, fail) :-
    \+ (MatchedExprType = list(_); MatchedExprType = any), % Fail only if not a list or any
    log(type_error, pattern_type_mismatch(plist, MatchedExprType)),
    fail.


refine_env_from_patterns([], [], Env, Env).
refine_env_from_patterns([P|Ps], [T|Ts], EnvIn, EnvOut) :-
    refine_env_from_pattern(P, T, EnvIn, EnvMid),
    ( EnvMid == fail -> EnvOut = fail, !
    ; refine_env_from_patterns(Ps, Ts, EnvMid, EnvOut)
    ).


% --- Type Unification ---
unify_types(T, T, T) :- !, log(unification, identical(T)).
unify_types(any, T, T) :- !, log(unification, any_with(T) -> T).
unify_types(T, any, T) :- !, log(unification, T -> T). % Corrected T_with_any to T
unify_types(never, _T, never) :- !, log(unification, 'never involved'). % Or should it be T? Depends on meaning.
unify_types(_T, never, never) :- !, log(unification, 'never involved').

unify_types(list(T1), list(T2), list(TU)) :- !,
    unify_types(T1, T2, TU),
    log(unification, list(T1, T2) -> list(TU)).

unify_types(tuple(Ts1), tuple(Ts2), tuple(TUs)) :- !,
    length(Ts1, L), length(Ts2, L), % Tuples must have same arity
    maplist(unify_types, Ts1, Ts2, TUs),
    log(unification, tuple(Ts1, Ts2) -> tuple(TUs)).

% Union Type Unification (simplified: create a canonical union)
unify_types(union(A, B), C, union(A, union(B,C))) :- \+ is_union(C), !. % Simplistic, needs canonical form
unify_types(A, union(B, C), union(A, union(B,C))) :- \+ is_union(A), !.
unify_types(union(A1,B1), union(A2,B2), union(A1,union(B1,union(A2,B2)))) :- !. % Very naive

is_union(union(_,_)).

% Intersection (placeholder)
unify_types(intersection(A,B), C, intersection(A,intersection(B,C))) :- !. % Needs proper logic

% Helper to reduce a list of types to a single type (e.g., for match clauses)
reduce_types([T], T) :- !.
reduce_types([T1, T2 | Ts], ResultType) :-
    unify_types(T1, T2, UnifiedHead),
    ( UnifiedHead == never -> ResultType = never, ! % Propagate failure
    ; reduce_types([UnifiedHead | Ts], ResultType)
    ).
reduce_types([], never). % Consistent with match behavior for no clauses.


% --- Example Predicate for Type Narrowing ---
% This would be part of function signature definitions.
% For 'validate_user(X)', if it returns true, X's type is refined.
% This is usually handled by the 'if' construct using the boolean result.
% e.g. if validate_user(x) then (env refined for x) else (env not refined or negated refinement)

% Example:
% ?- initial_env(Env), infer_type(if(is_number(id(x)), id(x), string_val("no")), [x:union(number,string)], Type).
% Type = number  (because string_val("no") is string, unified with number from then branch, if x is number then x is number, else x is string.
% The unification of number and string should fail, or result in union(number, string).
% The example above needs careful check of unify_types.
% Let's assume unify_types(T1,T2,union(T1,T2)) if they are different base types.

% Corrected unification for disparate types (common supertype or union)
unify_types(T1, T2, union(T1, T2)) :-
    % This is a fallback if no other rule matches and T1, T2 are not 'any' or 'never'
    T1 \= any, T2 \= any, T1 \= never, T2 \= never,
    T1 \= T2, % Not identical
    % Ensure not to double-wrap unions unnecessarily (basic check)
    \+ (T1 = union(_,_) ; T2 = union(_,_)),
    log(unification, disparate_types(T1, T2) -> union(T1,T2)).