Backlog/lib/zargs/src/ArgumentRegistry.zig

const std = @import("std");
const metadata = @import("metadata");
const ParsedValue = @import("ArgumentType").ParsedValue;

/// Central registry for all command-line arguments
/// Manages argument metadata, tracks modules, and provides lookup functionality
pub const ArgumentRegistry = struct {
    /// Memory allocator
    allocator: std.mem.Allocator,

    /// Map from argument name (e.g., "verbose", "v") to metadata
    /// Both long names and short flags are stored here
    /// Metadata is owned and must be freed
    arguments: std.StringHashMap(metadata.ArgumentMetadata),

    /// Map from argument name to list of modules that registered it
    /// Used for collision detection and help text generation
    modules_by_arg: std.StringHashMap(std.ArrayListUnmanaged([]const u8)),

    /// Set of struct type names that have been registered
    /// Prevents duplicate registration
    registered_types: std.StringHashMap(void),

    /// Cached argv for parsing
    /// Owned by this registry
    argv: ?[]const [:0]const u8 = null,

    /// Whether help was requested (--help or -h)
    help_requested: bool = false,

    /// Parsed values storage
    /// Maps argument name to parsed value
    parsed_values: std.StringHashMap(ParsedValue),

    /// Track which argument keys are allocated (short flags)
    /// Long argument names come from field names (comptime strings) and shouldn't be freed
    allocated_keys: std.StringHashMap(void),

    /// Initialize a new argument registry
    pub fn init(allocator: std.mem.Allocator) ArgumentRegistry {
        return .{
            .allocator = allocator,
            .arguments = std.StringHashMap(metadata.ArgumentMetadata).init(allocator),
            .modules_by_arg = std.StringHashMap(std.ArrayListUnmanaged([]const u8)).init(allocator),
            .registered_types = std.StringHashMap(void).init(allocator),
            .parsed_values = std.StringHashMap(ParsedValue).init(allocator),
            .allocated_keys = std.StringHashMap(void).init(allocator),
        };
    }

    /// Clean up all resources
    pub fn deinit(self: *ArgumentRegistry) void {
        // Clean up modules_by_arg lists
        var modules_iter = self.modules_by_arg.valueIterator();
        while (modules_iter.next()) |list| {
            list.deinit(self.allocator);
        }
        self.modules_by_arg.deinit();

        // Clean up argument keys (only short flags that were allocated)
        var key_iter = self.allocated_keys.keyIterator();
        while (key_iter.next()) |key| {
            self.allocator.free(key.*);
        }
        self.allocated_keys.deinit();
        self.arguments.deinit();
        self.registered_types.deinit();

        // Clean up parsed values
        var values_iter = self.parsed_values.valueIterator();
        while (values_iter.next()) |value| {
            // Free memory for string types
            switch (value.*) {
                .string => |str| self.allocator.free(str),
                .string_list => |list| {
                    for (list) |str| {
                        self.allocator.free(str);
                    }
                    self.allocator.free(list);
                },
                .enum_type => |enum_val| self.allocator.free(enum_val.name),
                else => {},
            }
        }
        self.parsed_values.deinit();

        // Free argv if we own it
        if (self.argv) |args| {
            for (args) |arg| {
                self.allocator.free(arg);
            }
            self.allocator.free(args);
        }
    }

    /// Check if a type has already been registered
    pub fn isTypeRegistered(self: *const ArgumentRegistry, comptime T: type) bool {
        const type_name = @typeName(T);
        return self.registered_types.contains(type_name);
    }

    /// Mark a type as registered
    pub fn markTypeRegistered(self: *ArgumentRegistry, comptime T: type) !void {
        const type_name = @typeName(T);
        try self.registered_types.put(type_name, {});
    }

    /// Check if help was requested
    pub fn isHelpRequested(self: *const ArgumentRegistry) bool {
        return self.help_requested;
    }

    /// Look up argument metadata by name (long or short form)
    pub fn getArgument(self: *const ArgumentRegistry, name: []const u8) ?*const metadata.ArgumentMetadata {
        if (self.arguments.getPtr(name)) |ptr| {
            return ptr;
        }
        return null;
    }

    /// Get list of modules that registered a specific argument
    pub fn getModulesForArg(self: *const ArgumentRegistry, name: []const u8) ?std.ArrayListUnmanaged([]const u8) {
        return self.modules_by_arg.get(name);
    }

    /// Get a parsed value by argument name
    pub fn getParsedValue(self: *const ArgumentRegistry, name: []const u8) ?ParsedValue {
        return self.parsed_values.get(name);
    }

    /// Store a parsed value
    /// Frees the old value if it exists and is a string type
    pub fn storeParsedValue(self: *ArgumentRegistry, name: []const u8, value: ParsedValue) !void {
        // Check if there's an old value we need to free
        if (self.parsed_values.get(name)) |old_value| {
            switch (old_value) {
                .string => |str| self.allocator.free(str),
                .string_list => |list| {
                    for (list) |str| {
                        self.allocator.free(str);
                    }
                    self.allocator.free(list);
                },
                .enum_type => |enum_val| self.allocator.free(enum_val.name),
                else => {},
            }
        }
        try self.parsed_values.put(name, value);
    }

    // ========================================================================
    // Registration Methods
    // ========================================================================

    /// Register metadata for a struct type
    /// Extracts all field metadata and registers each argument
    pub fn registerMetadata(
        self: *ArgumentRegistry,
        comptime T: type,
        comptime module_name: []const u8,
    ) !void {
        // Skip if already registered
        if (self.isTypeRegistered(T)) {
            return;
        }

        // Extract and register each field directly
        const type_info = @typeInfo(T);
        if (type_info != .@"struct") {
            @compileError("registerMetadata requires a struct type");
        }

        inline for (type_info.@"struct".fields) |field| {
            const field_meta = metadata.extractFieldMetadata(T, field);
            try self.registerArgument(&field_meta, module_name);
        }

        // Mark type as registered
        try self.markTypeRegistered(T);
    }

    /// Register a single argument with collision detection
    fn registerArgument(
        self: *ArgumentRegistry,
        arg_meta: *const metadata.ArgumentMetadata,
        module_name: []const u8,
    ) !void {
        // Check if argument already exists (long form)
        const long_exists = self.arguments.getPtr(arg_meta.arg_name);
        if (long_exists) |existing| {
            // Compatible collision: same type
            if (existing.arg_type == arg_meta.arg_type) {
                // Add this module to the list
                try self.addModuleForArg(arg_meta.arg_name, module_name);
                // Don't return yet - we might need to register short form
            } else {
                // Incompatible collision: different types
                return error.IncompatibleArgumentType;
            }
        } else {
            // Register the argument (long form) - store a copy
            try self.arguments.put(arg_meta.arg_name, arg_meta.*);
            try self.addModuleForArg(arg_meta.arg_name, module_name);
        }

        // Register short form if present
        if (arg_meta.short) |short_char| {
            // Create a persistent string for the short key
            const short_key = try self.allocator.alloc(u8, 1);
            short_key[0] = short_char;
            
            // Check for short flag collision
            if (self.arguments.getPtr(short_key)) |existing| {
                // Check if types are compatible
                if (existing.arg_type == arg_meta.arg_type) {
                    // Compatible collision
                    try self.addModuleForArg(short_key, module_name);
                    self.allocator.free(short_key); // Free the temporary key
                    return;
                }
                
                self.allocator.free(short_key); // Free the temporary key
                return error.IncompatibleArgumentType;
            }
            
            // No collision - register the short form (key will be owned by the hash map)
            try self.arguments.put(short_key, arg_meta.*);
            try self.addModuleForArg(short_key, module_name);
            
            // Track that this key was allocated and needs to be freed
            try self.allocated_keys.put(short_key, {});
        }
    }

    /// Add a module to the list for an argument
    fn addModuleForArg(self: *ArgumentRegistry, arg_name: []const u8, module_name: []const u8) !void {
        const entry = try self.modules_by_arg.getOrPut(arg_name);
        if (!entry.found_existing) {
            entry.value_ptr.* = std.ArrayListUnmanaged([]const u8){};
        }
        try entry.value_ptr.append(self.allocator, module_name);
    }

    /// Check if an argument is registered
    pub fn hasArgument(self: *const ArgumentRegistry, name: []const u8) bool {
        return self.arguments.contains(name);
    }

    /// Get the number of registered arguments
    pub fn argumentCount(self: *const ArgumentRegistry) usize {
        return self.arguments.count();
    }
};

// Compile-time validation
comptime {
    // Verify ArgumentRegistry can be created
    const allocator = std.heap.page_allocator;
    _ = ArgumentRegistry.init(allocator);
}