"""Nested config walker and constructor.

Replaces the flat dataclass pipeline. Walks the config tree directly to
register CLI arguments and constructs nested configs from parsed values.
No flattening, no reconstruction, no topological sort.
"""

from __future__ import annotations

import json
import logging
import os
import sys
from dataclasses import MISSING, fields, replace
from itertools import product
from typing import Any, Dict, List, Optional, Set, Tuple, Type, get_args

from vidhi.base_poly_config import BasePolyConfig
from vidhi.cli_parser import HelpNode, VidhiArgumentParser
from vidhi.constants import (
    DICT_KEY_TYPE,
    METADATA_KEY_ALIASES,
    METADATA_KEY_ARGNAME,
    METADATA_KEY_CHOICES,
    METADATA_KEY_HELP,
    METADATA_KEY_POSITIONAL,
    METADATA_KEY_SHORTCUTS,
)
from vidhi.utils import (
    _build_container_info,
    get_all_subclasses,
    get_inner_type,
    is_bool,
    is_container,
    is_dict,
    is_list,
    is_optional,
    is_primitive_type,
    is_subclass,
    load_yaml_config,
    resolve_container_value,
)

logger = logging.getLogger(__name__)


# =============================================================================
# Type key helpers (shared with validation.py)
# =============================================================================


def _get_type_key(type_val: Any) -> str:
    """Get a string type key from a polymorphic type value."""
    if hasattr(type_val, "value"):
        if isinstance(type_val.value, str):
            return type_val.value.lower()
        elif hasattr(type_val, "name"):
            return type_val.name.lower()
        else:
            return str(type_val.value).lower()
    elif hasattr(type_val, "name"):
        return type_val.name.lower()
    else:
        return str(type_val).lower()


def _get_field_type_info(field: Any) -> Tuple[Any, bool]:
    """Extract type information from a dataclass field."""
    if is_optional(field.type):
        return get_inner_type(field.type), True
    return field.type, False


def _get_default_value_for_poly_field(field: Any, field_type: Type[Any]) -> str:
    """Get the default variant type key for a polymorphic field."""
    if field.default_factory is not MISSING:
        default_instance = field.default_factory()
        type_val = default_instance.get_type()
    elif field.default is not MISSING:
        if field.default is None:
            return "None"
        type_val = field.default.get_type()
    else:
        raise ValueError(
            f"Field {field.name} of type {field_type} must have a default or default_factory"
        )
    return _get_type_key(type_val)


def _get_default_variant_class(field: Any, field_type: Type[Any]) -> Optional[Type]:
    """Get the default variant class for a polymorphic field."""
    if field.default_factory is not MISSING:
        return type(field.default_factory())
    elif field.default is not MISSING and field.default is not None:
        return type(field.default)
    return None


# =============================================================================
# ConfigWalker — walk tree, register CLI args
# =============================================================================


class ConfigWalker:
    """Walk a nested config tree and register CLI arguments."""

    def __init__(self, root_class: Type):
        self.root_class = root_class
        # Collected metadata needed for construction and help
        self.poly_children: Dict[str, Dict[str, Type]] = {}
        self.poly_default_types: Dict[str, str] = {}
        self.poly_variant_defaults: Dict[str, Dict[str, Dict[str, Any]]] = {}
        self.poly_field_variants: Dict[str, Dict[str, Any]] = {}
        self.poly_variant_fields: Set[str] = set()
        self.args_with_default_none: Set[str] = set()
        self.list_fields: Dict[str, Any] = {}
        self.container_fields: Dict[str, Any] = {}
        self.nested_group_docs: Dict[str, str] = {}
        self.nested_groups: Dict[str, str] = {}
        self._visited: Set[Type] = set()

    def build_parser(self) -> VidhiArgumentParser:
        """Walk the config tree and build a CLI parser."""
        description = None
        if self.root_class.__doc__:
            description = self.root_class.__doc__.strip().split("\n")[0]

        parser = VidhiArgumentParser(description=description)
        self._walk(parser, self.root_class, prefix="", group="")

        # Build the help tree from the class hierarchy
        parser.help_tree = self._build_help_tree(parser)

        return parser

    def _build_help_tree(self, parser: VidhiArgumentParser) -> HelpNode:
        """Build a help tree by walking the config class hierarchy."""
        root = HelpNode()
        if self.root_class.__doc__:
            root.doc = self.root_class.__doc__.strip().split("\n")[0]
        self._build_help_node(parser, self.root_class, "", root)
        return root

    def _collect_variant_defaults(self, cls: Type, prefix: str) -> Dict[str, Any]:
        """Collect default values for all primitive fields in a dataclass tree.

        For nested DCs, uses the variant class's default_factory instance
        (if any) to get the actual per-variant defaults, rather than
        falling back to the nested DC's own class-level defaults.
        """
        defaults: Dict[str, Any] = {}
        for f in fields(cls):
            flat_name = f"{prefix}.{f.name}" if prefix else f.name
            field_type, _ = _get_field_type_info(f)

            if is_subclass(field_type, BasePolyConfig):
                pass  # Poly defaults handled by _build_poly_help_node
            elif hasattr(field_type, "__dataclass_fields__"):
                # Get the variant's actual default instance for this nested DC
                default_instance = self._get_field_default(f)
                if default_instance is not None and hasattr(
                    default_instance, "__dataclass_fields__"
                ):
                    self._collect_instance_defaults(
                        default_instance, flat_name, defaults
                    )
                else:
                    defaults.update(
                        self._collect_variant_defaults(field_type, flat_name)
                    )
            else:
                defaults[flat_name] = self._get_field_default(f)
        return defaults

    def _collect_instance_defaults(
        self,
        instance: Any,
        prefix: str,
        defaults: Dict[str, Any],
    ) -> None:
        """Collect defaults from a dataclass instance, recursing into nested DCs."""
        for f in fields(type(instance)):
            flat_name = f"{prefix}.{f.name}"
            field_type, _ = _get_field_type_info(f)
            value = getattr(instance, f.name)

            if is_subclass(field_type, BasePolyConfig):
                pass  # Poly defaults handled by _build_poly_help_node
            elif hasattr(field_type, "__dataclass_fields__"):
                if value is not None and hasattr(value, "__dataclass_fields__"):
                    self._collect_instance_defaults(value, flat_name, defaults)
            else:
                defaults[flat_name] = value

    def _build_help_node(
        self,
        parser: VidhiArgumentParser,
        cls: Type,
        prefix: str,
        node: HelpNode,
        default_overrides: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Populate a HelpNode by walking the fields of a dataclass."""
        for f in fields(cls):
            flat_name = f"{prefix}.{f.name}" if prefix else f.name
            field_type, _ = _get_field_type_info(f)

            if is_subclass(field_type, BasePolyConfig):
                child = self._build_poly_help_node(parser, field_type, flat_name)
                node.groups[f.name] = child
            elif hasattr(field_type, "__dataclass_fields__"):
                child = HelpNode()
                if field_type.__doc__:
                    child.doc = field_type.__doc__.strip().split("\n")[0]
                self._build_help_node(
                    parser, field_type, flat_name, child, default_overrides
                )
                node.groups[f.name] = child
            else:
                arg = parser.arguments.get(flat_name)
                if arg is not None:
                    if default_overrides and flat_name in default_overrides:
                        arg = replace(arg, default=default_overrides[flat_name])
                    node.args.append(arg)

    def _build_poly_help_node(
        self,
        parser: VidhiArgumentParser,
        field_type: Type,
        flat_name: str,
    ) -> HelpNode:
        """Build a HelpNode for a polymorphic config field."""
        node = HelpNode()

        # Base class doc
        if field_type.__doc__:
            node.doc = field_type.__doc__.strip().split("\n")[0]

        # Type selector argument
        type_field_name = f"{flat_name}.type"
        node.type_arg = parser.arguments.get(type_field_name)

        # Default variant
        node.default_variant = self.poly_default_types.get(flat_name, "")

        # Build a sub-node for each variant class
        variant_classes = self.poly_children.get(flat_name, {})
        for type_key, variant_cls in sorted(variant_classes.items()):
            variant_node = HelpNode()
            if variant_cls.__doc__:
                variant_node.doc = variant_cls.__doc__.strip().split("\n")[0]

            # Collect correct defaults for THIS variant's class tree
            variant_defaults = self._collect_variant_defaults(variant_cls, flat_name)

            # Walk ALL fields of the variant class (inherited + own)
            for vf in fields(variant_cls):
                vf_flat = f"{flat_name}.{vf.name}"
                vf_type, _ = _get_field_type_info(vf)

                if is_subclass(vf_type, BasePolyConfig):
                    child = self._build_poly_help_node(parser, vf_type, vf_flat)
                    variant_node.groups[vf.name] = child
                elif hasattr(vf_type, "__dataclass_fields__"):
                    child = HelpNode()
                    if vf_type.__doc__:
                        child.doc = vf_type.__doc__.strip().split("\n")[0]
                    self._build_help_node(
                        parser, vf_type, vf_flat, child, variant_defaults
                    )
                    variant_node.groups[vf.name] = child
                else:
                    arg = parser.arguments.get(vf_flat)
                    if arg is not None:
                        if vf_flat in variant_defaults:
                            arg = replace(arg, default=variant_defaults[vf_flat])
                        variant_node.args.append(arg)

            node.variants[type_key] = variant_node

        return node

    def _walk(
        self,
        parser: VidhiArgumentParser,
        cls: Type,
        prefix: str,
        group: str,
    ) -> None:
        """Recursively walk a dataclass and register its fields."""
        for f in fields(cls):
            flat_name = f"{prefix}.{f.name}" if prefix else f.name
            field_type, field_is_optional = _get_field_type_info(f)

            if f.default is None:
                self.args_with_default_none.add(flat_name)

            if is_subclass(field_type, BasePolyConfig):
                self._register_poly(parser, f, field_type, flat_name, group, prefix)
            elif hasattr(field_type, "__dataclass_fields__"):
                # Record nested group docs
                if field_type.__doc__:
                    doc = field_type.__doc__.strip().split("\n")[0]
                    self.nested_group_docs[flat_name] = doc

                self._walk(parser, field_type, flat_name, group or flat_name)
            else:
                self._register_primitive(
                    parser, f, field_type, flat_name, group, prefix
                )

    def _register_poly(
        self,
        parser: VidhiArgumentParser,
        field: Any,
        field_type: Type,
        flat_name: str,
        group: str,
        prefix: str,
    ) -> None:
        """Register a polymorphic config field and all its variants."""
        subclasses = list(get_all_subclasses(field_type))
        if not subclasses:
            raise ValueError(
                f"Polymorphic config error for field '{field.name}': "
                f"{field_type.__name__} has no subclasses."
            )

        # Validate all subclasses implement get_type()
        for sc in subclasses:
            try:
                sc.get_type()
            except NotImplementedError:
                raise ValueError(
                    f"Polymorphic config error for field '{field.name}': "
                    f"{sc.__name__}.get_type() is not implemented."
                ) from None

        default_type_key = _get_default_value_for_poly_field(field, field_type)
        default_variant_cls = _get_default_variant_class(field, field_type)

        # Collect variant info
        valid_types = []
        variant_classes: Dict[str, Type] = {}
        for sc in subclasses:
            type_key = _get_type_key(sc.get_type())
            valid_types.append(type_key)
            variant_classes[type_key] = sc

        self.poly_children[flat_name] = variant_classes
        self.poly_default_types[flat_name] = default_type_key
        self.poly_variant_defaults[flat_name] = {}

        # Register type selector argument
        type_field_name = f"{flat_name}.type"
        type_help = f"Type of {field.name}. Choices: {', '.join(sorted(valid_types))}"

        # Set group and variant descriptions
        for base in (next(iter(subclasses))).__bases__:
            if hasattr(base, "__doc__") and base.__doc__:
                doc = base.__doc__.strip().split("\n")[0]
                parser.set_group_description(group or flat_name, doc)
                break
        for type_key, sc in variant_classes.items():
            if sc.__doc__:
                doc = sc.__doc__.strip().split("\n")[0]
                parser.set_variant_description(group or flat_name, type_key, doc)

        # Determine the group for the type selector
        type_group = group or flat_name
        is_top_level_type = "." not in flat_name
        type_selector_group = "" if is_top_level_type else type_group

        parser.add_argument(
            name=type_field_name,
            cli_name=type_field_name,
            arg_type=str,
            default=default_type_key,
            help_text=type_help,
            required=False,
            is_boolean=False,
            is_list=False,
            variants=None,
            type_field=None,
            group=type_selector_group,
            choices=sorted(valid_types),
            is_type_selector=True,
            poly_base_field=flat_name,
        )

        # Collect all fields from all variants, tracking which variants have each
        all_variant_fields: Dict[str, Dict[str, Any]] = {}

        for sc in subclasses:
            type_key = _get_type_key(sc.get_type())
            is_default_variant = sc is default_variant_cls

            for vf in fields(sc):
                vf_flat = f"{flat_name}.{vf.name}"
                vf_type, _ = _get_field_type_info(vf)

                # Get default
                if vf.default is not MISSING:
                    variant_default = vf.default
                elif vf.default_factory is not MISSING:
                    variant_default = vf.default_factory
                else:
                    variant_default = MISSING

                if vf.name not in all_variant_fields:
                    all_variant_fields[vf.name] = {
                        "type": vf_type,
                        "field": vf,
                        "variants": set(),
                        "variant_defaults": {},
                        "variant_field_types": {},
                    }
                elif is_default_variant:
                    # Prefer default variant's field for recursion
                    all_variant_fields[vf.name]["field"] = vf
                    all_variant_fields[vf.name]["type"] = vf_type

                all_variant_fields[vf.name]["variants"].add(type_key)
                all_variant_fields[vf.name]["variant_defaults"][
                    type_key
                ] = variant_default
                all_variant_fields[vf.name]["variant_field_types"][type_key] = (
                    vf,
                    vf_type,
                )

                # Store variant defaults for reconstruction
                if vf_flat not in self.poly_variant_defaults[flat_name]:
                    self.poly_variant_defaults[flat_name][vf_flat] = {}
                self.poly_variant_defaults[flat_name][vf_flat][
                    type_key
                ] = variant_default

        # Get all variant type keys for determining variant-specific fields
        all_type_keys = set(variant_classes.keys())

        # Register each field
        for field_name, info in all_variant_fields.items():
            vf = info["field"]
            vf_type = info["type"]
            variants = info["variants"]
            variant_defaults = info["variant_defaults"]
            vf_flat = f"{flat_name}.{field_name}"
            poly_group = group or flat_name

            self.poly_variant_fields.add(vf_flat)
            self.poly_field_variants[vf_flat] = {
                "variants": variants,
                "type_field": type_field_name,
            }

            is_nested_poly = is_subclass(vf_type, BasePolyConfig)
            is_nested_dc = (
                hasattr(vf_type, "__dataclass_fields__") and not is_nested_poly
            )

            if is_nested_poly:
                self._register_poly(parser, vf, vf_type, vf_flat, poly_group, flat_name)
            elif is_nested_dc:
                unique_types = {vft for _, vft in info["variant_field_types"].values()}
                if len(unique_types) == 1:
                    # Same type across all variants — walk once (fast path)
                    if vf_type.__doc__:
                        self.nested_group_docs[vf_flat] = vf_type.__doc__.strip().split(
                            "\n"
                        )[0]
                    self._walk(parser, vf_type, vf_flat, poly_group)
                    self._propagate_poly_variant_info(
                        vf_flat,
                        variants,
                        type_field_name,
                    )
                else:
                    # Different types — walk each variant's concrete type
                    self._walk_variant_nested_dc(
                        parser,
                        vf_flat,
                        poly_group,
                        info["variant_field_types"],
                        type_field_name,
                        all_type_keys,
                        default_type_key,
                    )
            else:
                # Primitive/list field
                field_default = self._get_poly_field_default(
                    default_type_key, variant_defaults
                )
                is_variant_specific = variants != all_type_keys

                self._register_poly_primitive(
                    parser,
                    vf,
                    vf_type,
                    vf_flat,
                    poly_group,
                    field_default,
                    is_variant_specific,
                    variants,
                    type_field_name,
                    all_type_keys,
                )

    def _propagate_poly_variant_info(
        self,
        vf_flat: str,
        variants: Set[str],
        type_field_name: str,
    ) -> None:
        """Tag nested fields under a poly variant with their variant info."""
        prefix_check = f"{vf_flat}."
        # Tag all fields we registered under this prefix
        for field_name in list(self.nested_groups.keys()):
            if field_name.startswith(prefix_check):
                self.poly_variant_fields.add(field_name)
                self.poly_field_variants[field_name] = {
                    "variants": variants,
                    "type_field": type_field_name,
                }

    def _walk_variant_nested_dc(
        self,
        parser: VidhiArgumentParser,
        vf_flat: str,
        poly_group: str,
        variant_field_types: Dict[str, Tuple[Any, Type]],
        type_field_name: str,
        all_type_keys: Set[str],
        default_type_key: str,
    ) -> None:
        """Walk a nested dataclass field that has different concrete types per variant.

        When different polymorphic variants use different concrete types for the
        same field name, we must walk each variant's type separately to discover
        all sub-fields and tag them with correct variant sets.
        """
        # Phase 1 — Collect all fields from all variant types
        all_sub_fields: Dict[str, Dict[str, Any]] = {}

        for type_key, (_, concrete_type) in variant_field_types.items():
            is_default = type_key == default_type_key

            # Set group docs from the default variant's type
            if is_default and concrete_type.__doc__:
                self.nested_group_docs[vf_flat] = concrete_type.__doc__.strip().split(
                    "\n"
                )[0]

            for sf in fields(concrete_type):
                sf_flat = f"{vf_flat}.{sf.name}"
                sf_type, _ = _get_field_type_info(sf)

                # Get default
                if sf.default is not MISSING:
                    sub_default = sf.default
                elif sf.default_factory is not MISSING:
                    sub_default = sf.default_factory
                else:
                    sub_default = MISSING

                if sf.name not in all_sub_fields:
                    all_sub_fields[sf.name] = {
                        "type": sf_type,
                        "field": sf,
                        "variants": set(),
                        "variant_defaults": {},
                        "variant_field_types": {},
                    }
                elif is_default:
                    all_sub_fields[sf.name]["field"] = sf
                    all_sub_fields[sf.name]["type"] = sf_type

                all_sub_fields[sf.name]["variants"].add(type_key)
                all_sub_fields[sf.name]["variant_defaults"][type_key] = sub_default
                all_sub_fields[sf.name]["variant_field_types"][type_key] = (sf, sf_type)

        # Populate nested_groups for help grouping
        self.nested_groups[vf_flat] = poly_group

        # Phase 2 — Register each sub-field
        for sub_name, sub_info in all_sub_fields.items():
            sf = sub_info["field"]
            sf_type = sub_info["type"]
            sf_variants = sub_info["variants"]
            sf_variant_defaults = sub_info["variant_defaults"]
            sf_flat = f"{vf_flat}.{sub_name}"

            self.poly_variant_fields.add(sf_flat)
            self.poly_field_variants[sf_flat] = {
                "variants": sf_variants,
                "type_field": type_field_name,
            }

            is_sub_poly = is_subclass(sf_type, BasePolyConfig)
            is_sub_dc = hasattr(sf_type, "__dataclass_fields__") and not is_sub_poly

            if is_sub_poly:
                self._register_poly(parser, sf, sf_type, sf_flat, poly_group, vf_flat)
            elif is_sub_dc:
                unique_sub_types = {
                    vft for _, vft in sub_info["variant_field_types"].values()
                }
                if len(unique_sub_types) == 1:
                    # Same nested DC type across variants — walk once
                    if sf_type.__doc__:
                        self.nested_group_docs[sf_flat] = sf_type.__doc__.strip().split(
                            "\n"
                        )[0]
                    self._walk(parser, sf_type, sf_flat, poly_group)
                    self._propagate_poly_variant_info(
                        sf_flat,
                        sf_variants,
                        type_field_name,
                    )
                else:
                    # Recurse for different nested DC types
                    self._walk_variant_nested_dc(
                        parser,
                        sf_flat,
                        poly_group,
                        sub_info["variant_field_types"],
                        type_field_name,
                        all_type_keys,
                        default_type_key,
                    )
            else:
                # Primitive/list field
                field_default = self._get_poly_field_default(
                    default_type_key, sf_variant_defaults
                )
                is_variant_specific = sf_variants != all_type_keys

                self._register_poly_primitive(
                    parser,
                    sf,
                    sf_type,
                    sf_flat,
                    poly_group,
                    field_default,
                    is_variant_specific,
                    sf_variants,
                    type_field_name,
                    all_type_keys,
                )

    def _get_poly_field_default(
        self, default_type_key: str, variant_defaults: Dict[str, Any]
    ) -> Any:
        """Get the default value from the default variant."""
        if default_type_key in variant_defaults:
            val = variant_defaults[default_type_key]
            if val is not MISSING:
                if callable(val) and not isinstance(val, type):
                    return val()
                return val
        # Fallback to any variant with a default
        for val in variant_defaults.values():
            if val is not MISSING:
                if callable(val) and not isinstance(val, type):
                    return val()
                return val
        return None

    def _register_poly_primitive(
        self,
        parser: VidhiArgumentParser,
        vf: Any,
        vf_type: Any,
        vf_flat: str,
        group: str,
        default: Any,
        is_variant_specific: bool,
        variants: Set[str],
        type_field: str,
        all_type_keys: Set[str],
    ) -> None:
        """Register a primitive field from a polymorphic variant."""
        from enum import Enum as EnumType

        metadata = vf.metadata
        help_text = metadata.get(METADATA_KEY_HELP, "")
        custom_argname = metadata.get(METADATA_KEY_ARGNAME)
        aliases = metadata.get(METADATA_KEY_ALIASES)
        metadata_choices = metadata.get(METADATA_KEY_CHOICES)
        shortcuts = metadata.get(METADATA_KEY_SHORTCUTS)
        positional = metadata.get(METADATA_KEY_POSITIONAL, False)

        cli_name = custom_argname if custom_argname else vf_flat

        is_boolean = is_bool(vf_type)
        is_list_field = is_list(vf_type)
        actual_type = vf_type

        choices = None
        if isinstance(vf_type, type) and issubclass(vf_type, EnumType):
            choices = [m.value for m in vf_type]
        if choices is None and metadata_choices:
            choices = metadata_choices

        if is_list_field:
            inner_type = get_args(vf_type)[0]
            if is_primitive_type(inner_type):
                actual_type = inner_type
            else:
                actual_type = json.loads
            self.list_fields[vf_flat] = get_args(vf_type)[0]
        elif is_dict(vf_type):
            actual_type = json.loads

        # Container metadata
        container_info = None
        is_container_field = False
        if is_container(vf_type):
            container_info = _build_container_info(vf_type)
            self.container_fields[vf_flat] = container_info
            is_container_field = True

        # Wrap default in variant dict for variant-specific fields
        display_default = default
        if is_variant_specific and default is not None:
            display_default = {v: default for v in variants}

        parser.add_argument(
            name=vf_flat,
            cli_name=cli_name,
            arg_type=actual_type,
            default=display_default,
            help_text=help_text,
            required=False,
            is_boolean=is_boolean,
            is_list=is_list_field,
            variants=variants if is_variant_specific else None,
            type_field=type_field if is_variant_specific else None,
            group=group,
            choices=choices,
            aliases=aliases,
            shortcuts=shortcuts,
            positional=positional,
            is_container=is_container_field,
            container_info=container_info,
        )

        self.nested_groups[vf_flat] = group

    def _register_primitive(
        self,
        parser: VidhiArgumentParser,
        f: Any,
        field_type: Any,
        flat_name: str,
        group: str,
        prefix: str,
    ) -> None:
        """Register a primitive (non-dataclass, non-poly) field."""
        from enum import Enum as EnumType

        metadata = f.metadata
        help_text = metadata.get(METADATA_KEY_HELP, "")
        custom_argname = metadata.get(METADATA_KEY_ARGNAME)
        aliases = metadata.get(METADATA_KEY_ALIASES)
        metadata_choices = metadata.get(METADATA_KEY_CHOICES)
        shortcuts = metadata.get(METADATA_KEY_SHORTCUTS)
        positional = metadata.get(METADATA_KEY_POSITIONAL, False)

        cli_name = custom_argname if custom_argname else flat_name

        # Get default
        default = self._get_field_default(f)

        is_boolean = is_bool(field_type)
        is_list_field = is_list(field_type)
        actual_type = field_type

        choices = None
        if isinstance(field_type, type) and issubclass(field_type, EnumType):
            choices = [m.value for m in field_type]
        if choices is None and metadata_choices:
            choices = metadata_choices

        if is_list_field:
            inner_type = get_args(field_type)
            if inner_type:
                inner = inner_type[0]
                if is_primitive_type(inner):
                    actual_type = inner
                else:
                    actual_type = json.loads
                self.list_fields[flat_name] = inner
        elif is_dict(field_type):
            actual_type = json.loads

        # Container metadata
        container_info = None
        is_container_field = False
        if is_container(field_type):
            container_info = _build_container_info(field_type)
            self.container_fields[flat_name] = container_info
            is_container_field = True

        # Determine argument group
        arg_group = group
        if not arg_group and "." in flat_name:
            arg_group = flat_name.split(".")[0]

        parser.add_argument(
            name=flat_name,
            cli_name=cli_name,
            arg_type=actual_type,
            default=default,
            help_text=help_text,
            required=False,
            is_boolean=is_boolean,
            is_list=is_list_field,
            variants=None,
            type_field=None,
            group=arg_group,
            choices=choices,
            aliases=aliases,
            shortcuts=shortcuts,
            positional=positional,
            is_container=is_container_field,
            container_info=container_info,
        )

        # Track nested group membership
        if prefix:
            self.nested_groups[flat_name] = prefix

    def _get_field_default(self, f: Any) -> Any:
        """Get the default value for a field."""
        if f.default is not MISSING:
            value = f.default
            if callable(value) and not isinstance(value, type):
                return value()
            return value
        elif f.default_factory is not MISSING:
            return f.default_factory()
        return None


# =============================================================================
# Flat-to-nested dict conversion
# =============================================================================


def _flat_to_nested(flat_dict: Dict[str, Any]) -> Dict[str, Any]:
    """Convert a flat dict with dot-separated keys to a nested dict.

    Example:
        {"a.b": 1, "a.c.type": "x"} -> {"a": {"b": 1, "c": {"type": "x"}}}
    """
    result: Dict[str, Any] = {}
    for key, value in flat_dict.items():
        parts = key.split(".")
        current = result
        for part in parts[:-1]:
            if part not in current:
                current[part] = {}
            elif not isinstance(current[part], dict):
                current[part] = {}
            current = current[part]
        current[parts[-1]] = value
    return result


def _deep_merge(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
    """Recursively merge override into base. Override wins on conflicts."""
    result = dict(base)
    for key, value in override.items():
        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
            result[key] = _deep_merge(result[key], value)
        else:
            result[key] = value
    return result


# =============================================================================
# CLI integration — config file loading + sweep support
# =============================================================================


def parse_cli_sweep_nested(
    config_class: Type,
    walker: ConfigWalker,
    parser: VidhiArgumentParser,
) -> list:
    """Parse CLI args and create config instances using the nested approach."""
    from vidhi.utils import create_class_from_dict

    # Handle built-in flags
    _handle_builtin_flags(config_class, walker)

    # Handle --config flag
    config_file = _extract_config_file_arg(sys.argv)

    parsed_args, provided_arg_names = parser.parse_args()

    # Track CLI-provided args
    cli_provided_args = {
        k: v for k, v in parsed_args.items() if k in provided_arg_names
    }

    # Process container and list fields in CLI args
    parsed_args = _process_iterable_args(
        parsed_args, walker.list_fields, walker.container_fields
    )
    cli_provided_args = _process_iterable_args(
        cli_provided_args, walker.list_fields, walker.container_fields
    )

    # Validate variant fields (skip when config file may override types)
    if not config_file:
        _validate_variant_fields(walker, [parsed_args], cli_provided_args)

    # Drop default-variant fields when a non-default variant is selected.
    # The walker only walks the default variant's config tree, so parsed_args
    # contains defaults from that tree. When a different variant is selected,
    # those defaults may reference fields that don't exist on the selected
    # variant's concrete classes. Removing them lets create_class_from_dict
    # use the actual class defaults instead.
    parsed_args = _filter_non_selected_variant_defaults(
        parsed_args, provided_arg_names, walker
    )

    # Convert flat CLI args (dot-separated) to nested dict
    cli_nested = _flat_to_nested(parsed_args)
    cli_provided_nested = _flat_to_nested(cli_provided_args)

    # Load config from --config flag
    loaded_configs: List[Dict[str, Any]] = [{}]
    if config_file:
        file_config = load_yaml_config(config_file)
        loaded_configs = _explode_config(walker, file_config)

    # Create config instances from combinations
    results = []
    for file_dict in loaded_configs:
        # Process container and list fields in file config
        file_dict = _process_iterable_args(
            file_dict, walker.list_fields, walker.container_fields
        )

        # Start with CLI defaults, overlay file config, then CLI overrides
        merged = _deep_merge(cli_nested, file_dict)
        # Re-apply CLI-provided args on top (they win over file config)
        merged = _deep_merge(merged, cli_provided_nested)

        config = create_class_from_dict(config_class, merged)
        results.append(config)

    return results


def with_cli_overrides_nested(
    config: Any,
    config_class: Type,
    walker: ConfigWalker,
    parser: VidhiArgumentParser,
) -> Any:
    """Override a config with CLI arguments using the nested approach."""
    from vidhi.utils import create_class_from_dict, dataclass_to_dict

    base_dict = dataclass_to_dict(config)

    argv_copy = sys.argv.copy()
    config_file = _extract_config_file_arg(argv_copy)

    _, cli_provided = parser.parse_args(argv_copy[1:])

    if not cli_provided and not config_file:
        return config

    # Full parse to get CLI config
    configs = parse_cli_sweep_nested(config_class, walker, parser)
    if len(configs) > 1:
        raise ValueError(
            f"CLI arguments produced {len(configs)} configs. "
            "with_cli_overrides() only supports single configs."
        )

    cli_config = configs[0]
    cli_dict = dataclass_to_dict(cli_config)

    provided_fields = set()
    for name in cli_provided:
        parts = tuple(name.split("."))
        provided_fields.add(parts)

    # YAML file values also override the base config
    if config_file:
        file_config = load_yaml_config(config_file)
        _collect_dict_paths(file_config, (), provided_fields)

    merged_dict = _merge_with_provided(base_dict, cli_dict, provided_fields)
    return create_class_from_dict(config_class, merged_dict)


def _collect_dict_paths(
    d: Dict[str, Any],
    path: tuple,
    result: Set[tuple],
) -> None:
    """Collect all leaf key paths from a nested dict into result."""
    for key, value in d.items():
        current = path + (key,)
        if isinstance(value, dict):
            _collect_dict_paths(value, current, result)
        else:
            result.add(current)


def _merge_with_provided(
    base: Dict[str, Any],
    cli: Dict[str, Any],
    provided: set,
    path: tuple = (),
) -> Dict[str, Any]:
    """Merge CLI values into base, but only for provided fields."""
    result = {}
    all_keys = set(base.keys()) | set(cli.keys())

    for key in all_keys:
        current_path = path + (key,)
        base_val = base.get(key)
        cli_val = cli.get(key)

        field_provided = any(p[: len(current_path)] == current_path for p in provided)

        if isinstance(base_val, dict) and isinstance(cli_val, dict):
            result[key] = _merge_with_provided(
                base_val, cli_val, provided, current_path
            )
        elif field_provided:
            result[key] = cli_val if cli_val is not None else base_val
        else:
            result[key] = base_val if key in base else cli_val

    return result


# =============================================================================
# Config explosion (sweep expansion)
# =============================================================================


def _explode_config(
    walker: ConfigWalker,
    config_dict: Dict[str, Any],
) -> List[Dict[str, Any]]:
    """Expand a config dict into nested dicts, handling sweeps."""
    # Handle _list wrapper from load_yaml_config
    if (
        isinstance(config_dict, dict)
        and len(config_dict) == 1
        and "_list" in config_dict
    ):
        list_data = config_dict["_list"]
        if not isinstance(list_data, list):
            return [{}]
        all_exploded = []
        for item in list_data:
            if isinstance(item, dict):
                exploded = _explode_recursive(walker, item, "")
                all_exploded.extend(exploded)
        return all_exploded if all_exploded else [{}]

    return _explode_recursive(walker, config_dict, "")


def _explode_recursive(
    walker: ConfigWalker,
    config_dict: Dict[str, Any],
    prefix: str,
) -> List[Dict[str, Any]]:
    """Recursively explode config dict into combinations (nested dicts)."""
    list_keys = []
    list_values = []
    non_list_items = {}
    dict_items = {}

    for key, value in config_dict.items():
        prefixed_key = f"{prefix}.{key}" if prefix else key

        if isinstance(value, list) and len(value) > 0:
            # Check if field is declared as a literal list
            if prefixed_key in walker.list_fields:
                non_list_items[key] = value
            else:
                list_keys.append(key)
                list_values.append(value)
        elif isinstance(value, dict):
            # Check if field is a literal dict
            if prefixed_key in walker.container_fields:
                non_list_items[key] = value
            else:
                dict_items[key] = value
        else:
            non_list_items[key] = value

    # Process dict items
    dict_combinations = [{}]
    for key, nested_dict in dict_items.items():
        nested_prefix = f"{prefix}.{key}" if prefix else key
        exploded = _explode_recursive(walker, nested_dict, nested_prefix)
        new_combos = []
        for base in dict_combinations:
            for nested in exploded:
                combo = base.copy()
                combo[key] = nested
                new_combos.append(combo)
        dict_combinations = new_combos

    if not list_keys:
        result = []
        for dc in dict_combinations:
            combined = non_list_items.copy()
            combined.update(dc)
            result.append(combined)
        return result

    # Handle list expansion
    processed_list_values = []
    for key, values in zip(list_keys, list_values):
        if values and isinstance(values[0], dict):
            nested_prefix = f"{prefix}.{key}" if prefix else key
            exploded = []
            for v in values:
                exploded.extend(_explode_recursive(walker, v, nested_prefix))
            processed_list_values.append(exploded)
        else:
            processed_list_values.append(values)

    result = []
    for combination in product(*processed_list_values):
        for dc in dict_combinations:
            new_config = non_list_items.copy()
            new_config.update(dc)
            for key, value in zip(list_keys, combination):
                new_config[key] = value
            result.append(new_config)

    return result


# =============================================================================
# Shared utilities
# =============================================================================


def _extract_config_file_arg(argv: List[str]) -> Optional[str]:
    """Extract and remove --config argument from argv."""
    config_file = None
    i = 0
    while i < len(argv):
        if argv[i] == "--config":
            if i + 1 < len(argv) and not argv[i + 1].startswith("--"):
                config_file = argv[i + 1]
                del argv[i : i + 2]
            else:
                print("Error: --config requires a file path", file=sys.stderr)
                sys.exit(1)
        elif argv[i].startswith("--config="):
            config_file = argv[i].split("=", 1)[1]
            del argv[i]
        else:
            i += 1
    return config_file


def _handle_builtin_flags(config_class: Type, walker: ConfigWalker) -> None:
    """Handle --install-shell-completions and --export-json-schema."""
    for i, arg in enumerate(sys.argv):
        if arg == "--install-shell-completions":
            shell_arg = None
            if i + 1 < len(sys.argv) and sys.argv[i + 1] in ("bash", "zsh", "fish"):
                shell_arg = sys.argv[i + 1]
            _handle_install_completions(config_class, walker, shell_arg)
        elif arg.startswith("--install-shell-completions="):
            shell_arg = arg.split("=", 1)[1]
            _handle_install_completions(config_class, walker, shell_arg)

    if "--export-json-schema" in sys.argv:
        _handle_export_schema(config_class)


def _handle_install_completions(
    config_class: Type, walker: ConfigWalker, shell_override: Optional[str] = None
) -> None:
    """Handle --install-shell-completions flag."""
    from vidhi.cli_completion import detect_shell, install_completion

    command_name = os.path.basename(sys.argv[0])
    if command_name.endswith(".py"):
        command_name = command_name[:-3]

    if shell_override:
        if shell_override not in ("bash", "zsh", "fish"):
            print(f"Unknown shell: {shell_override}")
            print("Supported shells: bash, zsh, fish")
            sys.exit(1)
        shell = shell_override
    else:
        shell = detect_shell()
        if shell is None:
            print("Could not detect shell from $SHELL environment variable.")
            sys.exit(1)

    print(f"Installing {shell} completions for '{command_name}'...")
    try:
        result = install_completion(config_class, command_name, shell, walker=walker)
        print(result)
        sys.exit(0)
    except Exception as e:
        print(f"Error installing completions: {e}")
        sys.exit(1)


def _handle_export_schema(config_class: Type) -> None:
    """Handle --export-json-schema flag."""
    from vidhi.schema import ConfigSchema

    try:
        idx = sys.argv.index("--export-json-schema")
        if idx + 1 >= len(sys.argv) or sys.argv[idx + 1].startswith("--"):
            output_path = "config.schema.json"
        else:
            output_path = sys.argv[idx + 1]
    except ValueError:
        output_path = "config.schema.json"

    output_path = os.path.realpath(output_path)

    try:
        schema = ConfigSchema(config_class)
        schema.export_json_schema(output_path)
        print(f"JSON Schema exported to: {output_path}")
        sys.exit(0)
    except Exception as e:
        print(f"Error exporting schema: {e}")
        sys.exit(1)


def _filter_non_selected_variant_defaults(
    parsed_args: Dict[str, Any],
    provided_arg_names: set,
    walker: ConfigWalker,
) -> Dict[str, Any]:
    """Remove fields that don't belong to the selected variant.

    For each polymorphic type selector, determine the selected variant
    (explicit or default) and drop non-provided fields that don't belong
    to that variant. This prevents cross-variant field leakage where
    fields from one variant contaminate another.
    """
    # Build mapping: type_field -> selected variant key
    selected_variants: Dict[str, str] = {}
    for type_prefix, default_type in walker.poly_default_types.items():
        type_key = f"{type_prefix}.type"
        selected = parsed_args.get(type_key, default_type)
        if selected is not None:
            selected_variants[type_key] = str(selected).lower()

    if not selected_variants:
        return parsed_args

    # Drop fields that don't belong to the selected variant
    to_drop: Set[str] = set()
    for field_name, info in walker.poly_field_variants.items():
        type_field = info["type_field"]
        valid_variants = info["variants"]

        selected = selected_variants.get(type_field)
        if selected is None:
            continue

        # If the selected variant is not in this field's valid variants,
        # and the user didn't explicitly provide it, drop it
        if selected not in valid_variants and field_name not in provided_arg_names:
            to_drop.add(field_name)

    return {k: v for k, v in parsed_args.items() if k not in to_drop}


def _validate_variant_fields(
    walker: ConfigWalker,
    all_values: List[Dict[str, Any]],
    cli_provided_args: Dict[str, Any],
) -> None:
    """Validate that CLI-provided variant fields match the final type."""
    for values in all_values:
        for field_name in cli_provided_args:
            poly_info = walker.poly_field_variants.get(field_name)
            if not poly_info:
                continue

            valid_variants = poly_info.get("variants", set())
            type_field = poly_info.get("type_field")
            if not valid_variants or not type_field:
                continue

            final_type = values.get(type_field)
            if final_type is None:
                continue

            if final_type.lower() not in valid_variants:
                valid_list = ", ".join(sorted(valid_variants))
                print(
                    f"Error: --{field_name} is only valid for {valid_list}, "
                    f"but {type_field}={final_type}",
                    file=sys.stderr,
                )
                sys.exit(1)


def _process_iterable_args(
    arg_map: Dict[str, Any],
    list_fields: Dict[str, Any],
    container_fields: Dict[str, Any],
) -> Dict[str, Any]:
    """Process container and list fields by converting raw values to typed instances."""
    result = {}
    for arg_name, arg_value in arg_map.items():
        if arg_name in container_fields and isinstance(arg_value, (dict, list)):
            result[arg_name] = resolve_container_value(
                arg_value, container_fields[arg_name]
            )
            continue

        if arg_name in list_fields:
            if arg_value is None:
                result[arg_name] = arg_value
                continue
            target_type = list_fields[arg_name]
            return_iterable = []
            if isinstance(target_type, type) and issubclass(
                target_type, BasePolyConfig
            ):
                subclasses = get_all_subclasses(target_type)
                for raw_value in arg_value:
                    assert (
                        DICT_KEY_TYPE in raw_value
                    ), f"Each raw value in an iterable of BasePolyConfigs must contain a '{DICT_KEY_TYPE}' key."
                    is_match = False
                    for subclass in subclasses:
                        if (
                            subclass.get_type().name.upper()
                            == raw_value[DICT_KEY_TYPE].upper()
                        ):
                            subclass_kwargs = {
                                k: v for k, v in raw_value.items() if k != DICT_KEY_TYPE
                            }
                            return_iterable.append(subclass(**subclass_kwargs))
                            is_match = True
                            break
                    assert is_match, f"No class found for type '{raw_value['type']}'"
            elif hasattr(target_type, "__dataclass_fields__"):
                for raw_value in arg_value:
                    return_iterable.append(target_type(**raw_value))
            elif isinstance(target_type, type):
                for raw_value in arg_value:
                    return_iterable.append(target_type(raw_value))
            else:
                raise ValueError(f"Unsupported target type: {target_type}")
            result[arg_name] = return_iterable
        else:
            result[arg_name] = arg_value
    return result
