using LLama.Common;
using LLama.Sampling;
using LLama.Transformers;
using LLama;
using System;
using System.IO;
using System.Net;
using System.Text;
using System.Text.Json;

public partial class Program
{
    public static async Task Main(string[] args)
    {
        string modelPath = "model gguf pass !!";
        var server = new LLMServer(modelPath);
        await server.StartServer(5000);
    }
}

public class LLMServerRequest
{
    public string UserInput { get; set; } = string.Empty;
    public string? SystemPrompt { get; set; }
    public bool Stateless { get; set; } = false;
    public bool discussion { get; set; } = false;
}

public class LLMServer
{
    private readonly LLamaWeights _model;
    private readonly ModelParams _modelParams;
    // private readonly  Dictionary<string, ChatSession> _sessions = new();
    private readonly LLamaContext _context;
    private readonly InteractiveExecutor _executor;
    // private readonly StatelessExecutor _statelessExecutor;
    private readonly ChatHistory _chatHistory;
    private readonly  ChatSession _session;
    private const string DefaultSystemPrompt = "ここにシステムプロンプト";
    private const string DefaultInstruction = "あなたはjson形式で返答であるmessageと今の感情であるemotionを「Joy、Trust、Fear、Surprise、Sadness、Disgust、Anger、Hope」のどれか1つを選択しなさい。";
    public LLMServer(string modelPath)
    {
        _modelParams = new ModelParams(modelPath)
        {
            ContextSize = 4096,
            GpuLayerCount = -1
        };
        _model = LLamaWeights.LoadFromFile(_modelParams);

        _context = _model.CreateContext(_modelParams);
        _executor = new InteractiveExecutor(_context);
        _chatHistory = new ChatHistory();
        _chatHistory.AddMessage(AuthorRole.System, DefaultSystemPrompt);
        _chatHistory.AddMessage(AuthorRole.System, DefaultInstruction);

        _session = new ChatSession(_executor, _chatHistory);
        _session.WithHistoryTransform(new PromptTemplateTransformer(_model, withAssistant: true));
        _session.WithOutputTransform(new LLamaTransforms.KeywordTextOutputStreamTransform([_model.Tokens.EndOfTurnToken ?? "User:", "�"],redundancyLength: 5));
    }

    public async Task StartServer(int port = 8000)
    {
        Console.WriteLine($"サーバーを起動中。ポート: {port}");
        var listener = new HttpListener();
        listener.Prefixes.Add($"http://localhost:{port}/");
        listener.Start();
        Console.WriteLine($"サーバーを起動しました。ポート: {port}");

        while (true)
        {
            var context = await listener.GetContextAsync();
            _ = ProcessRequestAsync(context);
        }
    }

    private async Task ProcessRequestAsync(HttpListenerContext context)
    {
        try
        {
            using var reader = new StreamReader(context.Request.InputStream, context.Request.ContentEncoding);
            var jsonInput = await reader.ReadToEndAsync();
            var request = JsonSerializer.Deserialize<LLMServerRequest>(jsonInput);

            if (request == null || string.IsNullOrEmpty(request.UserInput))
            {
                await SendResponseAsync(context, "無効なリクエストです。メッセージが必要です。", 400);
                return;
            }

            var response = await GenerateResponseAsync(request);
            await SendResponseAsync(context, response);
        }
        catch (Exception ex)
        {
            await SendResponseAsync(context, $"エラーが発生しました: {ex.Message}", 500);
        }
    }

    private async Task<string> GenerateResponseAsync(LLMServerRequest request)
    {
        InferenceParams inferenceParams;

        if(request.discussion)
        {
            var gbnf = (await File.ReadAllTextAsync("gbnf pass !!")).Trim();
            inferenceParams = new InferenceParams
            {
                SamplingPipeline = new DefaultSamplingPipeline
                {
                    Temperature = 0.6f,
                    Grammar = new(gbnf, "root"),
                },
                MaxTokens = -1,
                AntiPrompts = [_model.Tokens.EndOfTurnToken ?? "User:"]
            };
        }
        else
        {
            var emotion_gbnf = (await File.ReadAllTextAsync("gbnf pass !!")).Trim();
            inferenceParams = new InferenceParams
            {
                SamplingPipeline = new DefaultSamplingPipeline
                {
                    Temperature = 0.6f,
                    Grammar = new(emotion_gbnf, "root"),
                },
                MaxTokens = -1,
                AntiPrompts = [_model.Tokens.EndOfTurnToken ?? "User:"]
            };
        }

        var output = new StringBuilder();
        if (request.Stateless)
        {
            LLamaContext statelessContext = _model.CreateContext(_modelParams);
            InteractiveExecutor statelessExecutor = new InteractiveExecutor(statelessContext);

            ChatHistory statelessChatHistory = new ChatHistory();

            statelessChatHistory.AddMessage(AuthorRole.System, request.SystemPrompt ?? DefaultSystemPrompt);
            if(request.discussion){
                statelessChatHistory.AddMessage(AuthorRole.System, "あなたは次の議題に対する意見を述べてください。json形式での回答が必要です．その議題に対する賛否stanceをtrueまたはfalseで、その理由reasonを文字列で具体的に回答しなければなりません。");
            }
            ChatSession statelessSession = new ChatSession(statelessExecutor, statelessChatHistory);

            statelessSession.WithHistoryTransform(new PromptTemplateTransformer(_model, withAssistant: true));
            statelessSession.WithOutputTransform(new LLamaTransforms.KeywordTextOutputStreamTransform([_model.Tokens.EndOfTurnToken ?? "User:", "�"],redundancyLength: 5));
            await foreach (string token in statelessSession.ChatAsync(
                new ChatHistory.Message(AuthorRole.User, request.UserInput), inferenceParams))
            {
                output.Append(token);
            }
        }
        else
        {
            await foreach (string token in _session.ChatAsync(
                new ChatHistory.Message(AuthorRole.User, request.UserInput), inferenceParams))
            {
                output.Append(token);
            }
        }

        return output.ToString();
    }

    private async Task SendResponseAsync(HttpListenerContext context, string responseText, int statusCode = 200)
    {
        var response = context.Response;
        response.StatusCode = statusCode;
        response.ContentType = "application/json; charset=utf-8";

        var buffer = Encoding.UTF8.GetBytes(JsonSerializer.Serialize(new { reply = responseText }));
        response.ContentLength64 = buffer.Length;

        await response.OutputStream.WriteAsync(buffer, 0, buffer.Length);
        response.OutputStream.Close();
    }
}

root ::= message

message ::= "{" ws "\"emotion\"" ws ":" ws "\"" emotion "\"" ws "," ws "\"message\"" ws ":" ws string ws "}"

emotion ::= "Joy" | "Trust" | "Fear" | "Surprise" | "Sadness" | "Disgust"  | "Anger" | "Hope "

value ::= array | string | number | boolean | null

boolean ::= "true" | "false"

null ::= "null"

object ::= "{" ws ( string ws ":" ws value ( ws "," ws string ws ":" ws value )* )? "}" ws

array ::= "[" ws ( value ( ws "," ws value )* )? "]" ws

string ::= "\"" ( [^"\\] | "\\" (["\\/bfnrt] | "u" [0-9a-fA-F]{4}) )* "\""

number ::= "-"? ( "0" | [1-9][0-9]* ) ( "." [0-9]+ )? ( [eE] [+-]? [0-9]+ )?

ws ::= [ \t\n\r]*