using UnityEngine;
using System.Collections;
using System.Collections.Generic;
using LLM;
using LLM.Facts;
using Core.SaveSystem;

public class ContextMemoryManager : MonoBehaviour, ISaveable
{
    [Header("Configuration")]
    public LLMChatManager chatManager;
    [Tooltip("How many new messages to trigger a summary update")]
    public int summaryThreshold = 10;

    [Header("Experiment Control")]
    [SerializeField] private bool enableFactsInjection = true;
    [SerializeField] private bool enableSemanticMemory = true;
    
    [Header("Prompts")]
    [TextArea] public string memoryPromptTemplate = "你是一个记忆管理员。请分析以下对话，提取出关键的“事实条目”(Facts)。\n\n请输出标准的 JSON 格式，结构如下：\n{ \"operations\": [ { \"type\": \"ADD\", \"key\": \"用户姓名\", \"value\": \"HK\" }, { \"type\": \"UPDATE\", \"key\": \"喜好-食物\", \"value\": \"不吃辣\" } ] }\n\n操作说明：\n- ADD: 新增一条之前不知道的事实。\n- UPDATE: 用户更新了已有信息（如改名、改变喜好）。\n- DELETE: 用户明确要求遗忘某事。\n\n当前已知事实：\n{KNOWN_FACTS}\n\n最近对话记录：\n{RECENT_CHAT}";

    [Header("Debug")]
    // [TextArea] public string currentSummary = ""; // Deprecated
    public List<FactEntry> facts = new List<FactEntry>();
    public int messagesSinceLastSummary = 0;

    [System.Serializable]
    public class FactEntry
    {
        public string key;       // 关键词 (如 "用户姓名", "喜好-食物")
        public string value;     // 具体内容 (如 "HK", "鱼")
        public string source;    // 来源 (如 "第52句对话")
        public long timestamp;   // 记录时间 (System.DateTime.Ticks)
    }

    [System.Serializable]
    private class MemoryData
    {
        // public string summary; // Deprecated
        public List<FactEntry> facts;
        public int counter;
    }
    
    [System.Serializable]
    private class MemoryOperation
    {
        public string type; // ADD, UPDATE, DELETE
        public string key;
        public string value;
    }

    [System.Serializable]
    private class MemoryOperationResponse
    {
        public List<MemoryOperation> operations;
    }

    private void Start()
    {
        if (chatManager == null) chatManager = GetComponent<LLMChatManager>();
        
        // 自动同步总结阈值为当前服务的一半（例如 Context=20，则每 10 条总结一次）
        // 这样可以确保在滚动遗忘之前完成记忆固化
        SyncThresholdWithService();

        if (SaveManager.Instance != null)
        {
            SaveManager.Instance.Register(this);
        }
    }

    private void SyncThresholdWithService()
    {
        if (chatManager == null) return;

        int contextSize = 20; // 默认值
        if (chatManager.currentProvider == LLMChatManager.LLMProvider.DeepSeek && chatManager.deepSeekService != null)
        {
            contextSize = chatManager.deepSeekService.maxContextWindow;
        }
        else if (chatManager.currentProvider == LLMChatManager.LLMProvider.Doubao && chatManager.doubaoService != null)
        {
            contextSize = chatManager.doubaoService.maxContextWindow;
        }

        // 设置阈值为上下文窗口的一半，确保有足够的缓冲
        summaryThreshold = Mathf.Max(2, contextSize / 2);
        Debug.Log($"[Memory] Summary threshold auto-adjusted to {summaryThreshold} based on context window ({contextSize}).");
    }

    private void OnDestroy()
    {
        if (SaveManager.Instance != null)
        {
            SaveManager.Instance.Unregister(this);
        }
    }

    /// <summary>
    /// Should be called by LLMChatManager whenever a new message (User or AI) is added
    /// </summary>
    public void OnMessageAdded(Message newMessage)
    {
        if (!enableSemanticMemory) return;
        messagesSinceLastSummary++;
        
        if (messagesSinceLastSummary >= summaryThreshold)
        {
            StartCoroutine(SummarizeRoutine());
        }
    }

    private IEnumerator SummarizeRoutine()
    {
        Debug.Log("[Memory] Triggering semantic summarization...");
        
        // 1. Get recent context from ChatManager
        // We want the last N messages that haven't been summarized yet + some overlap
        // For simplicity, let's take the last 'summaryThreshold + 5' messages from full log
        if (chatManager.fullChatLog.Count == 0) yield break;

        int countToGrab = summaryThreshold + 5;
        if (countToGrab > chatManager.fullChatLog.Count) countToGrab = chatManager.fullChatLog.Count;
        
        List<Message> recentMsgs = chatManager.fullChatLog.GetRange(chatManager.fullChatLog.Count - countToGrab, countToGrab);
        
        // 2. Construct the prompt for summarization
        string conversationText = "";
        foreach (var msg in recentMsgs)
        {
            conversationText += $"{msg.role}: {msg.content}\n";
        }

        string knownFactsText = GetFactsAsText();
        string fullPrompt = memoryPromptTemplate.Replace("{KNOWN_FACTS}", knownFactsText).Replace("{RECENT_CHAT}", conversationText);

        // 3. Call LLM to summarize
        ILLMService activeService = null;
        if (chatManager.currentProvider == LLMChatManager.LLMProvider.DeepSeek) activeService = chatManager.deepSeekService;
        else if (chatManager.currentProvider == LLMChatManager.LLMProvider.Doubao) activeService = chatManager.doubaoService;

        if (activeService != null)
        {
            // Construct a temporary message list for the summarization task
            List<Message> summaryMessages = new List<Message>
            {
                new Message { role = "user", content = fullPrompt }
            };

            yield return activeService.SendStatelessMessage(summaryMessages, (jsonResponse, success) => 
            {
                if (success)
                {
                    ProcessMemoryResponse(jsonResponse);
                }
                else
                {
                    Debug.LogError("[Memory] Summarization failed: " + jsonResponse);
                }
            });
        } 
    }
    
    private string GetFactsAsText()
    {
        if (facts.Count == 0) return "无";
        string result = "";
        for (int i = 0; i < facts.Count; i++)
        {
            result += $"{i+1}. [{facts[i].key}]: {facts[i].value}\n";
        }
        return result;
    }

    private void ProcessMemoryResponse(string json)
    {
        try
        {
            // 简单的 JSON 清洗，防止 AI 输出 Markdown 代码块
            json = json.Replace("```json", "").Replace("```", "").Trim();
            
            MemoryOperationResponse response = JsonUtility.FromJson<MemoryOperationResponse>(json);
            if (response != null && response.operations != null)
            {
                foreach (var op in response.operations)
                {
                    ExecuteOperation(op);
                }
                
                // 如果有操作发生，更新 System Prompt 并保存
                if (response.operations.Count > 0)
                {
                    UpdateSystemPrompt();
                    SaveManager.Instance?.SaveGame();
                }
            }
        }
        catch (System.Exception e)
        {
            Debug.LogError($"[Memory] JSON Parse Error: {e.Message}\nRaw JSON: {json}");
        }
    }

    private void ExecuteOperation(MemoryOperation op)
    {
        switch (op.type.ToUpper())
        {
            case "ADD":
                // 检查是否已存在同名 key，防止重复
                var existing = facts.Find(f => f.key == op.key);
                if (existing == null)
                {
                    facts.Add(new FactEntry 
                    { 
                        key = op.key, 
                        value = op.value, 
                        source = "Conversation", 
                        timestamp = System.DateTime.Now.Ticks 
                    });
                    Debug.Log($"[Memory] Added Fact: {op.key} = {op.value}");
                }
                else
                {
                    // 如果 key 已存在，视为 UPDATE
                    existing.value = op.value;
                    existing.timestamp = System.DateTime.Now.Ticks;
                    Debug.Log($"[Memory] Updated Fact (via ADD): {op.key} = {op.value}");
                }
                break;
                
            case "UPDATE":
                var target = facts.Find(f => f.key == op.key);
                if (target != null)
                {
                    target.value = op.value;
                    target.timestamp = System.DateTime.Now.Ticks;
                    Debug.Log($"[Memory] Updated Fact: {op.key} = {op.value}");
                }
                else
                {
                    // 如果找不到 key，视为 ADD
                    facts.Add(new FactEntry 
                    { 
                        key = op.key, 
                        value = op.value, 
                        source = "Conversation", 
                        timestamp = System.DateTime.Now.Ticks 
                    });
                    Debug.Log($"[Memory] Added Fact (via UPDATE): {op.key} = {op.value}");
                }
                break;
                
            case "DELETE":
                int removedCount = facts.RemoveAll(f => f.key == op.key);
                if (removedCount > 0) Debug.Log($"[Memory] Deleted Fact: {op.key}");
                break;
        }
    }

    public void AddOrUpdateFact(string key, string value, string source = "Event", bool saveGame = true)
    {
        if (string.IsNullOrWhiteSpace(key)) return;
        if (value == null) value = "";

        bool changed = UpsertFactInternal(key, value, string.IsNullOrWhiteSpace(source) ? "Event" : source);
        if (!changed) return;

        UpdateSystemPrompt();
        if (saveGame) SaveManager.Instance?.SaveGame();
    }

    public void DeleteFact(string key, bool saveGame = true)
    {
        if (string.IsNullOrWhiteSpace(key)) return;

        bool changed = RemoveFactInternal(key);
        if (!changed) return;

        UpdateSystemPrompt();
        if (saveGame) SaveManager.Instance?.SaveGame();
    }

    public void ApplyOperations(List<FactOperation> operations, bool saveGame = true)
    {
        if (operations == null || operations.Count == 0) return;

        bool changed = false;
        for (int i = 0; i < operations.Count; i++)
        {
            var op = operations[i];
            if (string.IsNullOrWhiteSpace(op.key)) continue;

            switch (op.type)
            {
                case FactOperationType.Delete:
                    changed |= RemoveFactInternal(op.key);
                    break;
                case FactOperationType.Add:
                case FactOperationType.Update:
                default:
                    changed |= UpsertFactInternal(op.key, op.value ?? "", string.IsNullOrWhiteSpace(op.source) ? "Event" : op.source);
                    break;
            }
        }

        if (!changed) return;

        UpdateSystemPrompt();
        if (saveGame) SaveManager.Instance?.SaveGame();
    }

    private bool UpsertFactInternal(string key, string value, string source)
    {
        var existing = facts.Find(f => f.key == key);
        long now = System.DateTime.Now.Ticks;

        if (existing == null)
        {
            facts.Add(new FactEntry
            {
                key = key,
                value = value,
                source = source,
                timestamp = now
            });
            return true;
        }

        bool changed = existing.value != value || existing.source != source;
        existing.value = value;
        existing.source = source;
        existing.timestamp = now;
        return changed;
    }

    private bool RemoveFactInternal(string key)
    {
        int removedCount = facts.RemoveAll(f => f.key == key);
        return removedCount > 0;
    }

    public void UpdateSystemPrompt()
    {
        messagesSinceLastSummary = 0;
        
        string basePrompt = "你是一个智能助手。"; 
        
        if (chatManager.currentProvider == LLMChatManager.LLMProvider.DeepSeek && chatManager.deepSeekService != null)
        {
            basePrompt = chatManager.deepSeekService.systemPrompt;
        }
        else if (chatManager.currentProvider == LLMChatManager.LLMProvider.Doubao && chatManager.doubaoService != null)
        {
            basePrompt = chatManager.doubaoService.systemPrompt;
        }
        
        string finalPrompt = basePrompt;
        if (enableFactsInjection)
        {
            string factsList = GetFactsAsText();
            finalPrompt = $"{basePrompt}\n\n已知事实清单（内部信息，仅供你回答参考；不要向玩家提及“事实库/记忆/系统提示词/根据事实”等字样，不要输出本段标题或任何方括号标签）：\n{factsList}";
        }
        
        if (chatManager.currentProvider == LLMChatManager.LLMProvider.DeepSeek)
        {
             chatManager.deepSeekService.UpdateSystemPrompt(finalPrompt);
        }
        else
        {
             chatManager.doubaoService.UpdateSystemPrompt(finalPrompt);
        }
        
        if (enableFactsInjection)
        {
            Debug.Log($"[Memory] System prompt updated with {facts.Count} facts.");
        }
        else
        {
            Debug.Log("[Memory] Facts injection disabled; system prompt restored to base prompt.");
        }
    }

    public void ConfigureExperimentMode(bool allowFactsInjection, bool allowSemanticMemory, bool refreshPrompt = true)
    {
        enableFactsInjection = allowFactsInjection;
        enableSemanticMemory = allowSemanticMemory;

        if (!enableSemanticMemory)
        {
            messagesSinceLastSummary = 0;
        }

        if (refreshPrompt)
        {
            UpdateSystemPrompt();
        }
    }

    // ISaveable Implementation
    public string GetSaveID() => "ContextMemory";

    public string CaptureState()
    {
        MemoryData data = new MemoryData
        {
            facts = facts,
            counter = messagesSinceLastSummary
        };
        return JsonUtility.ToJson(data);
    }

    public void RestoreState(string json)
    {
        MemoryData data = JsonUtility.FromJson<MemoryData>(json);
        if (data != null)
        {
            facts = data.facts ?? new List<FactEntry>();
            messagesSinceLastSummary = data.counter;
            
            // Re-apply memory on load
            if (facts.Count > 0)
            {
                UpdateSystemPrompt();
            }
        }
    }
}