-module(jchat_db).

-export([init/0, 
         create_tables/0,
         create_table/2,
         init_state_counters/0,
         get_user/1, create_user/1, update_user/1, get_user_by_id/1, get_user_by_email/1,
         get_conversation/1, create_conversation/2, update_conversation/2,
         get_message/1, create_message/2, update_message/2,
         get_participant/1, create_participant/2, update_participant/2,
         get_presence/1, update_presence/2,
         query_conversations/2, query_messages/2]).

-include("jchat.hrl").

%% Initialize database
init() ->
    case mnesia:create_schema([node()]) of
        ok -> ok;
        {error, {_, {already_exists, _}}} -> ok;
        Error -> error({schema_creation_failed, Error})
    end,
    mnesia:start(),
    create_tables(),
    ok.

%% Create database tables
create_tables() ->
    % Users table
    create_table(user, [
        {attributes, record_info(fields, user)},
        {ram_copies, [node()]},
        {type, set}
    ]),
    
    % Conversations table
    create_table(conversation, [
        {attributes, record_info(fields, conversation)},
        {ram_copies, [node()]},
        {type, set}
    ]),
    
    % Messages table
    create_table(message, [
        {attributes, record_info(fields, message)},
        {ram_copies, [node()]},
        {type, set}
    ]),
    
    % Participants table
    create_table(participant, [
        {attributes, record_info(fields, participant)},
        {ram_copies, [node()]},
        {type, set}
    ]),
    
    % Presence table
    create_table(presence, [
        {attributes, record_info(fields, presence)},
        {ram_copies, [node()]},
        {type, set}
    ]),
    
    % State counters for JMAP
    create_table(state_counter, [
        {attributes, record_info(fields, state_counter)},
        {ram_copies, [node()]},
        {type, set}
    ]),
    
    mnesia:wait_for_tables([user, conversation, message, participant, 
                           presence, state_counter], 5000),
    
    % Initialize state counters
    init_state_counters().

%% Helper function to create table with error handling
create_table(TableName, Options) ->
    case mnesia:create_table(TableName, Options) of
        {atomic, ok} -> ok;
        {aborted, {already_exists, TableName}} -> ok;
        Error -> error({table_creation_failed, TableName, Error})
    end.

%% User operations
get_user(Id) ->
    get_user_by_id(Id).

get_user_by_id(Id) ->
    case mnesia:dirty_read(user, Id) of
        [User] -> {ok, User};
        [] -> {error, not_found}
    end.

get_user_by_email(Email) ->
    MatchSpec = [{#user{email = '$1', _ = '_'}, [{'=:=', '$1', Email}], ['$_']}],
    case mnesia:dirty_select(user, MatchSpec) of
        [User] -> {ok, User};
        [] -> {error, not_found};
        Users -> {ok, hd(Users)} % Take first match if multiple (shouldn't happen)
    end.

create_user(User) when is_record(User, user) ->
    case mnesia:dirty_write(User) of
        ok -> {ok, User};
        Error -> {error, Error}
    end.

update_user(User) when is_record(User, user) ->
    case mnesia:dirty_write(User) of
        ok -> {ok, User};
        Error -> {error, Error}
    end.

%% Conversation operations
get_conversation(Id) ->
    case mnesia:dirty_read(conversation, Id) of
        [Conv] -> {ok, Conv};
        [] -> {error, not_found}
    end.

create_conversation(Id, Attrs) ->
    Now = jchat_utils:now_iso8601(),
    Conv = #conversation{
        id = Id,
        title = maps:get(<<"title">>, Attrs, null),
        description = maps:get(<<"description">>, Attrs, null),
        created_at = Now,
        updated_at = Now,
        is_archived = maps:get(<<"isArchived">>, Attrs, false),
        is_muted = maps:get(<<"isMuted">>, Attrs, false),
        participant_ids = maps:get(<<"participantIds">>, Attrs, []),
        last_message_id = null,
        last_message_at = null,
        unread_count = 0,
        message_count = 0,
        metadata = maps:get(<<"metadata">>, Attrs, null)
    },
    case mnesia:dirty_write(Conv) of
        ok -> 
            update_state_counter(conversation),
            {ok, Conv};
        Error -> Error
    end.

update_conversation(Id, Updates) ->
    case mnesia:dirty_read(conversation, Id) of
        [Conv] ->
            UpdatedConv = apply_updates(Conv, Updates),
            UpdatedConv2 = UpdatedConv#conversation{updated_at = jchat_utils:now_iso8601()},
            case mnesia:dirty_write(UpdatedConv2) of
                ok -> 
                    update_state_counter(conversation),
                    {ok, UpdatedConv2};
                Error -> Error
            end;
        [] -> 
            {error, not_found}
    end.

%% Message operations
get_message(Id) ->
    case mnesia:dirty_read(message, Id) of
        [Msg] -> {ok, Msg};
        [] -> {error, not_found}
    end.

create_message(Id, Attrs) ->
    Now = jchat_utils:now_iso8601(),
    ConvId = maps:get(<<"conversationId">>, Attrs),
    Msg = #message{
        id = Id,
        conversation_id = ConvId,
        sender_id = maps:get(<<"senderId">>, Attrs, <<"unknown">>),
        sent_at = Now,
        received_at = Now,
        edited_at = null,
        body = maps:get(<<"body">>, Attrs),
        body_type = maps:get(<<"bodyType">>, Attrs, <<"text/plain">>),
        attachments = maps:get(<<"attachments">>, Attrs, null),
        reply_to_message_id = maps:get(<<"replyToMessageId">>, Attrs, null),
        is_system_message = maps:get(<<"isSystemMessage">>, Attrs, false),
        is_deleted = false,
        reactions = null,
        delivery_status = <<"sent">>,
        read_by = [],
        metadata = maps:get(<<"metadata">>, Attrs, null)
    },
    case mnesia:dirty_write(Msg) of
        ok ->
            % Update conversation with new message
            update_conversation_last_message(ConvId, Id, Now),
            update_state_counter(message),
            {ok, Msg};
        Error -> Error
    end.

update_message(Id, Updates) ->
    case mnesia:dirty_read(message, Id) of
        [Msg] ->
            UpdatedMsg = apply_updates(Msg, Updates),
            UpdatedMsg2 = case maps:get(body, Updates, undefined) of
                undefined -> UpdatedMsg;
                _ -> UpdatedMsg#message{edited_at = jchat_utils:now_iso8601()}
            end,
            case mnesia:dirty_write(UpdatedMsg2) of
                ok -> 
                    update_state_counter(message),
                    {ok, UpdatedMsg2};
                Error -> Error
            end;
        [] -> 
            {error, not_found}
    end.

%% Participant operations
get_participant(Id) ->
    case mnesia:dirty_read(participant, Id) of
        [Part] -> {ok, Part};
        [] -> {error, not_found}
    end.

create_participant(Id, Attrs) ->
    Now = jchat_utils:now_iso8601(),
    Part = #participant{
        id = Id,
        conversation_id = maps:get(conversation_id, Attrs),
        user_id = maps:get(user_id, Attrs),
        display_name = maps:get(display_name, Attrs),
        avatar_blob_id = maps:get(avatar_blob_id, Attrs, null),
        role = maps:get(role, Attrs, <<"member">>),
        joined_at = Now,
        last_active_at = null,
        is_active = true,
        permissions = maps:get(permissions, Attrs, [<<"send">>]),
        metadata = maps:get(metadata, Attrs, null)
    },
    case mnesia:dirty_write(Part) of
        ok -> 
            update_state_counter(participant),
            {ok, Part};
        Error -> Error
    end.

update_participant(Id, Updates) ->
    case mnesia:dirty_read(participant, Id) of
        [Part] ->
            UpdatedPart = apply_updates(Part, Updates),
            case mnesia:dirty_write(UpdatedPart) of
                ok -> 
                    update_state_counter(participant),
                    {ok, UpdatedPart};
                Error -> Error
            end;
        [] -> 
            {error, not_found}
    end.

%% Presence operations
get_presence(UserId) ->
    case mnesia:dirty_read(presence, UserId) of
        [Pres] -> {ok, Pres};
        [] -> {error, not_found}
    end.

update_presence(UserId, Updates) ->
    Now = jchat_utils:now_iso8601(),
    Presence = case mnesia:dirty_read(presence, UserId) of
        [Existing] -> apply_updates(Existing, Updates);
        [] -> #presence{
            user_id = UserId,
            status = maps:get(status, Updates, <<"offline">>),
            status_message = maps:get(status_message, Updates, null),
            last_seen_at = null,
            updated_at = Now
        }
    end,
    UpdatedPresence = Presence#presence{updated_at = Now},
    case mnesia:dirty_write(UpdatedPresence) of
        ok -> 
            update_state_counter(presence),
            {ok, UpdatedPresence};
        Error -> Error
    end.

%% Query operations
query_conversations(UserId, Filter) ->
    MatchSpec = build_conversation_match_spec(UserId, Filter),
    Conversations = mnesia:dirty_select(conversation, MatchSpec),
    {ok, Conversations}.

query_messages(Filter, Sort) ->
    MatchSpec = build_message_match_spec(Filter),
    Messages = mnesia:dirty_select(message, MatchSpec),
    SortedMessages = sort_messages(Messages, Sort),
    {ok, SortedMessages}.

%% Internal functions
apply_updates(Record, Updates) ->
    Fields = case Record of
        #conversation{} -> record_info(fields, conversation);
        #message{} -> record_info(fields, message);
        #participant{} -> record_info(fields, participant);
        #presence{} -> record_info(fields, presence)
    end,
    apply_updates(Record, Updates, Fields, 2).

apply_updates(Record, _Updates, [], _Index) ->
    Record;
apply_updates(Record, Updates, [Field|Rest], Index) ->
    UpdatedRecord = case maps:get(Field, Updates, undefined) of
        undefined -> Record;
        Value -> setelement(Index, Record, Value)
    end,
    apply_updates(UpdatedRecord, Updates, Rest, Index + 1).

update_conversation_last_message(ConvId, MsgId, Timestamp) ->
    case mnesia:dirty_read(conversation, ConvId) of
        [Conv] ->
            UpdatedConv = Conv#conversation{
                last_message_id = MsgId,
                last_message_at = Timestamp,
                message_count = Conv#conversation.message_count + 1,
                updated_at = jchat_utils:now_iso8601()
            },
            mnesia:dirty_write(UpdatedConv);
        [] ->
            ok
    end.

update_state_counter(Type) ->
    Key = {<<"default">>, Type}, % {account_id, object_type}
    case mnesia:dirty_read(state_counter, Key) of
        [#state_counter{state = State}] ->
            NewState = integer_to_binary(binary_to_integer(State) + 1),
            Now = jchat_utils:now_iso8601(),
            mnesia:dirty_write(#state_counter{
                account_id = <<"default">>,
                object_type = Type,
                state = NewState,
                updated_at = Now
            });
        [] ->
            Now = jchat_utils:now_iso8601(),
            mnesia:dirty_write(#state_counter{
                account_id = <<"default">>,
                object_type = Type,
                state = <<"1">>,
                updated_at = Now
            })
    end.

build_conversation_match_spec(_UserId, _Filter) ->
    % Simplified - in production would build proper match specs
    [{'$1', [], ['$1']}].

build_message_match_spec(Filter) ->
    case maps:get(<<"inConversation">>, Filter, null) of
        null ->
            % No filter - return all messages
            [{'$1', [], ['$1']}];
        ConversationId ->
            % Filter by conversation ID
            % Match: #message{conversation_id = ConversationId, _ = '_'}
            [{{message, '_', ConversationId, '_', '_', '_', '_', '_', '_', '_', '_', '_', '_', '_', '_', '_', '_'}, [], ['$_']}]
    end.

sort_messages(Messages, _Sort) ->
    % Simplified - in production would implement proper sorting
    Messages.

%% Initialize state counters for JMAP
init_state_counters() ->
    Types = [conversation, message, participant, presence],
    lists:foreach(fun(Type) ->
        Key = {<<"default">>, Type}, % {account_id, object_type}
        case mnesia:dirty_read(state_counter, Key) of
            [] ->
                Now = jchat_utils:now_iso8601(),
                Counter = #state_counter{
                    account_id = <<"default">>,
                    object_type = Type,
                    state = <<"0">>,
                    updated_at = Now
                },
                mnesia:dirty_write(Counter);
            _ ->
                ok % Already exists
        end
    end, Types).