using System;
using System.Collections.Generic;
using System.Text;
using System.Threading.Tasks;
using UnityEngine;
namespace OnDeviceAgent.AgentCore
{

    public sealed class AgentBuilder
    {
        public string Endpoint = "http://localhost:11434";
        public string Model = "gemma4:e2b";
        public string AgentName = "UnityToolCallingAgent";
        public string Instructions = "You are a concise, helpful on-device assistant. Use a tool when it helps.";

        // Restated last in the prompt (just before the user turn) for short-context SLMs. Empty disables it.
        // Default OFF: on very small models (2B), an end-of-prompt meta-instruction can suppress tool calls
        // via recency bias. Set a concise, tool-encouraging value only after verifying it on the target model.
        public string TurnReminder = string.Empty;
        public bool Streaming = true;
        public bool Verbose = false;
        public int MaxRecentTurns = 6;
        public string MemoryDirectoryName = "AgentMemory";
        public int ContextWindowTokens = 8192;
        public int CompactionTriggerPercent = 95;
        public double CharsPerToken = ContextTokenEstimator.DefaultCharsPerToken;
        public Action<string> Log;

        public BuiltinToolDependencies ToolDependencies;

        public Action<ToolRegistry> ConfigureTools;

        public IReadOnlyList<string> Skills;

        public bool LoadProjectSkills = true;

        public string ProjectAssetsDirectory = "Agent";

        // Optional: reuse the RAG E5 embedder to gate project skills by relevance to each turn instead of
        // baking every skill into the prompt. When null, all project skills are baked in (legacy behaviour).
        public ITextEmbedder SkillEmbedder;

        public int MaxActiveSkills = 2;

        // E5 query/passage cosine threshold; multilingual-e5-small scores relevant pairs well above unrelated ones.
        public float SkillRelevanceThreshold = 0.78f;

        public AgentRuntime Build(IUnityMainThreadDispatcher dispatcher)
        {
            var setup = PrepareSetup(dispatcher);

            // GetResult() is safe: FileStreamingAssetReader returns Task.FromResult synchronously. A truly-async
            // IStreamingAssetReader would deadlock here — those consumers must call BuildAsync instead.
            if (UseSkillSelector)
            {
                var ctx = AgentInstructionBuilder
                    .LoadProjectContextAsync(ComposeInstructions(), setup.Assets, ProjectAssetsDirectory, System.Threading.CancellationToken.None)
                    .GetAwaiter().GetResult();
                return CreateRuntime(setup, dispatcher, ctx.BaseInstructions, BuildSkillSelector(ctx.Skills, dispatcher));
            }

            var projectContext = AgentInstructionBuilder
                .BuildProjectContextAsync(ComposeInstructions(), setup.Assets, ProjectAssetsDirectory, System.Threading.CancellationToken.None)
                .GetAwaiter().GetResult();

            return CreateRuntime(setup, dispatcher, projectContext, null);
        }

        public async Task<AgentRuntime> BuildAsync(IUnityMainThreadDispatcher dispatcher)
        {
            var setup = PrepareSetup(dispatcher);

            // Await the project-context read instead of blocking, so a truly-async IStreamingAssetReader
            // (a consumer-supplied reader) does not deadlock the way Build's GetResult() would.
            if (UseSkillSelector)
            {
                var ctx = await AgentInstructionBuilder
                    .LoadProjectContextAsync(ComposeInstructions(), setup.Assets, ProjectAssetsDirectory, System.Threading.CancellationToken.None)
                    .ConfigureAwait(false);
                return CreateRuntime(setup, dispatcher, ctx.BaseInstructions, BuildSkillSelector(ctx.Skills, dispatcher));
            }

            var projectContext = await AgentInstructionBuilder
                .BuildProjectContextAsync(ComposeInstructions(), setup.Assets, ProjectAssetsDirectory, System.Threading.CancellationToken.None)
                .ConfigureAwait(false);

            return CreateRuntime(setup, dispatcher, projectContext, null);
        }

        bool UseSkillSelector => SkillEmbedder != null && SkillEmbedder.Ready;

        // Embeds each project skill as an E5 passage (on the main thread - Build runs there) for the selector.
        SkillSelector BuildSkillSelector(IReadOnlyList<SkillEntry> skills, IUnityMainThreadDispatcher dispatcher)
        {
            if (skills == null || skills.Count == 0)
                return null;
            // Match on the description (short, query-shaped), not the body, which embeds far from a user's question.
            foreach (var skill in skills)
            {
                var matchText = string.IsNullOrWhiteSpace(skill.Description) ? skill.Body : skill.Description;
                skill.Embedding = SkillEmbedder.EmbedText(skill.Name + "\n" + matchText, false);
            }
            return new SkillSelector(skills, SkillEmbedder, dispatcher, MaxActiveSkills, SkillRelevanceThreshold);
        }

        sealed class BuildSetup
        {
            public JsonlConversationStore Store;
            public IStreamingAssetReader Assets;
            public ToolRegistry Registry;
            public AgentRunOptions Options;
            public IAgentTransport Transport;
            public ToolApprovalGate Approval;
            public AgentToolGuardrailController Guardrail;
            public AgentTraceStore Trace;
        }

        BuildSetup PrepareSetup(IUnityMainThreadDispatcher dispatcher)
        {
            if (dispatcher == null) throw new ArgumentNullException(nameof(dispatcher));

            var policy = new AgentMemoryCompactionPolicy { RecentTurnCount = MaxRecentTurns };
            var pathRoot = System.IO.Path.Combine(Application.persistentDataPath, MemoryDirectoryName);
            var store = new JsonlConversationStore(new FilePersistentPathProvider(pathRoot), policy);

            IStreamingAssetReader assets = LoadProjectSkills
                ? new FileStreamingAssetReader(Application.streamingAssetsPath)
                : (IStreamingAssetReader)new EmptyStreamingAssetReader();

            var registry = new ToolRegistry();
            var deps = ToolDependencies ?? new BuiltinToolDependencies();
            BuiltinToolFactory.RegisterAll(registry, true, deps);
            ConfigureTools?.Invoke(registry);

            var options = new AgentRunOptions(
                Endpoint, Model, AgentName, Streaming, true, false, MaxRecentTurns, Verbose,
                ProjectAssetsDirectory)
            {
                ContextWindowTokens = ContextWindowTokens,
                CompactionTriggerPercent = CompactionTriggerPercent,
                CharsPerToken = CharsPerToken,
                TurnReminder = TurnReminder,
            };
    #if UNITY_ANDROID && !UNITY_EDITOR
            // Pick NPU only on SoCs verified to run the QNN export (Fold 7 / SM8750); every other device uses GPU
            // (with CPU fallback in the transport). The model file is matched to this same decision in the
            // provisioner — keep them in sync, since the wrong-SoC NPU export aborts nativeCreateEngine natively.
            // UseNpu is read here, on the Unity main thread, so its JNI SoC query runs safely.
            var llmBackend = LiteRtModelProvisioner.UseNpu ? "NPU" : "GPU";
            IAgentTransport transport = new AndroidLlmTransport(dispatcher, llmBackend, Log);
    #else
            IAgentTransport transport = new OllamaTransport(Endpoint, Model, options.AgentName, Log);
    #endif

            return new BuildSetup
            {
                Store = store,
                Assets = assets,
                Registry = registry,
                Options = options,
                Transport = transport,
                Approval = new ToolApprovalGate(),
                Guardrail = new AgentToolGuardrailController(),
                Trace = new AgentTraceStore(),
            };
        }

        AgentRuntime CreateRuntime(
            BuildSetup setup, IUnityMainThreadDispatcher dispatcher, string projectContext, SkillSelector skillSelector)
        {
            var runtime = new AgentRuntime(
                setup.Transport, setup.Registry, setup.Store, dispatcher, setup.Assets, setup.Options, Log,
                false,
                _ => Task.FromResult(true),
                new MarkdownPropertyAccessor(null),
                setup.Approval, setup.Guardrail, setup.Trace,
                true,
                projectContext);
            runtime.SkillSelector = skillSelector;
            return runtime;
        }

        string ComposeInstructions()
        {
            if (Skills == null || Skills.Count == 0)
                return Instructions;
            var sb = new StringBuilder(Instructions ?? string.Empty);
            foreach (var s in Skills)
            {
                if (string.IsNullOrWhiteSpace(s)) continue;
                sb.AppendLine();
                sb.AppendLine();
                sb.AppendLine("Skill instruction:");
                sb.AppendLine(s.Trim());
            }
            return sb.ToString().Trim();
        }
    }
}