const std = @import("std");
const metadata = @import("metadata.zig");
const ParsedValue = @import("ArgumentType.zig").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),

    /// Owned copy of argv (only if setArgv was called)
    argv: []const [:0]u8,

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

    /// Track if we've done the initial argv scan for help flag
    argv_scanned: 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 .{
            .argv = std.process.argsAlloc(allocator) catch unreachable,
            .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 (only if setArgv was called)
        std.process.argsFree(self.allocator, self.argv);
    }

    /// 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);
    }

    /// Scan argv for help flag without full parsing
    pub fn scanForHelp(self: *ArgumentRegistry) void {
        if (self.argv_scanned) return;
        self.argv_scanned = true;

        const argv = self.argv;
        for (argv[1..]) |arg| {
            // arg is already [:0]const u8, no need to span it
            if (std.mem.eql(u8, arg, "--help") or
                std.mem.eql(u8, arg, "-h"))
            {
                self.help_requested = true;
                return;
            }
        }
    }

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

    /// 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, {});
    }

    /// 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);
    }

    /// Lazy populate: register metadata, parse argv, and populate struct
    /// This is the main entry point for lazy parsing
    pub fn populate(
        self: *ArgumentRegistry,
        comptime T: type,
        comptime module_name: []const u8,
        allocator: std.mem.Allocator,
    ) !T {
        // Register metadata if not already done
        if (!self.isTypeRegistered(T)) {
            try self.registerMetadata(T, module_name);
        }

        // Parse argv on-demand for this type only
        const argv = self.argv;
        try self.parseArgvForType(T, argv);

        // Populate and return the struct
        const parsing = @import("parsing.zig");
        return parsing.populateStruct(T, self, allocator);
    }

    /// Parse argv only for arguments relevant to a specific type
    /// Ignores unknown arguments (they may belong to other modules)
    fn parseArgvForType(self: *ArgumentRegistry, comptime T: type, argv: []const [:0]const u8) !void {
        _ = T; // Type is used implicitly via registered metadata

        const parsing = @import("parsing.zig");
        // Parse argv, ignoring unknown arguments
        try parsing.parseArgv(self, argv);
    }

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

    /// Register metadata for a struct type
    /// INTERNAL USE ONLY: For normal use, call populate() instead
    /// This is only public for testing and internal library use
    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();
    }
};