"""
Conversion utilities from Python dataclass definitions to Investigation DoE config.

The central entry point is :func:`args2config`, which accepts a list of
:class:`~investigation.doeargs.args.ExampleArgs` subclasses and produces the
nested JSON structure that ``ivst init`` writes to ``doe.json``.  All type
information and defaults are read directly from the dataclass — no runtime
execution of user code is required.
"""
from investigation.doeargs.args import ExampleArgs
from typing import List, Dict, Type
from dataclasses import is_dataclass, MISSING


def _get_field_default(field):
    """
    Get the default value for a dataclass field.

    Args:
        field: A dataclass field object

    Returns:
        The default value for the field, or None if no default is available
    """
    if field.default is not MISSING:
        return field.default
    elif field.default_factory is not MISSING:
        return field.default_factory()
    else:
        return None


def extract_args(arg: Type):
    """
    Extract arguments from a dataclass, recursively flattening nested dataclasses.

    Nested dataclass fields are flattened using dot notation.
    For example, if a dataclass has a field 'data' of type Data (another dataclass),
    and Data has a field 'batch_size', the result will include 'data.batch_size'.

    Args:
        arg: A dataclass type to extract arguments from

    Returns:
        A dictionary mapping field names (possibly with dots for nested fields)
        to their type and default value information
    """
    args_dict = {}

    for key, field in arg.__dataclass_fields__.items():
        field_type = field.type

        # Check if this field is a nested dataclass
        if is_dataclass(field_type):
            # Recursively extract nested fields
            nested_dict = extract_args(field_type)  # type: ignore

            # Add nested fields with dot notation
            for nested_key, nested_value in nested_dict.items():
                args_dict[f"{key}.{nested_key}"] = nested_value
        else:
            # Regular field - add it directly
            args_dict[key] = {
                "default": _get_field_default(field),
            }

    return args_dict


def find_dict_cut(args_dicts: List[Dict[str, Dict]]):
    """Partition extracted arg dicts into shared and per-model fields.

    Fields present in *all* arg dicts with identical defaults become
    ``shared``; fields unique to a model or having different defaults remain
    per-model.

    Args:
        args_dicts: List of dicts as returned by :func:`extract_args`, one per
            model variant.  Each variant must have a unique ``name`` default.

    Returns:
        A tuple ``(joint_dict, cutted_dicts)`` where *joint_dict* contains the
        shared fields and *cutted_dicts* is the list of per-model remainders.

    Raises:
        AssertionError: If duplicate ``name`` defaults are found, or if a
            nominally shared field has inconsistent defaults across variants.
    """
    names = set(map(lambda x: x["name"]["default"], args_dicts))
    assert len(names) == len(
        args_dicts
    ), f"The default name for all args must be unqiue, found {len(args_dicts)} args_dicts mismatch with {names}."

    controlled_fields = [
        "default",
    ]
    joint_keys = set(args_dicts[0].keys())
    # input(args_dicts)
    for args_dict in args_dicts:
        joint_keys = joint_keys.intersection(set(args_dict.keys()))
    joint_keys.remove("name")

    joint_dict = {key: args_dicts[0][key] for key in joint_keys}

    mismatchs = {}
    for args_dict in args_dicts:
        model_name = args_dict["name"]["default"]
        for key, fields in joint_dict.items():
            for field in controlled_fields:
                if (fields[field] != args_dict[key][field]):
                    mismatchs[model_name] = mismatchs.get(model_name, []) + [key]

    assert len(mismatchs) == 0, f"Found mismatch fields that are shared by all configures {mismatchs}"

    def cut_dict(args_dict: Dict[str, Dict]):
        return {key: val for key, val in args_dict.items() if key not in joint_keys}

    cutted_dicts = list(map(cut_dict, args_dicts))
    return joint_dict, cutted_dicts


def format_doe(joint_dict: Dict, cuts: List[Dict]):
    """Assemble the final DoE JSON structure from shared and per-model fields.

    Produces the ``{"meta": {...}, "shared": {...}, "<ModelName>": {...}}``
    structure that Investigation commands consume.  Seeds and job names are
    stripped from shared because Investigation manages them automatically.

    Args:
        joint_dict: Shared parameters (output of :func:`find_dict_cut`).
        cuts: Per-model parameter dicts (output of :func:`find_dict_cut`).

    Returns:
        Nested dict ready to be serialised as ``doe.json``.
    """
    for term in ["seed", "job_name"]:
        joint_dict.pop(term, "")
    total_doe = {
        "meta": {
            "num_seeds": {
                "type": str(int),
                "default": 1,
                "value": 1,
            },
            "include": {
                "type": "List[List[Dict[str, Dict]]]",
                "default": [],
                "value": [],
            }
        },
        "shared": {
            key: {
                **val, "range": []
            }
            for key, val in sorted(joint_dict.items(), key=lambda x: x[0])
        },
        **{
            cut["name"]["default"]: {
                key: {
                    **val, "range": []
                }
                for key, val in sorted(cut.items(), key=lambda x: x[0])
            }
            for cut in cuts
        }
    }

    return total_doe


def args2config(args: List[Type[ExampleArgs]]):
    """Convert a list of :class:`~investigation.doeargs.args.ExampleArgs` subclasses to a DoE config.

    This is the primary entry point called by ``ivst init``.  It inspects the
    dataclass fields at import time — not at runtime — so all type annotations
    and default values are available to the Python type checker and IDE
    immediately.

    Args:
        args: One or more :class:`~investigation.doeargs.args.ExampleArgs`
            subclasses representing the model variants to sweep over.

    Returns:
        Nested dict in Investigation DoE JSON format, ready to write to
        ``doe.json``.

    Example:
        .. code-block:: python

            from investigation.doeargs.args2config import args2config
            from src.args import Args
            import json

            config = args2config([Args])
            print(json.dumps(config, indent=2))
    """
    return format_doe(*(find_dict_cut([*map(extract_args, args)])))
