% PUBLIC DOMAIN - NO LICENSE, NO WARRANTY
% Copyright 2025 TimeHexOn & foxhop & russell@unturf
% https://www.permacomputer.com

%% UncloseAI Erlang Library
%% OpenAI-compatible API client with streaming support
%% Compatible with vLLM, Ollama, and OpenAI-compatible endpoints

-module(uncloseai).
-export([main/0, init/0, init/1, chat/2, chat/3, chat_stream/2, chat_stream/3, tts/2, tts/3]).

-record(client, {models = [], tts_endpoints = [], timeout = 30000}).
-record(model_info, {id, endpoint, max_tokens}).

%%%===================================================================
%%% LIBRARY API
%%%===================================================================

%% Initialize client with auto-discovery
init() ->
    init([]).

init(Opts) ->
    application:ensure_all_started(hackney),
    Timeout = proplists:get_value(timeout, Opts, 30000),

    io:format("Initializing UncloseAI client...~n"),

    % Discover chat/code models
    Models = discover_models_loop(1, []),

    % Discover TTS endpoints
    TtsEndpoints = discover_tts_loop(1, []),

    io:format("Discovered ~p models, ~p TTS endpoints~n~n",
              [length(Models), length(TtsEndpoints)]),

    #client{models = Models, tts_endpoints = TtsEndpoints, timeout = Timeout}.

%% Non-streaming chat completion
chat(Client, Messages) ->
    chat(Client, Messages, []).

chat(Client, Messages, Opts) ->
    ModelIdx = proplists:get_value(model_idx, Opts, 0),
    MaxTokens = proplists:get_value(max_tokens, Opts, 100),
    Temperature = proplists:get_value(temperature, Opts, 0.7),

    #client{models = Models, timeout = Timeout} = Client,

    case ModelIdx < length(Models) of
        false ->
            {error, "Invalid model index"};
        true ->
            Model = lists:nth(ModelIdx + 1, Models),
            #model_info{id = ModelId, endpoint = Endpoint} = Model,
            Url = <<Endpoint/binary, "/chat/completions">>,

            Request = jsx:encode(#{
                model => ModelId,
                messages => Messages,
                stream => false,
                max_tokens => MaxTokens,
                temperature => Temperature
            }),

            case hackney:post(Url,
                             [{<<"Content-Type">>, <<"application/json">>}],
                             Request,
                             [{timeout, Timeout}]) of
                {ok, 200, _Headers, ClientRef} ->
                    {ok, Body} = hackney:body(ClientRef),
                    Data = jsx:decode(Body, [return_maps]),
                    Choices = maps:get(<<"choices">>, Data),
                    [FirstChoice | _] = Choices,
                    Message = maps:get(<<"message">>, FirstChoice),
                    Content = maps:get(<<"content">>, Message),
                    {ok, Content};
                {error, Reason} ->
                    {error, Reason}
            end
    end.

%% Streaming chat completion - returns a function that sends chunks to caller
chat_stream(Client, Messages) ->
    chat_stream(Client, Messages, []).

chat_stream(Client, Messages, Opts) ->
    ModelIdx = proplists:get_value(model_idx, Opts, 0),
    MaxTokens = proplists:get_value(max_tokens, Opts, 500),
    Temperature = proplists:get_value(temperature, Opts, 0.7),

    #client{models = Models, timeout = Timeout} = Client,

    case ModelIdx < length(Models) of
        false ->
            {error, "Invalid model index"};
        true ->
            Model = lists:nth(ModelIdx + 1, Models),
            #model_info{id = ModelId, endpoint = Endpoint} = Model,
            Url = <<Endpoint/binary, "/chat/completions">>,

            Request = jsx:encode(#{
                model => ModelId,
                messages => Messages,
                stream => true,
                max_tokens => MaxTokens,
                temperature => Temperature
            }),

            % Start streaming in a spawned process
            Caller = self(),
            spawn(fun() ->
                case hackney:post(Url,
                                 [{<<"Content-Type">>, <<"application/json">>}],
                                 Request,
                                 [{timeout, Timeout}, {async, once}]) of
                    {ok, ClientRef} ->
                        stream_loop(ClientRef, Caller, <<>>);
                    {error, Reason} ->
                        Caller ! {stream_error, Reason}
                end
            end),
            ok
    end.

%% Text-to-speech generation
tts(Client, Text) ->
    tts(Client, Text, []).

tts(Client, Text, Opts) ->
    Voice = proplists:get_value(voice, Opts, <<"alloy">>),
    OutputFile = proplists:get_value(output_file, Opts, "speech.mp3"),

    #client{tts_endpoints = TtsEndpoints, timeout = Timeout} = Client,

    case TtsEndpoints of
        [] ->
            {error, "No TTS endpoints available"};
        [Endpoint | _] ->
            Url = <<Endpoint/binary, "/audio/speech">>,
            Request = jsx:encode(#{
                model => <<"tts-1">>,
                voice => Voice,
                input => Text
            }),

            case hackney:post(Url,
                             [{<<"Content-Type">>, <<"application/json">>}],
                             Request,
                             [{timeout, Timeout}]) of
                {ok, 200, _Headers, ClientRef} ->
                    {ok, Body} = hackney:body(ClientRef),
                    file:write_file(OutputFile, Body),
                    {ok, OutputFile};
                {error, Reason} ->
                    {error, Reason}
            end
    end.

%%%===================================================================
%%% Private Functions
%%%===================================================================

discover_models_from_endpoint(Endpoint) ->
    Url = <<Endpoint/binary, "/models">>,
    case hackney:get(Url, [], <<>>, [{timeout, 10000}]) of
        {ok, 200, _Headers, ClientRef} ->
            {ok, Body} = hackney:body(ClientRef),
            Data = jsx:decode(Body, [return_maps]),
            case maps:get(<<"data">>, Data, undefined) of
                undefined -> [];
                ModelList ->
                    % Filter out modelperm-* entries
                    lists:filtermap(fun(Model) ->
                        ModelId = maps:get(<<"id">>, Model),
                        case binary:match(ModelId, <<"modelperm-">>) of
                            {0, _} -> false;
                            nomatch ->
                                MaxTokens = maps:get(<<"max_model_len">>, Model, 8192),
                                {true, #model_info{
                                    id = ModelId,
                                    endpoint = Endpoint,
                                    max_tokens = MaxTokens
                                }}
                        end
                    end, ModelList)
            end;
        _ ->
            []
    end.

discover_models_loop(I, Acc) when I > 9999 ->
    lists:flatten(lists:reverse(Acc));
discover_models_loop(I, Acc) ->
    VarName = "MODEL_ENDPOINT_" ++ integer_to_list(I),
    case os:getenv(VarName) of
        false ->
            lists:flatten(lists:reverse(Acc));
        Endpoint ->
            EndpointBin = list_to_binary(Endpoint),
            io:format("Endpoint ~p: ~s~n", [I, EndpointBin]),
            NewModels = discover_models_from_endpoint(EndpointBin),
            discover_models_loop(I + 1, [NewModels | Acc])
    end.

discover_tts_loop(I, Acc) when I > 9999 ->
    lists:reverse(Acc);
discover_tts_loop(I, Acc) ->
    VarName = "TTS_ENDPOINT_" ++ integer_to_list(I),
    case os:getenv(VarName) of
        false ->
            lists:reverse(Acc);
        Endpoint ->
            EndpointBin = list_to_binary(Endpoint),
            discover_tts_loop(I + 1, [EndpointBin | Acc])
    end.

%% Stream processing loop
stream_loop(ClientRef, Caller, Buffer) ->
    case hackney:stream_body(ClientRef) of
        {ok, Data} ->
            NewBuffer = <<Buffer/binary, Data/binary>>,
            {Lines, Remaining} = extract_lines(NewBuffer),

            lists:foreach(fun(Line) ->
                case parse_sse_line(Line) of
                    {ok, Content} -> Caller ! {stream_chunk, Content};
                    done -> Caller ! stream_done;
                    skip -> ok
                end
            end, Lines),

            stream_loop(ClientRef, Caller, Remaining);
        done ->
            Caller ! stream_done;
        {error, Reason} ->
            Caller ! {stream_error, Reason}
    end.

extract_lines(Binary) ->
    Lines = binary:split(Binary, <<"\n">>, [global]),
    case lists:last(Lines) of
        <<>> ->
            {lists:droplast(Lines), <<>>};
        Partial ->
            {lists:droplast(Lines), Partial}
    end.

parse_sse_line(Line) ->
    case binary:match(Line, <<"data: ">>) of
        {0, 6} ->
            Data = binary:part(Line, 6, byte_size(Line) - 6),
            case Data of
                <<"[DONE]">> -> done;
                _ ->
                    try
                        Json = jsx:decode(Data, [return_maps]),
                        case maps:find(<<"choices">>, Json) of
                            {ok, [Choice | _]} ->
                                case maps:find(<<"delta">>, Choice) of
                                    {ok, Delta} ->
                                        case maps:find(<<"content">>, Delta) of
                                            {ok, Content} -> {ok, Content};
                                            error -> skip
                                        end;
                                    error -> skip
                                end;
                            error -> skip
                        end
                    catch
                        _:_ -> skip
                    end
            end;
        nomatch ->
            skip
    end.

%%%===================================================================
%%% Demo Program
%%%===================================================================

main() ->
    io:format("=== UncloseAI Erlang Client (with Streaming) ===~n~n"),

    % Initialize client
    Client = init(),

    #client{models = Models, tts_endpoints = TtsEndpoints} = Client,

    case Models of
        [] ->
            io:format("ERROR: No models discovered~n"),
            halt(1);
        _ ->
            % Non-streaming chat example
            io:format("=== Non-Streaming Chat ===~n"),
            [FirstModel | _] = Models,
            #model_info{id = ModelId} = FirstModel,
            io:format("Model: ~s~n", [ModelId]),

            Messages = [#{role => <<"user">>,
                         content => <<"Explain quantum computing in one sentence">>}],
            case chat(Client, Messages) of
                {ok, Response} ->
                    io:format("Response: ~s~n~n", [Response]);
                {error, Reason} ->
                    io:format("Error: ~p~n~n", [Reason])
            end,

            % Streaming chat example
            ModelIdx = case length(Models) >= 2 of
                true -> 1;
                false -> 0
            end,
            Model = lists:nth(ModelIdx + 1, Models),
            #model_info{id = ModelId2} = Model,
            io:format("=== Streaming Chat ===~n"),
            io:format("Model: ~s~n", [ModelId2]),
            io:format("Response: ", []),

            StreamMessages = [#{role => <<"user">>,
                               content => <<"Write a hello world program in Erlang">>}],
            chat_stream(Client, StreamMessages, [{model_idx, ModelIdx}]),

            % Receive and print streaming chunks
            stream_receive_loop(),
            io:format("~n~n"),

            % TTS example
            case TtsEndpoints of
                [] ->
                    ok;
                _ ->
                    io:format("=== TTS Speech Generation ===~n"),
                    io:format("Model: tts-1~n"),

                    case tts(Client, <<"Hello from UncloseAI Erlang client!">>,
                            [{output_file, "/tmp/speech.mp3"}]) of
                        {ok, File} ->
                            io:format("Audio saved to ~s~n", [File]);
                        {error, TtsReason} ->
                            io:format("TTS failed: ~p~n", [TtsReason])
                    end
            end,

            io:format("~n=== Examples Complete ===~n")
    end.

stream_receive_loop() ->
    receive
        {stream_chunk, Content} ->
            io:format("~s", [Content]),
            stream_receive_loop();
        stream_done ->
            ok;
        {stream_error, Reason} ->
            io:format("~nError: ~p~n", [Reason])
    after 30000 ->
        io:format("~nTimeout~n")
    end.
