-module(erlmcp_stdio_server).
-behaviour(gen_server).

-include("erlmcp.hrl").

%% API exports
-export([
    start_link/1,
    add_tool/3, add_tool/4,
    add_resource/3, add_resource/4,
    add_prompt/3, add_prompt/4,
    stop/0
]).

%% gen_server callbacks
-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]).

%% Types
-type tool_handler() :: fun((map()) -> binary() | iolist()).
-type resource_handler() :: fun((binary()) -> binary() | iolist()).
-type prompt_handler() :: fun((map()) -> binary() | iolist() | [map()]).

-record(state, {
    tools = #{} :: #{binary() => {binary(), tool_handler(), map() | undefined}},
    resources = #{} :: #{binary() => {binary(), resource_handler()} | {binary(), resource_handler(), binary()}},
    prompts = #{} :: #{binary() => {binary(), prompt_handler(), [map()] | undefined}},
    initialized = false :: boolean(),
    reader_pid :: pid() | undefined
}).

-type state() :: #state{}.

%%====================================================================
%% API Functions
%%====================================================================

-spec start_link(map()) -> {ok, pid()} | {error, term()}.
start_link(Options) ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [Options], []).

-spec add_tool(binary(), binary(), tool_handler()) -> ok | {error, term()}.
add_tool(Name, Description, Handler) ->
    add_tool(Name, Description, Handler, undefined).

-spec add_tool(binary(), binary(), tool_handler(), map() | undefined) -> ok | {error, term()}.
add_tool(Name, Description, Handler, Schema) when is_binary(Name), is_binary(Description), is_function(Handler, 1) ->
    gen_server:call(?MODULE, {add_tool, Name, Description, Handler, Schema}).

-spec add_resource(binary(), binary(), resource_handler()) -> ok | {error, term()}.
add_resource(Uri, Description, Handler) ->
    add_resource(Uri, Description, Handler, <<"text/plain">>).

-spec add_resource(binary(), binary(), resource_handler(), binary()) -> ok | {error, term()}.
add_resource(Uri, Description, Handler, MimeType) when is_binary(Uri), is_binary(Description), is_function(Handler, 1) ->
    gen_server:call(?MODULE, {add_resource, Uri, Description, Handler, MimeType}).

-spec add_prompt(binary(), binary(), prompt_handler()) -> ok | {error, term()}.
add_prompt(Name, Description, Handler) ->
    add_prompt(Name, Description, Handler, undefined).

-spec add_prompt(binary(), binary(), prompt_handler(), [map()] | undefined) -> ok | {error, term()}.
add_prompt(Name, Description, Handler, Arguments) when is_binary(Name), is_binary(Description), is_function(Handler, 1) ->
    gen_server:call(?MODULE, {add_prompt, Name, Description, Handler, Arguments}).

-spec stop() -> ok.
stop() ->
    gen_server:stop(?MODULE).

%%====================================================================
%% gen_server callbacks
%%====================================================================

-spec init([map()]) -> {ok, state()} | {stop, term()}.
init([_Options]) ->
    process_flag(trap_exit, true),

    logger:info("Initializing erlmcp_stdio_server"),

    % Ensure jsx is available
    case catch application:ensure_all_started(jsx) of
        {'EXIT', Reason} ->
            logger:error("Failed to start jsx application: ~p", [Reason]),
            {stop, {jsx_start_error, Reason}};
        {error, Reason} ->
            logger:error("Failed to start jsx application: ~p", [Reason]),
            {stop, {jsx_start_error, Reason}};
        _ ->
            logger:info("jsx started successfully"),
            % Start the stdin reader process
            logger:info("Starting stdin reader process"),
            ServerPid = self(),
            ReaderPid = spawn_link(fun() -> stdin_reader_loop(ServerPid) end),
            logger:info("Stdin reader process started: ~p", [ReaderPid]),
            {ok, #state{reader_pid = ReaderPid}}
    end.

-spec handle_call(term(), {pid(), term()}, state()) -> {reply, term(), state()}.
handle_call({add_tool, Name, Description, Handler, Schema}, _From, State) ->
    NewTools = maps:put(Name, {Description, Handler, Schema}, State#state.tools),
    {reply, ok, State#state{tools = NewTools}};

handle_call({add_resource, Uri, Description, Handler, MimeType}, _From, State) ->
    NewResources = maps:put(Uri, {Description, Handler, MimeType}, State#state.resources),
    {reply, ok, State#state{resources = NewResources}};

handle_call({add_prompt, Name, Description, Handler, Arguments}, _From, State) ->
    NewPrompts = maps:put(Name, {Description, Handler, Arguments}, State#state.prompts),
    {reply, ok, State#state{prompts = NewPrompts}};

handle_call(_Request, _From, State) ->
    {reply, {error, unknown_request}, State}.

-spec handle_cast(term(), state()) -> {noreply, state()}.
handle_cast(_Msg, State) ->
    {noreply, State}.

-spec handle_info(term(), state()) -> {noreply, state()}.
handle_info({stdin_line, Line}, State) ->
    logger:info("Received stdin line: ~s", [Line]),
    NewState = handle_json_rpc_message(Line, State),
    {noreply, NewState};

handle_info({'EXIT', Pid, Reason}, #state{reader_pid = Pid} = State) ->
    logger:info("Stdin reader process ~p exited with reason: ~p", [Pid, Reason]),
    case Reason of
        normal ->
            logger:info("Stdin reader finished normally"),
            {stop, normal, State};
        _ ->
            logger:error("Stdin reader died: ~p", [Reason]),
            {stop, {reader_died, Reason}, State}
    end;

handle_info(Info, State) ->
    logger:warning("Unexpected info message: ~p", [Info]),
    {noreply, State}.

-spec terminate(term(), state()) -> ok.
terminate(_Reason, #state{reader_pid = ReaderPid}) when is_pid(ReaderPid) ->
    logger:info("Terminating, killing reader process ~p", [ReaderPid]),
    exit(ReaderPid, shutdown),
    ok;
terminate(_Reason, _State) ->
    logger:info("Terminating (no reader process)"),
    ok.

-spec code_change(term(), state(), term()) -> {ok, state()}.
code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%%====================================================================
%% Internal Functions - Stdin Reader
%%====================================================================

-spec stdin_reader_loop(pid()) -> no_return().
stdin_reader_loop(Parent) ->
    logger:info("Stdin reader starting, parent: ~p", [Parent]),

    case read_stdin_line() of
        {ok, Line} ->
            CleanLine = string:trim(Line),
            case CleanLine of
                "" ->
                    stdin_reader_loop(Parent);
                _ ->
                    logger:info("Stdin reader got line: ~s", [CleanLine]),
                    Parent ! {stdin_line, CleanLine},
                    stdin_reader_loop(Parent)
            end;
        eof ->
            logger:info("EOF received from stdin, stopping reader"),
            exit(normal);
        {error, Reason} ->
            logger:error("Stdin read error: ~p", [Reason]),
            exit({read_error, Reason})
    end.

-spec read_stdin_line() -> {ok, string()} | eof | {error, term()}.
read_stdin_line() ->
    try
        % Try the standard approach first
        case io:get_line("") of
            eof ->
                eof;
            {error, ReadReason} ->
                {error, ReadReason};
            Line when is_list(Line) ->
                {ok, Line};
            Line when is_binary(Line) ->
                {ok, binary_to_list(Line)};
            Other ->
                logger:warning("Unexpected io:get_line result: ~p", [Other]),
                {error, {unexpected_result, Other}}
        end
    catch
        error:enotsup ->
            % stdin might not be available in this environment
            logger:error("io:get_line not supported in this environment"),
            {error, enotsup};
        Error:ExceptionReason:Stack ->
            logger:error("Exception in read_stdin_line: ~p:~p~n~p", [Error, ExceptionReason, Stack]),
            {error, {exception, Error, ExceptionReason}}
    end.

%%====================================================================
%% Internal Functions - JSON-RPC Message Handling
%%====================================================================

-spec handle_json_rpc_message(string(), state()) -> state().
handle_json_rpc_message(Message, State) ->
    logger:info("Parsing JSON-RPC message: ~s", [Message]),
    try jsx:decode(list_to_binary(Message), [return_maps]) of
        #{<<"method">> := <<"initialize">>, <<"id">> := Id} = Request ->
            logger:info("Handling initialize request with ID: ~p", [Id]),
            handle_initialize(Id, Request, State);
        #{<<"method">> := <<"notifications/initialized">>} ->
            logger:info("Received initialized notification"),
            State;
        #{<<"method">> := Method, <<"id">> := Id} = Request ->
            logger:info("Handling method call: ~s with ID: ~p", [Method, Id]),
            handle_method_call(Method, Id, Request, State);
        Other ->
            logger:warning("Unknown message format: ~p", [Other]),
            State
    catch
        Error:Reason:Stack ->
            logger:error("JSON parse error: ~p:~p~n~p", [Error, Reason, Stack]),
            State
    end.

-spec handle_initialize(term(), map(), state()) -> state().
handle_initialize(Id, _Request, State) ->
    logger:info("Building initialize response for ID: ~p", [Id]),
    Response = #{
        <<"jsonrpc">> => <<"2.0">>,
        <<"id">> => Id,
        <<"result">> => #{
            <<"protocolVersion">> => ?MCP_VERSION,
            <<"capabilities">> => #{
                <<"tools">> => #{<<"listChanged">> => false},
                <<"resources">> => #{<<"subscribe">> => false, <<"listChanged">> => false},
                <<"prompts">> => #{<<"listChanged">> => false}
            },
            <<"serverInfo">> => #{
                <<"name">> => <<"erlmcp-stdio">>,
                <<"version">> => <<"1.0.0">>
            }
        }
    },
    logger:info("Sending initialize response: ~p", [Response]),
    send_response(Response),
    logger:info("Initialize response sent, marking as initialized"),
    State#state{initialized = true}.

-spec handle_method_call(binary(), term(), map(), state()) -> state().
handle_method_call(<<"tools/list">>, Id, _Request, State) ->
    logger:info("Handling tools/list request"),
    Tools = build_tools_list(State),
    Response = #{
        <<"jsonrpc">> => <<"2.0">>,
        <<"id">> => Id,
        <<"result">> => #{<<"tools">> => Tools}
    },
    send_response(Response),
    State;

handle_method_call(<<"tools/call">>, Id, #{<<"params">> := Params}, State) ->
    logger:info("Handling tools/call request with params: ~p", [Params]),
    handle_tool_call(Id, Params, State),
    State;

handle_method_call(<<"resources/list">>, Id, _Request, State) ->
    logger:info("Handling resources/list request"),
    Resources = build_resources_list(State),
    Response = #{
        <<"jsonrpc">> => <<"2.0">>,
        <<"id">> => Id,
        <<"result">> => #{<<"resources">> => Resources}
    },
    send_response(Response),
    State;

handle_method_call(<<"resources/read">>, Id, #{<<"params">> := Params}, State) ->
    logger:info("Handling resources/read request with params: ~p", [Params]),
    handle_resource_read(Id, Params, State),
    State;

handle_method_call(<<"prompts/list">>, Id, _Request, State) ->
    logger:info("Handling prompts/list request"),
    Prompts = build_prompts_list(State),
    Response = #{
        <<"jsonrpc">> => <<"2.0">>,
        <<"id">> => Id,
        <<"result">> => #{<<"prompts">> => Prompts}
    },
    send_response(Response),
    State;

handle_method_call(<<"prompts/get">>, Id, #{<<"params">> := Params}, State) ->
    logger:info("Handling prompts/get request with params: ~p", [Params]),
    handle_prompt_get(Id, Params, State),
    State;

handle_method_call(Method, Id, _Request, State) ->
    logger:warning("Unknown method: ~s", [Method]),
    Error = #{
        <<"jsonrpc">> => <<"2.0">>,
        <<"id">> => Id,
        <<"error">> => #{
            <<"code">> => -32601,
            <<"message">> => <<"Method not found: ", Method/binary>>
        }
    },
    send_response(Error),
    State.

%%====================================================================
%% Internal Functions - Response Building
%%====================================================================

-spec build_tools_list(state()) -> [map()].
build_tools_list(State) ->
    maps:fold(fun(Name, {Description, _Handler, Schema}, Acc) ->
        Tool = #{
            <<"name">> => Name,
            <<"description">> => Description
        },
        ToolWithSchema = case Schema of
            undefined -> Tool;
            _ -> Tool#{<<"inputSchema">> => Schema}
        end,
        [ToolWithSchema | Acc]
    end, [], State#state.tools).

-spec build_resources_list(state()) -> [map()].
build_resources_list(State) ->
    maps:fold(fun
        (Uri, {Description, _Handler}, Acc) ->
            Resource = #{
                <<"uri">> => Uri,
                <<"name">> => Description,
                <<"mimeType">> => <<"text/plain">>
            },
            [Resource | Acc];
        (Uri, {Description, _Handler, MimeType}, Acc) ->
            Resource = #{
                <<"uri">> => Uri,
                <<"name">> => Description,
                <<"mimeType">> => MimeType
            },
            [Resource | Acc]
    end, [], State#state.resources).

-spec build_prompts_list(state()) -> [map()].
build_prompts_list(State) ->
    maps:fold(fun(Name, {Description, _Handler, Arguments}, Acc) ->
        Prompt = #{
            <<"name">> => Name,
            <<"description">> => Description
        },
        PromptWithArgs = case Arguments of
            undefined -> Prompt;
            _ -> Prompt#{<<"arguments">> => Arguments}
        end,
        [PromptWithArgs | Acc]
    end, [], State#state.prompts).

%%====================================================================
%% Internal Functions - Handler Execution
%%====================================================================

-spec handle_tool_call(term(), map(), state()) -> ok.
handle_tool_call(Id, #{<<"name">> := Name, <<"arguments">> := Arguments}, State) ->
    logger:info("Executing tool: ~s with arguments: ~p", [Name, Arguments]),
    case maps:get(Name, State#state.tools, undefined) of
        undefined ->
            logger:warning("Tool not found: ~s", [Name]),
            Error = #{
                <<"jsonrpc">> => <<"2.0">>,
                <<"id">> => Id,
                <<"error">> => #{
                    <<"code">> => -32601,
                    <<"message">> => <<"Tool not found">>
                }
            },
            send_response(Error);
        {_Description, Handler, _Schema} ->
            try
                Result = Handler(Arguments),
                ResultBinary = to_binary(Result),
                logger:info("Tool execution result: ~s", [ResultBinary]),
                Response = #{
                    <<"jsonrpc">> => <<"2.0">>,
                    <<"id">> => Id,
                    <<"result">> => #{
                        <<"content">> => [#{
                            <<"type">> => <<"text">>,
                            <<"text">> => ResultBinary
                        }]
                    }
                },
                send_response(Response)
            catch
                Class:Reason:Stack ->
                    logger:error("Tool handler crashed: ~p:~p~n~p", [Class, Reason, Stack]),
                    Error = #{
                        <<"jsonrpc">> => <<"2.0">>,
                        <<"id">> => Id,
                        <<"error">> => #{
                            <<"code">> => -32603,
                            <<"message">> => <<"Internal error">>
                        }
                    },
                    send_response(Error)
            end
    end;
handle_tool_call(Id, Params, _State) ->
    logger:warning("Invalid tool call parameters: ~p", [Params]),
    Error = #{
        <<"jsonrpc">> => <<"2.0">>,
        <<"id">> => Id,
        <<"error">> => #{
            <<"code">> => -32602,
            <<"message">> => <<"Invalid parameters">>
        }
    },
    send_response(Error).

-spec handle_resource_read(term(), map(), state()) -> ok.
handle_resource_read(Id, #{<<"uri">> := Uri}, State) ->
    logger:info("Reading resource: ~s", [Uri]),
    case maps:get(Uri, State#state.resources, undefined) of
        undefined ->
            logger:warning("Resource not found: ~s", [Uri]),
            Error = #{
                <<"jsonrpc">> => <<"2.0">>,
                <<"id">> => Id,
                <<"error">> => #{
                    <<"code">> => -32602,
                    <<"message">> => <<"Resource not found">>
                }
            },
            send_response(Error);
        {_Description, Handler} ->
            handle_resource_with_handler(Id, Uri, Handler, <<"text/plain">>);
        {_Description, Handler, MimeType} ->
            handle_resource_with_handler(Id, Uri, Handler, MimeType)
    end;
handle_resource_read(Id, Params, _State) ->
    logger:warning("Invalid resource read parameters: ~p", [Params]),
    Error = #{
        <<"jsonrpc">> => <<"2.0">>,
        <<"id">> => Id,
        <<"error">> => #{
            <<"code">> => -32602,
            <<"message">> => <<"Invalid parameters">>
        }
    },
    send_response(Error).

-spec handle_resource_with_handler(term(), binary(), fun(), binary()) -> ok.
handle_resource_with_handler(Id, Uri, Handler, MimeType) ->
    try
        Content = Handler(Uri),
        ContentBinary = to_binary(Content),
        logger:info("Resource content generated for ~s: ~s", [Uri, ContentBinary]),
        Response = #{
            <<"jsonrpc">> => <<"2.0">>,
            <<"id">> => Id,
            <<"result">> => #{
                <<"contents">> => [#{
                    <<"uri">> => Uri,
                    <<"mimeType">> => MimeType,
                    <<"text">> => ContentBinary
                }]
            }
        },
        send_response(Response)
    catch
        Class:Reason:Stack ->
            logger:error("Resource handler crashed: ~p:~p~n~p", [Class, Reason, Stack]),
            Error = #{
                <<"jsonrpc">> => <<"2.0">>,
                <<"id">> => Id,
                <<"error">> => #{
                    <<"code">> => -32603,
                    <<"message">> => <<"Internal error">>
                }
            },
            send_response(Error)
    end.

-spec handle_prompt_get(term(), map(), state()) -> ok.
handle_prompt_get(Id, #{<<"name">> := Name} = Params, State) ->
    logger:info("Getting prompt: ~s", [Name]),
    case maps:get(Name, State#state.prompts, undefined) of
        undefined ->
            logger:warning("Prompt not found: ~s", [Name]),
            Error = #{
                <<"jsonrpc">> => <<"2.0">>,
                <<"id">> => Id,
                <<"error">> => #{
                    <<"code">> => -32602,
                    <<"message">> => <<"Prompt not found">>
                }
            },
            send_response(Error);
        {_Description, Handler, _Arguments} ->
            try
                Arguments = maps:get(<<"arguments">>, Params, #{}),
                Result = Handler(Arguments),
                Messages = normalize_prompt_result(Result),
                logger:info("Prompt result generated for ~s", [Name]),
                Response = #{
                    <<"jsonrpc">> => <<"2.0">>,
                    <<"id">> => Id,
                    <<"result">> => #{
                        <<"messages">> => Messages
                    }
                },
                send_response(Response)
            catch
                Class:Reason:Stack ->
                    logger:error("Prompt handler crashed: ~p:~p~n~p", [Class, Reason, Stack]),
                    Error = #{
                        <<"jsonrpc">> => <<"2.0">>,
                        <<"id">> => Id,
                        <<"error">> => #{
                            <<"code">> => -32603,
                            <<"message">> => <<"Internal error">>
                        }
                    },
                    send_response(Error)
            end
    end;
handle_prompt_get(Id, Params, _State) ->
    logger:warning("Invalid prompt get parameters: ~p", [Params]),
    Error = #{
        <<"jsonrpc">> => <<"2.0">>,
        <<"id">> => Id,
        <<"error">> => #{
            <<"code">> => -32602,
            <<"message">> => <<"Invalid parameters">>
        }
    },
    send_response(Error).

%%====================================================================
%% Internal Functions - Utilities
%%====================================================================

-spec send_response(map()) -> ok.
send_response(Response) ->
    Json = jsx:encode(Response),
    logger:info("Sending JSON response: ~s", [Json]),
    io:format("~s~n", [Json]),
    logger:info("Response sent to stdout"),
    ok.

-spec to_binary(term()) -> binary().
to_binary(B) when is_binary(B) -> B;
to_binary(L) when is_list(L) -> iolist_to_binary(L);
to_binary(A) when is_atom(A) -> atom_to_binary(A, utf8);
to_binary(I) when is_integer(I) -> integer_to_binary(I);
to_binary(F) when is_float(F) -> float_to_binary(F, [{decimals, 2}]);
to_binary(Other) -> iolist_to_binary(io_lib:format("~p", [Other])).

-spec normalize_prompt_result(term()) -> [map()].
normalize_prompt_result(Result) when is_binary(Result) ->
    [#{
        <<"role">> => <<"user">>,
        <<"content">> => #{
            <<"type">> => <<"text">>,
            <<"text">> => Result
        }
    }];
normalize_prompt_result(Result) when is_list(Result) ->
    case Result of
        [#{} | _] -> Result;  % Already a list of message maps
        _ -> normalize_prompt_result(iolist_to_binary(Result))
    end;
normalize_prompt_result(Result) ->
    normalize_prompt_result(to_binary(Result)).