// File: WebServer.cs
using System;
using System.IO;
using System.Net;
using System.Net.Http;
using System.Text;
using System.Text.Json.Nodes;
using System.Text.RegularExpressions;
using System.Threading.Tasks;

namespace AI_C2_Server
{
    public class WebServer
    {
        public static WebServer Instance { get; } = new WebServer();

        private HttpListener _listener;
        private bool _isRunning;
        private static readonly HttpClient _httpClient = new HttpClient();

        private WebServer() { }

        public void Start(int port = 8080)
        {
            if (_isRunning) return;

            _listener = new HttpListener();
            _listener.Prefixes.Add($"http://localhost:{port}/");
            _listener.Prefixes.Add($"http://127.0.0.1:{port}/");

            try
            {
                _listener.Start();
                _isRunning = true;
                Task.Run(ListenLoop);
                Logger.Log($"[SYSTEM] Initializing C2 Server on Port {port}...");
            }
            catch (Exception ex)
            {
                Logger.Log($"[ERROR] Failed to start server: {ex.Message}");
            }
        }

        public void Stop()
        {
            if (!_isRunning) return;
            _isRunning = false;
            _listener?.Stop();
            _listener?.Close();
            Logger.Log("[SYSTEM] Server stopped.");
        }

        private async Task ListenLoop()
        {
            while (_isRunning)
            {
                try
                {
                    var context = await _listener.GetContextAsync();
                    _ = Task.Run(() => ProcessRequestAsync(context));
                }
                catch (HttpListenerException)
                {
                    // Thrown when listener is stopped, safe to ignore
                }
                catch (Exception ex)
                {
                    Logger.Log($"[ERROR] Listener loop error: {ex.Message}");
                }
            }
        }

        private async Task ProcessRequestAsync(HttpListenerContext context)
        {
            var request = context.Request;
            var response = context.Response;

            // CORS Headers
            response.AddHeader("Access-Control-Allow-Origin", "*");
            response.AddHeader("Access-Control-Allow-Methods", "GET, POST, OPTIONS");
            response.AddHeader("Access-Control-Allow-Headers", "Content-Type");

            if (request.HttpMethod == "OPTIONS")
            {
                response.StatusCode = 200;
                response.Close();
                return;
            }

            try
            {
                string path = request.Url.AbsolutePath.TrimEnd('/');

                // Extract instance ID for logging
                string instanceId = "";
                if (path.StartsWith("/api/relay/state/")) instanceId = path.Substring("/api/relay/state/".Length);
                else if (path.StartsWith("/api/relay/command/")) instanceId = path.Substring("/api/relay/command/".Length);
                else if (path.StartsWith("/api/relay/result/")) instanceId = path.Substring("/api/relay/result/".Length);

                if (!string.IsNullOrEmpty(instanceId))
                {
                    string platform = OperationCenter.Instance.GetInstancePlatform(instanceId);
                    Logger.Log($"[HTTP] [{platform} | {instanceId}] {request.HttpMethod} {path}");
                }
                else
                {
                    Logger.Log($"[HTTP] {request.HttpMethod} {path}");
                }

                string body = "";

                if (request.HasEntityBody)
                {
                    using var reader = new StreamReader(request.InputStream, request.ContentEncoding);
                    body = await reader.ReadToEndAsync();
                }

                JsonNode payload = null;
                if (!string.IsNullOrWhiteSpace(body))
                {
                    try { payload = JsonNode.Parse(body); } catch { }
                }

                // --- RELAY ENDPOINTS ---
                if (path.StartsWith("/api/relay/state/"))
                {
                    OperationCenter.Instance.UpdateInstanceState(instanceId, payload as JsonObject);
                    SendJson(response, new JsonObject { ["status"] = "ok" });
                }
                else if (path.StartsWith("/api/relay/command/"))
                {
                    OperationCenter.Instance.UpdateInstanceState(instanceId, new JsonObject());
                    var cmd = OperationCenter.Instance.GetPendingCommand(instanceId);
                    SendJson(response, cmd ?? new JsonObject());
                }
                else if (path.StartsWith("/api/relay/result/"))
                {
                    OperationCenter.Instance.LogResult(instanceId, payload as JsonObject);
                    SendJson(response, new JsonObject { ["status"] = "acknowledged" });
                }
                // --- ADMIN ENDPOINTS ---
                else if (path == "/api/admin/instances")
                {
                    SendJson(response, OperationCenter.Instance.GetActiveInstances());
                }
                else if (path == "/api/admin/inject")
                {
                    await HandleCommand(response, "GENERATE", payload?["prompt"]?.DeepClone(), GetBool(payload, "blocking", true), GetInt(payload, "timeout", 600), GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/history")
                {
                    string id = request.QueryString["instance_id"];
                    await HandleCommand(response, "SCRAPE_HISTORY", null, true, 10, id);
                }
                else if (path == "/api/admin/session/analyze")
                {
                    await HandleCommand(response, "ANALYZE_SESSION", null, true, 5, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/upload_media")
                {
                    if (payload?["data"] == null || payload?["mime"] == null)
                    {
                        SendJson(response, new JsonObject { ["error"] = "Invalid Media Payload" }, 400);
                        return;
                    }
                    await HandleCommand(response, "UPLOAD_MEDIA", payload, true, 30, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/action/new_chat")
                {
                    await HandleCommand(response, "NEW_CHAT", null, true, 30, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/action/toggle_url_context")
                {
                    await HandleCommand(response, "SET_URL_CONTEXT", GetBool(payload, "enabled", true), true, 30, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/action/resolution")
                {
                    await HandleCommand(response, "SET_RESOLUTION", GetString(payload, "level", "Default"), true, 30, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/action/delete")
                {
                    await HandleCommand(response, "DELETE_TURN", GetInt(payload, "index", -1), true, 30, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/action/settings")
                {
                    await HandleCommand(response, "SET_SETTINGS", payload, true, 45, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/action/get_settings")
                {
                    await HandleCommand(response, "GET_SETTINGS", null, true, 60, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/action/get_chat_history")
                {
                    string id = GetString(payload, "instance_id");
                    bool forceSync = GetBool(payload, "force_sync", false);

                    if (!forceSync && OperationCenter.Instance.HasChatHistory(id))
                    {
                        SendJson(response, new JsonObject
                        {
                            ["status"] = "completed",
                            ["output"] = OperationCenter.Instance.GetChatHistory(id).ToJsonString(),
                            ["timestamp"] = DateTime.Now.ToString("o"),
                            ["instance_id"] = id
                        });
                        return;
                    }
                    await HandleCommand(response, "GET_CHAT_HISTORY", null, true, 120, id);
                }
                else if (path == "/api/admin/action/create_shard")
                {
                    await HandleCommand(response, "CREATE_SHARD", null, true, 600, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/action/inject_shard")
                {
                    if (payload?["shard_data"] == null)
                    {
                        SendJson(response, new JsonObject { ["error"] = "Missing shard_data" }, 400);
                        return;
                    }
                    await HandleCommand(response, "INJECT_SHARD", payload["shard_data"]?.DeepClone(), true, 600, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/action/inject_doc")
                {
                    if (payload?["content"] == null)
                    {
                        SendJson(response, new JsonObject { ["error"] = "Missing content" }, 400);
                        return;
                    }
                    await HandleCommand(response, "INJECT_DOC", payload, true, 30, GetString(payload, "instance_id"));
                }
                else if (path == "/api/admin/openai/system_instructions")
                {
                    string instructions = GetString(payload, "instructions", "");
                    OperationCenter.Instance.OpenAISystemInstructions = instructions;
                    SendJson(response, new JsonObject { ["status"] = "completed", ["output"] = "System instructions updated." });
                }
                else if (path.StartsWith("/api/admin/result/"))
                {
                    string taskId = path.Substring("/api/admin/result/".Length);
                    var res = OperationCenter.Instance.GetResult(taskId);
                    if (res != null)
                    {
                        SendJson(response, new JsonObject { ["status"] = "completed", ["result"] = res.DeepClone() });
                    }
                    else
                    {
                        SendJson(response, new JsonObject { ["status"] = "pending" }, 202);
                    }
                }
                else
                {
                    SendJson(response, new JsonObject { ["error"] = "Not Found" }, 404);
                }
            }
            catch (Exception ex)
            {
                Logger.Log($"[ERROR] Request processing failed: {ex.Message}");
                SendJson(response, new JsonObject { ["error"] = "Internal Server Error" }, 500);
            }
        }

        private async Task HandleCommand(HttpListenerResponse response, string commandType, JsonNode payload, bool blocking, int timeout, string instanceId)
        {
            if (string.IsNullOrEmpty(instanceId))
            {
                instanceId = OperationCenter.Instance.GetDefaultInstanceId();
            }

            // --- OPENAI API INTERCEPTION ---
            if (OperationCenter.Instance.IsOpenAIInstance(instanceId))
            {
                if (commandType == "GENERATE")
                {
                    try
                    {
                        OperationCenter.Instance.SetInstanceBusy(instanceId, true);
                        string prompt = payload?.ToString() ?? "";

                        // Auto-detect model if not cached
                        string modelName = OperationCenter.Instance.OpenAIModelName;
                        if (string.IsNullOrEmpty(modelName))
                        {
                            string modelsUrl = OperationCenter.Instance.OpenAIApiUrl.Replace("/chat/completions", "/models");
                            try
                            {
                                var modelReq = new HttpRequestMessage(HttpMethod.Get, modelsUrl);
                                if (!string.IsNullOrEmpty(OperationCenter.Instance.OpenAIApiKey))
                                {
                                    modelReq.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", OperationCenter.Instance.OpenAIApiKey);
                                }
                                var modelRes = await _httpClient.SendAsync(modelReq);
                                if (modelRes.IsSuccessStatusCode)
                                {
                                    string modelStr = await modelRes.Content.ReadAsStringAsync();
                                    var modelJson = JsonNode.Parse(modelStr);
                                    var dataArray = modelJson?["data"] as JsonArray;
                                    if (dataArray != null && dataArray.Count > 0)
                                    {
                                        modelName = dataArray[0]?["id"]?.ToString();
                                        OperationCenter.Instance.OpenAIModelName = modelName;
                                        Logger.Log($"[OPENAI] Auto-detected model: {modelName}");
                                    }
                                }
                            }
                            catch (Exception ex)
                            {
                                Logger.Log($"[OPENAI] Warning: Failed to fetch models from {modelsUrl}. Using 'default'. Error: {ex.Message}");
                            }

                            if (string.IsNullOrEmpty(modelName))
                            {
                                modelName = "default";
                            }
                        }

                        // Add user prompt to history
                        OperationCenter.Instance.AddChatHistory(instanceId, "User", prompt);

                        // Build messages array from history
                        var messagesArray = new JsonArray();

                        if (!string.IsNullOrEmpty(OperationCenter.Instance.OpenAISystemInstructions))
                        {
                            messagesArray.Add(new JsonObject { ["role"] = "system", ["content"] = OperationCenter.Instance.OpenAISystemInstructions });
                        }

                        var history = OperationCenter.Instance.GetChatHistory(instanceId);
                        foreach (var msg in history)
                        {
                            string role = msg["role"]?.ToString().ToLower() == "user" ? "user" : "assistant";
                            messagesArray.Add(new JsonObject { ["role"] = role, ["content"] = msg["text"]?.ToString() });
                        }

                        var openAiRequest = new JsonObject
                        {
                            ["model"] = modelName,
                            ["messages"] = messagesArray
                        };

                        var requestMsg = new HttpRequestMessage(HttpMethod.Post, OperationCenter.Instance.OpenAIApiUrl);
                        if (!string.IsNullOrEmpty(OperationCenter.Instance.OpenAIApiKey))
                        {
                            requestMsg.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", OperationCenter.Instance.OpenAIApiKey);
                        }
                        requestMsg.Content = new StringContent(openAiRequest.ToJsonString(), Encoding.UTF8, "application/json");

                        Logger.Log($"[OPENAI] Sending request to {OperationCenter.Instance.OpenAIApiUrl}");
                        var responseMsg = await _httpClient.SendAsync(requestMsg);
                        string resultStr = await responseMsg.Content.ReadAsStringAsync();

                        OperationCenter.Instance.SetInstanceBusy(instanceId, false);

                        if (!responseMsg.IsSuccessStatusCode)
                        {
                            Logger.Log($"[OPENAI ERROR] {responseMsg.StatusCode}: {resultStr}");
                            SendJson(response, new JsonObject { ["error"] = $"OpenAI API Error: {responseMsg.StatusCode}" }, 500);
                            return;
                        }

                        var jsonResult = JsonNode.Parse(resultStr);
                        string outputText = jsonResult?["choices"]?[0]?["message"]?["content"]?.ToString() ?? resultStr;

                        // Strip out <think>...</think> blocks generated by DeepSeek R1 models
                        outputText = Regex.Replace(outputText, @"<think>.*?</think>\s*", "", RegexOptions.Singleline | RegexOptions.IgnoreCase);

                        // Handle edge cases where the model omits the opening <think> tag but includes the closing </think> tag
                        outputText = Regex.Replace(outputText, @"^.*?<\/think>\s*", "", RegexOptions.Singleline | RegexOptions.IgnoreCase);

                        outputText = outputText.Trim();

                        // Add assistant response to history
                        OperationCenter.Instance.AddChatHistory(instanceId, "Model", outputText);

                        SendJson(response, new JsonObject
                        {
                            ["status"] = "completed",
                            ["output"] = outputText,
                            ["timestamp"] = DateTime.Now.ToString("o"),
                            ["instance_id"] = instanceId
                        }, 200);
                        return;
                    }
                    catch (Exception ex)
                    {
                        OperationCenter.Instance.SetInstanceBusy(instanceId, false);
                        Logger.Log($"[ERROR] OpenAI API call failed: {ex.Message}");
                        SendJson(response, new JsonObject { ["error"] = "OpenAI API Error: " + ex.Message }, 500);
                        return;
                    }
                }
                else if (commandType == "NEW_CHAT")
                {
                    OperationCenter.Instance.ClearChatHistory(instanceId);
                    SendJson(response, new JsonObject { ["status"] = "completed", ["output"] = "Chat reset.", ["timestamp"] = DateTime.Now.ToString("o"), ["instance_id"] = instanceId }, 200);
                    return;
                }
                else if (commandType == "SHUTDOWN")
                {
                    OperationCenter.Instance.ShutdownInstance(instanceId);
                    SendJson(response, new JsonObject { ["status"] = "completed", ["output"] = "Harness shut down.", ["timestamp"] = DateTime.Now.ToString("o"), ["instance_id"] = instanceId }, 200);
                    return;
                }
                else if (commandType == "GET_CHAT_HISTORY")
                {
                    SendJson(response, new JsonObject
                    {
                        ["status"] = "completed",
                        ["output"] = OperationCenter.Instance.GetChatHistory(instanceId).ToJsonString(),
                        ["timestamp"] = DateTime.Now.ToString("o"),
                        ["instance_id"] = instanceId
                    }, 200);
                    return;
                }
                else
                {
                    SendJson(response, new JsonObject { ["error"] = $"Command {commandType} not supported in OpenAI API mode.", ["instance_id"] = instanceId }, 400);
                    return;
                }
            }
            // -------------------------------

            var cmdData = new JsonObject
            {
                ["type"] = commandType,
                ["payload"] = payload?.DeepClone()
            };

            string taskId = OperationCenter.Instance.QueueCommand(cmdData, instanceId);

            if (!blocking)
            {
                SendJson(response, new JsonObject { ["status"] = "queued", ["task_id"] = taskId, ["instance_id"] = instanceId }, 200);
                return;
            }

            DateTime startTime = DateTime.Now;
            while ((DateTime.Now - startTime).TotalSeconds < timeout)
            {
                var result = OperationCenter.Instance.GetResult(taskId);
                if (result != null)
                {
                    SendJson(response, new JsonObject
                    {
                        ["status"] = "completed",
                        ["output"] = result["data"]?["output"]?.ToString(),
                        ["timestamp"] = result["received_at"]?.ToString(),
                        ["instance_id"] = result["instance_id"]?.ToString()
                    }, 200);
                    return;
                }
                await Task.Delay(500);
            }

            // Timeout handling: Cancel the task so it doesn't permanently block the FIFO queue
            OperationCenter.Instance.CancelPendingTask(instanceId, taskId);
            SendJson(response, new JsonObject { ["error"] = "Timeout", ["task_id"] = taskId, ["instance_id"] = instanceId }, 504);
        }

        private void SendJson(HttpListenerResponse response, JsonNode data, int statusCode = 200)
        {
            response.StatusCode = statusCode;
            response.ContentType = "application/json";
            byte[] buffer = Encoding.UTF8.GetBytes(data.ToJsonString());
            response.ContentLength64 = buffer.Length;
            try
            {
                response.OutputStream.Write(buffer, 0, buffer.Length);
            }
            catch { }
            finally
            {
                response.Close();
            }
        }

        private string GetString(JsonNode node, string key, string def = null)
        {
            if (node is JsonObject obj && obj.TryGetPropertyValue(key, out var val) && val != null)
                return val.ToString();
            return def;
        }

        private bool GetBool(JsonNode node, string key, bool def = false)
        {
            if (node is JsonObject obj && obj.TryGetPropertyValue(key, out var val) && val != null)
            {
                try { return val.GetValue<bool>(); } catch { return def; }
            }
            return def;
        }

        private int GetInt(JsonNode node, string key, int def = 0)
        {
            if (node is JsonObject obj && obj.TryGetPropertyValue(key, out var val) && val != null)
            {
                try { return val.GetValue<int>(); } catch { return def; }
            }
            return def;
        }
    }
}