using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.Text;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Text;

namespace RhMcp.Router.Codegen;

// Incremental source generator that scans plugin tool source files (passed as
// AdditionalFiles from the router project) and emits a router-side proxy class
// for each [McpServerTool] method. Generated proxies prepend a `slot` argument
// and forward the call to ProxyDispatcher.CallToolAsync.
//
// Parameter type rule: primitive types (string/bool/int/long/double/float and
// their nullables/arrays) pass through unchanged. Anything else — including
// Rhino-specific types like Vector3d and the plugin's record structs — maps to
// an "open object" Dictionary<string, JsonElement>? so the router doesn't have to
// reference RhinoCommon or duplicate types. The plugin deserialises on its end.
//
// Plugin tools take RhinoDoc as their first parameter (auto-injected from DI in
// the plugin's MCP server). That parameter is skipped in the proxy signature.
[Generator]
public class RouterToolGenerator : IIncrementalGenerator
{
    // Router-side CLR type for any non-passthrough (complex/object) tool parameter.
    private const string OpenObjectType =
        "global::System.Collections.Generic.Dictionary<string, global::System.Text.Json.JsonElement>?";

    private static readonly HashSet<string> PassThroughTypes = new()
    {
        "string", "string?",
        "string[]", "string?[]", "string[]?", "string?[]?",
        "bool", "bool?",
        "int", "int?", "int[]",
        "long", "long?",
        "double", "double?",
        "float", "float?",
    };

    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        var toolFiles = context.AdditionalTextsProvider
            .Where(static t => IsToolFile(t.Path));

        var parsed = toolFiles.Select(static (file, ct) =>
        {
            var text = file.GetText(ct);
            if (text is null) return default;
            var tree = CSharpSyntaxTree.ParseText(text, cancellationToken: ct);
            return new ParsedFile(file.Path, tree);
        }).Where(static p => p.Tree is not null);

        var tools = parsed.SelectMany(static (file, _) => ExtractTools(file.Tree!));

        context.RegisterSourceOutput(tools.Collect(), static (spc, all) =>
        {
            // Build a single combined file containing every proxy. Easier to debug than
            // 25 separate .g.cs files and avoids name-collision pitfalls.
            var sb = new StringBuilder();
            sb.AppendLine("// <auto-generated />");
            sb.AppendLine("#nullable enable");
            sb.AppendLine("namespace RhMcp.Router.Tools.Generated;");
            sb.AppendLine();

            foreach (var tool in all)
            {
                EmitProxy(sb, tool);
                sb.AppendLine();
            }

            // Emit a static registration helper Program.cs can call instead of
            // WithToolsFromAssembly() (which uses reflection and is AOT-hostile).
            EmitRegistrar(sb, all);

            spc.AddSource("RouterToolProxies.g.cs", SourceText.From(sb.ToString(), Encoding.UTF8));
        });
    }

    private readonly record struct ParsedFile(string Path, SyntaxTree? Tree);

    private static bool IsToolFile(string path)
    {
        var normalised = path.Replace('\\', '/');
        return normalised.EndsWith(".cs", System.StringComparison.Ordinal)
            && normalised.Contains("/plugin/Tools/");
    }

    private static IEnumerable<ToolInfo> ExtractTools(SyntaxTree tree)
    {
        foreach (var classDecl in tree.GetRoot().DescendantNodes().OfType<ClassDeclarationSyntax>())
        {
            if (!HasAttribute(classDecl.AttributeLists, "McpServerToolType")) continue;
            var className = classDecl.Identifier.ValueText;

            foreach (var method in classDecl.Members.OfType<MethodDeclarationSyntax>())
            {
                var toolAttr = FindAttribute(method.AttributeLists, "McpServerTool");
                if (toolAttr is null) continue;

                // McpServerToolAttribute ctor (plugin Server/Attributes.cs):
                //   (string name, string? title, bool readOnly, bool destructive)
                // Every plugin site writes these positionally, so match by index and
                // let a named override (e.g. `Name = "x"`) win if one is present.
                var toolName = ExtractStringArg(toolAttr, "Name", 0);
                if (toolName is null) continue;

                var toolTitle = ExtractStringArg(toolAttr, "Title", 1);
                var readOnly = ExtractBoolArg(toolAttr, "ReadOnly", 2) ?? false;
                var destructive = ExtractBoolArg(toolAttr, "Destructive", 3) ?? false;

                var descAttr = FindAttribute(method.AttributeLists, "Description");
                var description = descAttr is null ? "" : ExtractStringArg(descAttr, positionalIndex: 0) ?? "";

                var parameters = new List<ParameterInfo>();
                foreach (var p in method.ParameterList.Parameters)
                {
                    var typeText = p.Type?.ToString() ?? "object";
                    // Skip auto-injected, non-MCP-surface parameters. RhinoDoc and the
                    // request CancellationToken are filled by the plugin's SDK, never by
                    // the caller; the proxy supplies its own trailing `ct` regardless, so
                    // letting a plugin `ct` through would also collide with it.
                    if (typeText is "RhinoDoc" or "CancellationToken" or "System.Threading.CancellationToken") continue;

                    var paramName = p.Identifier.ValueText;
                    var paramDescAttr = FindAttribute(p.AttributeLists, "Description");
                    var paramDesc = paramDescAttr is null ? "" : ExtractStringArg(paramDescAttr, positionalIndex: 0) ?? "";
                    var defaultValue = p.Default?.Value?.ToString();
                    var routerType = MapType(typeText);

                    parameters.Add(new ParameterInfo(paramName, routerType, paramDesc, defaultValue));
                }

                yield return new ToolInfo(
                    className, toolName, toolTitle, readOnly, destructive,
                    description, parameters.ToImmutableArray());
            }
        }
    }

    private static string MapType(string typeText)
    {
        // Normalise whitespace inside the syntactic type. Arrays/nullables come through
        // as `string?`, `string[]`, `double?` etc.
        var t = typeText.Replace(" ", "");
        return PassThroughTypes.Contains(t)
            ? t
            : OpenObjectType;
    }

    private static void EmitProxy(StringBuilder sb, ToolInfo tool)
    {
        var className = $"{tool.ClassName}Proxy";

        sb.AppendLine("[global::ModelContextProtocol.Server.McpServerToolType]");
        sb.AppendLine($"public class {className}(global::RhMcp.Router.ProxyDispatcher proxy)");
        sb.AppendLine("{");
        // Always emit ReadOnly + Destructive so every router proxy carries the
        // hint annotations Anthropic's connector-submission policy requires.
        sb.Append($"    [global::ModelContextProtocol.Server.McpServerTool(Name = \"{tool.Name}\"");
        if (tool.Title is not null) sb.Append($", Title = \"{EscapeString(tool.Title)}\"");
        sb.Append($", ReadOnly = {(tool.ReadOnly ? "true" : "false")}");
        sb.Append($", Destructive = {(tool.Destructive ? "true" : "false")}");
        sb.AppendLine(")]");
        sb.AppendLine($"    [global::System.ComponentModel.Description(\"{EscapeString(tool.Description)}\")]");
        sb.AppendLine($"    public global::System.Threading.Tasks.Task<string> InvokeAsync(");

        // Emit original plugin parameters first (required ones may have no default,
        // optional ones carry their plugin-side default through). `slot` follows as
        // an optional trailing arg — must come after any required param per C# rules.
        foreach (var p in tool.Parameters)
        {
            sb.Append($"        [global::System.ComponentModel.Description(\"{EscapeString(p.Description)}\")] {p.Type} {p.Name}");
            if (p.Default is not null) sb.Append($" = {p.Default}");
            sb.AppendLine(",");
        }

        sb.AppendLine("        [global::System.ComponentModel.Description(\"Slot ID returned by spawn_slot. Omit to use the Rhino you're already working in (the one you last used, or one you have open); a new Rhino is auto-spawned only if none is running.\")] string? slot = null,");

        sb.AppendLine("        global::System.Threading.CancellationToken ct = default)");
        sb.AppendLine("    {");

        // Build args as a JsonObject so we never need reflection over an anonymous
        // type — AOT-safe. JsonElement? params parse-through. Array types build a
        // JsonArray element-by-element (JsonValue.Create<T[]> is reflection-based
        // and AOT-hostile). Scalar primitives use JsonValue.Create which has
        // type-specific overloads for string/int/bool/double/float/long.
        sb.AppendLine("        var args = new global::System.Text.Json.Nodes.JsonObject();");
        foreach (var p in tool.Parameters)
        {
            if (p.Type == OpenObjectType)
            {
                // Open object → rebuild a JsonObject from the dictionary entries so the
                // plugin receives a real JSON object (e.g. {"x":..,"y":..,"z":..}), not
                // a stringified one. Each value round-trips through its raw JSON text.
                sb.AppendLine($"        if ({p.Name} is not null)");
                sb.AppendLine("        {");
                sb.AppendLine($"            var __obj_{p.Name} = new global::System.Text.Json.Nodes.JsonObject();");
                sb.AppendLine($"            foreach (var __kv in {p.Name}) __obj_{p.Name}[__kv.Key] = global::System.Text.Json.Nodes.JsonNode.Parse(__kv.Value.GetRawText());");
                sb.AppendLine($"            args[\"{p.Name}\"] = __obj_{p.Name};");
                sb.AppendLine("        }");
            }
            else if (p.Type.Contains("[]"))
            {
                // Array → JsonArray. Works for string[], int[], string[]? etc. The
                // `JsonArray` ctor accepts `params JsonNode?[]`, but we want a loop
                // so per-element implicit conversions kick in.
                var indent = "        ";
                var nullable = p.Type.EndsWith("?");
                if (nullable)
                {
                    sb.AppendLine($"{indent}if ({p.Name} is not null)");
                    sb.AppendLine($"{indent}{{");
                    indent = "            ";
                }
                sb.AppendLine($"{indent}var __arr_{p.Name} = new global::System.Text.Json.Nodes.JsonArray();");
                // Wrap each element with the type-specific JsonValue.Create overload
                // so the resulting JsonNode? selects JsonArray.Add(JsonNode?) — the
                // non-generic, AOT-safe overload. Bare Add(item) routes through
                // generic Add<T>(T) which has RequiresUnreferencedCode.
                sb.AppendLine($"{indent}foreach (var __item in {p.Name}) __arr_{p.Name}.Add((global::System.Text.Json.Nodes.JsonNode?)global::System.Text.Json.Nodes.JsonValue.Create(__item));");
                sb.AppendLine($"{indent}args[\"{p.Name}\"] = __arr_{p.Name};");
                if (nullable)
                {
                    sb.AppendLine("        }");
                }
            }
            else
            {
                sb.AppendLine($"        args[\"{p.Name}\"] = global::System.Text.Json.Nodes.JsonValue.Create({p.Name});");
            }
        }
        // GH2_* tools only work in Rhino WIP (Grasshopper 2 only ships there),
        // so when no slot is passed the router auto-spawns WIP instead of the
        // configured default. Non-GH2 tools pass null and use the configured default.
        if (tool.ClassName.StartsWith("GH2_"))
        {
            sb.AppendLine($"        return proxy.CallToolAsync(slot, \"{tool.Name}\", args, ct, defaultVersionOverride: \"WIP\");");
        }
        else
        {
            sb.AppendLine($"        return proxy.CallToolAsync(slot, \"{tool.Name}\", args, ct);");
        }

        sb.AppendLine("    }");
        sb.AppendLine("}");
    }

    // Static helper Program.cs calls instead of WithToolsFromAssembly().
    // Emits one WithTools<T>() per proxy plus the hand-written SpawnSlotTool.
    // Keeps Program.cs from having to track each new proxy by hand.
    private static void EmitRegistrar(StringBuilder sb, ImmutableArray<ToolInfo> tools)
    {
        sb.AppendLine();
        sb.AppendLine("public static class RouterToolRegistrar");
        sb.AppendLine("{");
        sb.AppendLine("    public static global::Microsoft.Extensions.DependencyInjection.IMcpServerBuilder RegisterAll(global::Microsoft.Extensions.DependencyInjection.IMcpServerBuilder builder, global::System.Text.Json.JsonSerializerOptions? options = null)");
        sb.AppendLine("    {");
        sb.AppendLine("        builder = global::Microsoft.Extensions.DependencyInjection.McpServerBuilderExtensions.WithTools<global::RhMcp.Router.Tools.SpawnSlotTool>(builder, options);");
        foreach (var tool in tools)
        {
            sb.AppendLine($"        builder = global::Microsoft.Extensions.DependencyInjection.McpServerBuilderExtensions.WithTools<global::RhMcp.Router.Tools.Generated.{tool.ClassName}Proxy>(builder, options);");
        }
        sb.AppendLine("        return builder;");
        sb.AppendLine("    }");
        sb.AppendLine("}");
    }

    private static bool HasAttribute(SyntaxList<AttributeListSyntax> lists, string name)
        => FindAttribute(lists, name) is not null;

    private static AttributeSyntax? FindAttribute(SyntaxList<AttributeListSyntax> lists, string name)
    {
        foreach (var list in lists)
        {
            foreach (var attr in list.Attributes)
            {
                var attrName = attr.Name.ToString();
                // Handle `[McpServerTool]`, `[McpServerToolAttribute]`, `[Foo.McpServerTool]`.
                var lastDot = attrName.LastIndexOf('.');
                if (lastDot >= 0) attrName = attrName.Substring(lastDot + 1);
                if (attrName == name || attrName == name + "Attribute") return attr;
            }
        }
        return null;
    }

    /// <summary>
    /// Extract a string argument from an attribute, honouring both named and positional
    /// styles. A named match (`Name = "value"`) always wins if present; otherwise the
    /// argument at <paramref name="positionalIndex"/> (skipping any named args) is read.
    /// Pass <paramref name="namedArg"/> = null to match positional only.
    /// </summary>
    private static string? ExtractStringArg(AttributeSyntax attr, string? namedArg = null, int positionalIndex = -1)
    {
        if (attr.ArgumentList is null) return null;

        if (TryFindArg(attr.ArgumentList, namedArg, positionalIndex, out AttributeArgumentSyntax matched)
            && TryEvaluateConstantString(matched.Expression, out string value))
        {
            return value;
        }
        return null;
    }

    /// <summary>
    /// Fold a compile-time-constant string expression to its text. Handles a bare string
    /// literal and a `+` chain of string literals (e.g. AskUserTool's multi-line
    /// `[Description("a" + "b" + ...)]`), so concatenated descriptions don't emit as "".
    /// </summary>
    private static bool TryEvaluateConstantString(ExpressionSyntax expr, out string value)
    {
        switch (expr)
        {
            case LiteralExpressionSyntax lit when lit.IsKind(SyntaxKind.StringLiteralExpression):
                value = lit.Token.ValueText;
                return true;
            case ParenthesizedExpressionSyntax paren:
                return TryEvaluateConstantString(paren.Expression, out value);
            case BinaryExpressionSyntax bin when bin.IsKind(SyntaxKind.AddExpression)
                && TryEvaluateConstantString(bin.Left, out string left)
                && TryEvaluateConstantString(bin.Right, out string right):
                value = left + right;
                return true;
            default:
                value = "";
                return false;
        }
    }

    private static bool? ExtractBoolArg(AttributeSyntax attr, string namedArg, int positionalIndex = -1)
    {
        if (attr.ArgumentList is null) return null;

        if (TryFindArg(attr.ArgumentList, namedArg, positionalIndex, out AttributeArgumentSyntax matched)
            && matched.Expression is LiteralExpressionSyntax lit)
        {
            if (lit.IsKind(SyntaxKind.TrueLiteralExpression)) return true;
            if (lit.IsKind(SyntaxKind.FalseLiteralExpression)) return false;
        }
        return null;
    }

    /// <summary>
    /// Resolve a single attribute argument by name (preferred) or by positional index.
    /// Named args are excluded from the positional count so `[Foo("a", Bar = "b")]`
    /// still treats "a" as positional 0.
    /// </summary>
    private static bool TryFindArg(AttributeArgumentListSyntax argList, string? namedArg, int positionalIndex, out AttributeArgumentSyntax matched)
    {
        if (namedArg is not null)
        {
            foreach (AttributeArgumentSyntax arg in argList.Arguments)
            {
                if (arg.NameEquals?.Name.Identifier.ValueText == namedArg)
                {
                    matched = arg;
                    return true;
                }
            }
        }

        if (positionalIndex >= 0)
        {
            int seen = 0;
            foreach (AttributeArgumentSyntax arg in argList.Arguments)
            {
                if (arg.NameEquals is not null) continue;
                if (seen == positionalIndex)
                {
                    matched = arg;
                    return true;
                }
                seen++;
            }
        }

        matched = null!;
        return false;
    }

    private static string EscapeString(string s)
    {
        return s
            .Replace("\\", "\\\\")
            .Replace("\"", "\\\"")
            .Replace("\r", "\\r")
            .Replace("\n", "\\n");
    }

    private readonly record struct ToolInfo(
        string ClassName,
        string Name,
        string? Title,
        bool ReadOnly,
        bool Destructive,
        string Description,
        ImmutableArray<ParameterInfo> Parameters);

    private readonly record struct ParameterInfo(
        string Name,
        string Type,
        string Description,
        string? Default);
}