// File: OperationCenter.cs
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json.Nodes;

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

        private readonly Dictionary<string, Queue<JsonObject>> commandQueues = new();
        private readonly Dictionary<string, JsonObject> results = new();
        private readonly Dictionary<string, List<string>> pendingTaskIds = new();
        private readonly Dictionary<string, JsonObject> instances = new();
        private readonly Dictionary<string, JsonArray> chatHistory = new();
        private readonly Dictionary<string, JsonObject> tasks = new();

        // Tracks instances that have been shut down so they aren't recreated by lingering heartbeats
        private readonly HashSet<string> blacklistedInstances = new();

        private readonly object queueLock = new object();

        // OpenAI Configuration
        public bool UseOpenAIMode { get; set; } = false;
        public string OpenAIApiUrl { get; set; } = "http://192.168.12.181:8080/v1/chat/completions";
        public string OpenAIApiKey { get; set; } = "";
        public string OpenAIInstanceId { get; private set; }
        public string OpenAIModelName { get; set; } = "";
        public string OpenAISystemInstructions { get; set; } = "";

        private OperationCenter() { }

        public void ClearAll()
        {
            lock (queueLock)
            {
                commandQueues.Clear();
                results.Clear();
                pendingTaskIds.Clear();
                instances.Clear();
                chatHistory.Clear();
                tasks.Clear();
                blacklistedInstances.Clear();
                UseOpenAIMode = false;
                OpenAIInstanceId = null;
                OpenAIModelName = "";
                OpenAISystemInstructions = "";
            }
        }

        public bool IsOpenAIInstance(string instanceId)
        {
            return !string.IsNullOrEmpty(OpenAIInstanceId) && instanceId == OpenAIInstanceId;
        }

        public string EnableOpenAIMode(string url, string key)
        {
            lock (queueLock)
            {
                UseOpenAIMode = true;
                OpenAIApiUrl = url;
                OpenAIApiKey = key;
                OpenAIModelName = ""; // Reset model name so it fetches on next run

                if (string.IsNullOrEmpty(OpenAIInstanceId))
                {
                    OpenAIInstanceId = $"inst_openai_{DateTimeOffset.UtcNow.ToUnixTimeSeconds()}";
                }

                instances[OpenAIInstanceId] = new JsonObject
                {
                    ["last_seen"] = DateTime.Now.ToString("o"),
                    ["url"] = url,
                    ["turn_count"] = 0,
                    ["is_busy"] = false,
                    ["platform"] = "OpenAI API"
                };

                if (!chatHistory.ContainsKey(OpenAIInstanceId))
                {
                    chatHistory[OpenAIInstanceId] = new JsonArray();
                }

                Logger.Log($"[OP-CENTER] OpenAI API instance registered: {OpenAIInstanceId}");
                return OpenAIInstanceId;
            }
        }

        public void DisableOpenAIMode()
        {
            lock (queueLock)
            {
                UseOpenAIMode = false;
                OpenAIModelName = "";
                if (!string.IsNullOrEmpty(OpenAIInstanceId))
                {
                    instances.Remove(OpenAIInstanceId);
                    Logger.Log($"[OP-CENTER] OpenAI API instance removed: {OpenAIInstanceId}");
                    OpenAIInstanceId = null;
                }
            }
        }

        public void UpdateOpenAIUrl(string url)
        {
            lock (queueLock)
            {
                OpenAIModelName = ""; // Reset model name so it fetches again for the new URL
                if (!string.IsNullOrEmpty(OpenAIInstanceId) && instances.ContainsKey(OpenAIInstanceId))
                {
                    instances[OpenAIInstanceId]["url"] = url;
                }
            }
        }

        private void EnsureInstance(string instanceId)
        {
            bool isNew = false;
            if (!commandQueues.ContainsKey(instanceId))
            {
                commandQueues[instanceId] = new Queue<JsonObject>();
                pendingTaskIds[instanceId] = new List<string>();
                isNew = true;
            }

            if (!blacklistedInstances.Contains(instanceId) && !instances.ContainsKey(instanceId))
            {
                instances[instanceId] = new JsonObject
                {
                    ["last_seen"] = DateTime.Now.ToString("o"),
                    ["url"] = "",
                    ["turn_count"] = 0,
                    ["is_busy"] = false,
                    ["platform"] = "Unknown"
                };

                if (isNew)
                {
                    Logger.Log($"[OP-CENTER] New instance registered: {instanceId}");
                }
            }

            if (!chatHistory.ContainsKey(instanceId))
            {
                chatHistory[instanceId] = new JsonArray();
            }
        }

        public void UpdateInstanceState(string instanceId, JsonObject stateData)
        {
            lock (queueLock)
            {
                if (blacklistedInstances.Contains(instanceId)) return;

                EnsureInstance(instanceId);

                if (instances.ContainsKey(instanceId))
                {
                    instances[instanceId]["last_seen"] = DateTime.Now.ToString("o");

                    if (stateData != null)
                    {
                        if (stateData.ContainsKey("url"))
                            instances[instanceId]["url"] = stateData["url"]?.ToString();
                        if (stateData.ContainsKey("turn_count"))
                            instances[instanceId]["turn_count"] = stateData["turn_count"]?.GetValue<int>() ?? 0;
                        if (stateData.ContainsKey("is_busy"))
                            instances[instanceId]["is_busy"] = stateData["is_busy"]?.GetValue<bool>() ?? false;
                        if (stateData.ContainsKey("platform"))
                            instances[instanceId]["platform"] = stateData["platform"]?.ToString();
                    }
                }
            }
        }

        public JsonObject GetActiveInstances()
        {
            lock (queueLock)
            {
                var result = new JsonObject();
                foreach (var kvp in instances)
                {
                    result[kvp.Key] = kvp.Value.DeepClone();
                }
                return result;
            }
        }

        public string GetDefaultInstanceId()
        {
            lock (queueLock)
            {
                if (instances.Count == 0) return "default";
                return instances.OrderByDescending(x => DateTime.Parse(x.Value["last_seen"].ToString())).First().Key;
            }
        }

        public string GetInstancePlatform(string instanceId)
        {
            lock (queueLock)
            {
                if (instances.TryGetValue(instanceId, out var data))
                {
                    return data?["platform"]?.ToString() ?? "Unknown";
                }
                return "Unknown";
            }
        }

        public string QueueCommand(JsonObject commandData, string instanceId = null)
        {
            lock (queueLock)
            {
                if (string.IsNullOrEmpty(instanceId))
                {
                    instanceId = GetDefaultInstanceId();
                }

                EnsureInstance(instanceId);

                string taskId = $"CMD-{instanceId.Substring(0, Math.Min(8, instanceId.Length))}-{DateTimeOffset.UtcNow.ToUnixTimeSeconds()}";

                var command = new JsonObject
                {
                    ["id"] = taskId,
                    ["command"] = commandData.DeepClone(),
                    ["timestamp"] = DateTime.Now.ToString("o"),
                    ["instance_id"] = instanceId
                };

                tasks[taskId] = commandData.DeepClone().AsObject();

                string cmdType = commandData["type"]?.ToString();
                if (cmdType == "GENERATE")
                {
                    chatHistory[instanceId].Add(new JsonObject
                    {
                        ["role"] = "User",
                        ["text"] = commandData["payload"]?.ToString() ?? ""
                    });
                }

                commandQueues[instanceId].Enqueue(command);
                pendingTaskIds[instanceId].Add(taskId);

                return taskId;
            }
        }

        public JsonObject GetPendingCommand(string instanceId)
        {
            lock (queueLock)
            {
                if (!commandQueues.ContainsKey(instanceId)) return null;

                if (commandQueues[instanceId].Count > 0)
                {
                    if (instances.ContainsKey(instanceId))
                    {
                        instances[instanceId]["is_busy"] = true;
                    }
                    return commandQueues[instanceId].Dequeue();
                }
                return null;
            }
        }

        public bool LogResult(string instanceId, JsonObject resultData)
        {
            lock (queueLock)
            {
                if (instances.ContainsKey(instanceId))
                {
                    instances[instanceId]["is_busy"] = false;
                }

                if (pendingTaskIds.ContainsKey(instanceId) && pendingTaskIds[instanceId].Count > 0)
                {
                    string matchedTaskId = pendingTaskIds[instanceId][0];
                    pendingTaskIds[instanceId].RemoveAt(0);

                    if (tasks.TryGetValue(matchedTaskId, out var cmdData))
                    {
                        string cmdType = cmdData["type"]?.ToString();
                        string output = resultData["output"]?.ToString() ?? "";

                        if (cmdType == "GET_CHAT_HISTORY")
                        {
                            try { chatHistory[instanceId] = JsonNode.Parse(output).AsArray(); } catch { }
                        }
                        else if ((cmdType == "GENERATE" || cmdType == "INJECT_SHARD") && !output.StartsWith("Error"))
                        {
                            chatHistory[instanceId].Add(new JsonObject { ["role"] = "Model", ["text"] = output });
                        }
                        else if (cmdType == "NEW_CHAT")
                        {
                            chatHistory[instanceId].Clear();
                        }
                        else if (cmdType == "DELETE_TURN" && !output.StartsWith("Error"))
                        {
                            try
                            {
                                int idx = cmdData["payload"]?.GetValue<int>() ?? -1;
                                if (idx < 0) idx = chatHistory[instanceId].Count + idx;
                                if (idx >= 0 && idx < chatHistory[instanceId].Count)
                                {
                                    chatHistory[instanceId].RemoveAt(idx);
                                }
                            }
                            catch { }
                        }
                    }

                    results[matchedTaskId] = new JsonObject
                    {
                        ["received_at"] = DateTime.Now.ToString("o"),
                        ["data"] = resultData.DeepClone(),
                        ["instance_id"] = instanceId
                    };

                    Logger.Log($"[OP-CENTER] Task {matchedTaskId} Completed on Instance {instanceId}.");
                    return true;
                }
                else
                {
                    Logger.Log($"[OP-CENTER] Warning: Result received from {instanceId} but no tasks were pending.");
                    return false;
                }
            }
        }

        public JsonObject GetResult(string taskId)
        {
            lock (queueLock)
            {
                if (results.TryGetValue(taskId, out var res))
                {
                    return res.DeepClone().AsObject();
                }
                return null;
            }
        }

        public void CancelPendingTask(string instanceId, string taskId)
        {
            lock (queueLock)
            {
                if (pendingTaskIds.ContainsKey(instanceId))
                {
                    if (pendingTaskIds[instanceId].Remove(taskId))
                    {
                        Logger.Log($"[OP-CENTER] Task {taskId} timed out and was removed from pending queue for {instanceId}.");
                    }
                }
            }
        }

        public void ShutdownInstance(string instanceId)
        {
            lock (queueLock)
            {
                if (IsOpenAIInstance(instanceId))
                {
                    DisableOpenAIMode();
                    return;
                }

                // Queue the shutdown command so the harness receives it on its next poll
                QueueCommand(new JsonObject { ["type"] = "SHUTDOWN" }, instanceId);

                // Blacklist and remove from active instances so it disappears from the UI
                // and cannot be recreated by lingering heartbeats
                blacklistedInstances.Add(instanceId);
                instances.Remove(instanceId);
            }
        }

        public void SetInstanceBusy(string instanceId, bool isBusy)
        {
            lock (queueLock)
            {
                if (instances.ContainsKey(instanceId))
                {
                    instances[instanceId]["is_busy"] = isBusy;
                    instances[instanceId]["last_seen"] = DateTime.Now.ToString("o");
                }
            }
        }

        public void AddChatHistory(string instanceId, string role, string text)
        {
            lock (queueLock)
            {
                if (!chatHistory.ContainsKey(instanceId)) chatHistory[instanceId] = new JsonArray();
                chatHistory[instanceId].Add(new JsonObject { ["role"] = role, ["text"] = text });
            }
        }

        public void ClearChatHistory(string instanceId)
        {
            lock (queueLock)
            {
                if (chatHistory.ContainsKey(instanceId)) chatHistory[instanceId].Clear();
            }
        }

        public bool HasChatHistory(string instanceId)
        {
            lock (queueLock)
            {
                return chatHistory.ContainsKey(instanceId) && chatHistory[instanceId].Count > 0;
            }
        }

        public JsonArray GetChatHistory(string instanceId)
        {
            lock (queueLock)
            {
                if (chatHistory.TryGetValue(instanceId, out var history))
                {
                    return history.DeepClone().AsArray();
                }
                return new JsonArray();
            }
        }
    }
}