terra-llm/Assets/LLM/Editor/Commands/GatherContextCommand.cs

using UnityEditor;
using UnityEngine;
using LLM.Editor.Data;
using LLM.Editor.Helper;
using LLM.Editor.Analysis;
using JetBrains.Annotations;
using System.Collections.Generic;

namespace LLM.Editor.Commands
{
    /// <summary>
    /// The master "tool" command. It receives a batch of data requests from the LLM,
    /// dispatches them to the appropriate Context Providers, and aggregates the results.
    /// </summary>
    [UsedImplicitly]
    public class GatherContextCommand : ICommand
    {
        private readonly GatherContextParams _params;

        public GatherContextCommand(string jsonParams)
        {
            _params = jsonParams?.FromJson<GatherContextParams>();
        }

        public CommandOutcome Execute(CommandContext context)
        {
            if (_params?.requests == null || _params.requests.Count == 0)
            {
                Debug.LogError("[GatherContextCommand] No requests found in parameters.");
                return CommandOutcome.Error;
            }

            var results = new Dictionary<string, object>();

            foreach (var request in _params.requests)
            {
                var targetObject = ResolveIdentifier(request.subjectIdentifier);

                var provider = ContextProviderRegistry.GetProvider(request.dataType);
                if (provider == null)
                {
                    Debug.LogWarning($"[GatherContextCommand] No context provider found for dataType '{request.dataType}'. Skipping request for key '{request.contextKey}'.");
                    results[request.contextKey] = $"Error: No provider for dataType '{request.dataType}'";
                    continue;
                }
                
                // The qualifier is now passed as a JToken, which can be either a string or a complex object.
                // The provider itself will handle the specific type it expects. For most providers, this
                // will just be a string. For GetDataFromPath, it will be a PathRequest object.
                var qualifierString = request.qualifier?.Type == Newtonsoft.Json.Linq.JTokenType.String 
                    ? request.qualifier.ToString() 
                    : request.qualifier?.ToString(Newtonsoft.Json.Formatting.None);

                var contextData = provider.GetContext(targetObject, qualifierString);
                results[request.contextKey] = contextData;
            }

            var aggregatedContextJson = results.ToJson(true);
            context.CurrentSubject = aggregatedContextJson;
            
            Debug.Log($"[GatherContextCommand] Successfully gathered context for {_params.requests.Count} request(s). Pausing for next API call.");
            Debug.Log($"Aggregated Context:\n{aggregatedContextJson}");

            return CommandOutcome.AwaitingNextTurn;
        }

        /// <summary>
        /// Resolves an identifier string to a Unity Object.
        /// The identifier can be an InstanceID, a GUID, or an asset path.
        /// </summary>
        private static Object ResolveIdentifier(string identifier)
        {
            if (string.IsNullOrEmpty(identifier)) return null;

            // Priority 1: Try parsing as an InstanceID for scene objects
            if (int.TryParse(identifier, out var instanceId))
            {
                var obj = EditorUtility.InstanceIDToObject(instanceId);
                if (obj) return obj;
            }

            // Priority 2: Try resolving as a GUID for project assets
            var path = AssetDatabase.GUIDToAssetPath(identifier);
            if (!string.IsNullOrEmpty(path))
            {
                return AssetDatabase.LoadAssetAtPath<Object>(path);
            }

            // Priority 3: Assume it's a direct asset path
            if (System.IO.File.Exists(identifier) || System.IO.Directory.Exists(identifier))
            {
                 return AssetDatabase.LoadAssetAtPath<Object>(identifier);
            }

            Debug.LogWarning($"[GatherContextCommand] Could not resolve identifier '{identifier}' to a Unity Object.");
            return null;
        }
    }
}