trestle.core.generators

def generate_sample_model(
    model: Union[Type[TG], List[TG], Dict[str, TG]], include_optional: bool = False, depth: int = -1
) -> TG:
    """Given a model class, generate an object of that class with sample values.

    Can generate optional variables with an enabled flag. Any array objects will have a single entry injected into it.

    Note: Trestle generate will not activate recursive loops irrespective of the depth flag.

    Args:
        model: The model type provided. Typically for a user as an OscalBaseModel Subclass.
        include_optional: Whether or not to generate optional fields.
        depth: Depth of the tree at which optional fields are generated. Negative values (default) removes the limit.

    Returns:
        The generated instance with a pro-forma values filled out as best as possible.
    """
    effective_optional = include_optional and not depth == 0

    model_type = model
    # This block normalizes model type down to
    if utils.is_collection_field_type(model):
        model_type = utils.get_origin(model)
        model = utils.get_inner_type(model)

    # Handle Union types at the top level (e.g., when model is Union[Parameter1, Parameter2])
    # This can happen when get_inner_type returns a Union from a list[Union[...]]
    origin = utils.get_origin(model)
    is_union = origin == Union or str(origin) == "<class 'types.UnionType'>"
    if is_union:
        union_args = typing.get_args(model)
        # Find first non-None OscalBaseModel type in the union
        for arg in union_args:
            if arg is not type(None) and safe_is_sub(arg, OscalBaseModel):
                model = arg
                break
        else:
            # If no OscalBaseModel found, use first non-None type
            model = next((arg for arg in union_args if arg is not type(None)), union_args[0])

    model = cast(TG, model)

    model_dict = {}
    # this block is needed to avoid situations where an inbuilt is inside a list / dict.
    # the only time dict ever appears is with include_all, which is handled specially
    # the only type of collection possible after OSCAL 1.0.0 is list
    if safe_is_sub(model, OscalBaseModel):
        for field in model.__fields__:
            if model_type in [OscalVersion]:
                model_dict[field] = OSCAL_VERSION
                break
            # Special handling for include_all field - only skip if it's optional
            if field == 'include_all':
                if model.__fields__[field].required:  # type: ignore
                    # Field is required, generate it
                    model_dict[field] = {}
                elif include_optional:
                    # Field is optional and we want to include optional fields
                    model_dict[field] = {}
                continue
            outer_type = model.__fields__[field].outer_type_  # type: ignore

            # Skip fields with unresolved ForwardRefs, but if required, provide empty list
            if isinstance(outer_type, (str, ForwardRef)):
                # If it's a required field, we need to provide something
                # Assume it's a list type and provide an empty list
                if model.__fields__[field].required:  # type: ignore
                    model_dict[field] = []
                continue

            # Handle both typing.Union and types.UnionType (Python 3.10+ uses | operator)
            origin = utils.get_origin(outer_type)
            is_union = origin == Union or str(origin) == "<class 'types.UnionType'>"
            if is_union:
                # For Union types, prefer Enum types over other types for sample generation
                # This handles fields like Union[ConstrainedStr, Enum, None]
                union_args = typing.get_args(outer_type)
                enum_type = None
                for arg in union_args:
                    if arg is not type(None) and safe_is_sub(arg, Enum):
                        enum_type = arg
                        break
                # Use the enum type if found, otherwise fall back to first non-None, non-ForwardRef type
                if enum_type:
                    outer_type = enum_type
                else:
                    # Get first non-None, non-ForwardRef type
                    # Skip ForwardRef types as they haven't been resolved yet
                    outer_type = None
                    for arg in union_args:
                        # Check if arg is not None and not a ForwardRef
                        if arg is not type(None) and not isinstance(arg, (str, ForwardRef)):
                            outer_type = arg
                            break
                    if outer_type is None:
                        # If all types are ForwardRefs or None, skip this field
                        continue
            if model.__fields__[field].required or effective_optional:  # type: ignore
                # FIXME could be ForwardRef('SystemComponentStatus')
                if utils.is_collection_field_type(outer_type):
                    inner_type = utils.get_inner_type(outer_type)
                    # Check for circular reference: inner_type might be a Union containing model
                    if inner_type == model:
                        continue
                    # Also check if inner_type is a Union and model is one of its variants
                    inner_origin = utils.get_origin(inner_type)
                    is_inner_union = inner_origin == Union or str(inner_origin) == "<class 'types.UnionType'>"
                    if is_inner_union:
                        union_args = typing.get_args(inner_type)
                        if model in union_args:
                            continue  # Circular reference detected
                    # Skip recursion if depth is 0 (but allow -1 for unlimited)
                    if depth == 0:
                        model_dict[field] = []
                    else:
                        model_dict[field] = generate_sample_model(
                            outer_type, include_optional=include_optional, depth=depth - 1
                        )
                elif is_by_type(outer_type):
                    model_dict[field] = generate_sample_value_by_type(outer_type, field)
                elif safe_is_sub(outer_type, OscalBaseModel):
                    # Skip recursion if depth is 0 (but allow -1 for unlimited)
                    # But always generate required fields even at depth 0
                    if depth == 0 and not model.__fields__[field].required:  # type: ignore
                        continue  # Skip optional nested models at depth 0
                    else:
                        model_dict[field] = generate_sample_model(
                            outer_type, include_optional=include_optional, depth=depth - 1
                        )
                else:
                    # Handle special cases (hacking)
                    if model_type in [Base64Datatype]:
                        model_dict[field] = sample_base64_value
                    elif model_type in [Base64]:
                        if field == 'filename':
                            model_dict[field] = sample_base64.filename
                        elif field == 'media_type':
                            model_dict[field] = sample_base64.media_type
                        elif field == 'value':
                            model_dict[field] = sample_base64.value
                    elif model_type in [DateDatatype]:
                        model_dict[field] = sample_date_value
                    # Hacking here:
                    # Root models should ideally not exist, however, sometimes we are stuck with them.
                    # If that is the case we need sufficient information on the type in order to generate a model.
                    # E.g. we need the type of the container.
                    elif field == '__root__' and hasattr(model, '__name__'):
                        model_dict[field] = generate_sample_value_by_type(
                            outer_type, str_utils.classname_to_alias(model.__name__, AliasMode.FIELD)
                        )
                    else:
                        model_dict[field] = generate_sample_value_by_type(outer_type, field)
        # Note: this assumes list constrains in oscal are always 1 as a minimum size. if two this may still fail.
    else:
        if model_type is list:
            return [generate_sample_value_by_type(model, '')]  # type: ignore
        if model_type is dict:
            return {const.REPLACE_ME: generate_sample_value_by_type(model, '')}  # type: ignore
        raise err.TrestleError('Unhandled collection type.')
    if model_type is list:
        return [model(**model_dict)]  # type: ignore
    if model_type is dict:
        return {const.REPLACE_ME: model(**model_dict)}  # type: ignore
    return model(**model_dict)  # type: ignore

def generate_sample_value_by_type(type_: type, field_name: str) -> Union[datetime, bool, int, str, float, Enum, Base64]:
    """Given a type, return sample value.

    Includes the Optional use of passing down a parent_model
    """
    # FIXME: Should be in separate generator module as it inherits EVERYTHING
    if is_enum_method(type_):
        return sample_method
    if is_enum_task_valid_value(type_):
        return sample_task_valid_value
    if is_enum_observation_type_valid_value(type_):
        return sample_observation_type_valid_value
    if type_ is Base64:
        return sample_base64
    if type_ is datetime:
        return datetime.now().astimezone()
    if type_ is bool:
        return False
    if type_ is int:
        return 0
    if type_ is float:
        return 0.00
    if safe_is_sub(type_, ConstrainedStr) or (hasattr(type_, '__name__') and 'ConstrainedStr' in type_.__name__):
        # This code here is messy. we need to meet a set of constraints. If we do
        # TODO: handle regex directly
        if 'uuid' == field_name:
            return str(uuid.uuid4())
        # some things like location_uuid in lists arrive here with field_name=''
        if type_.regex and type_.regex.pattern.startswith('^[0-9A-Fa-f]{8}'):
            return const.SAMPLE_UUID_STR
        if field_name == 'date_authorized':
            return str(date.today().isoformat())
        if field_name == 'oscal_version':
            return OSCAL_VERSION
        if 'uuid' in field_name:
            return const.SAMPLE_UUID_STR
        # Only case where are UUID is required but not in name.
        if field_name.rstrip('s') == 'member_of_organization':
            return const.SAMPLE_UUID_STR
        return const.REPLACE_ME
    if hasattr(type_, '__name__') and 'ConstrainedIntValue' in type_.__name__:
        # create an int value as close to the floor as possible does not test upper bound
        multiple = type_.multiple_of if type_.multiple_of else 1  # default to every integer
        # this command is a bit of a problem
        floor = type_.ge if type_.ge else 0
        floor = type_.gt + 1 if type_.gt else floor
        if math.remainder(floor, multiple) == 0:
            return floor
        return (floor + 1) * multiple
    if safe_is_sub(type_, Enum):
        # keys and values diverge due to hypens in oscal names
        return type_(list(type_.__members__.values())[0])
    if type_ is str:
        if field_name == 'oscal_version':
            return OSCAL_VERSION
        return const.REPLACE_ME
    if type_ is pydantic.v1.networks.EmailStr:
        return pydantic.v1.networks.EmailStr('dummy@sample.com')
    if type_ is pydantic.v1.networks.AnyUrl:
        # TODO: Cleanup: this should be usable from a url.. but it's not inuitive.
        return pydantic.v1.networks.AnyUrl('https://sample.com/replaceme.html', scheme='http', host='sample.com')
    if type_ is list:
        raise err.TrestleError(f'Unable to generate sample for type {type_}')
    # default to empty dict for dict types, string for anything else
    # Check for both dict and generic dict types like dict[str, Any]
    if type_ is dict or (hasattr(type_, '__origin__') and type_.__origin__ is dict):
        return {}  # type: ignore[return-value]
    return const.REPLACE_ME

def is_by_type(model_type: Union[Type[TG], List[TG], Dict[str, TG]]) -> bool:
    """Check for by type."""
    rval = False
    if model_type == type_base64:
        rval = True
    return rval

def is_enum_method(type_: type) -> bool:
    """Test for method."""
    rval = False
    if utils.get_origin(type_) == Union:
        args = typing.get_args(type_)
        for arg in args:
            if "<enum 'Methods'>" == f'{arg}':
                rval = True
                break
    return rval

def is_enum_observation_type_valid_value(type_: type) -> bool:
    """Test for observation type valid value."""
    rval = False
    if utils.get_origin(type_) == Union:
        args = typing.get_args(type_)
        for arg in args:
            if "<enum 'ObservationTypeValidValues'>" == f'{arg}':
                rval = True
                break
    return rval

def is_enum_task_valid_value(type_: type) -> bool:
    """Test for task valid value."""
    rval = False
    if utils.get_origin(type_) == Union:
        args = typing.get_args(type_)
        for arg in args:
            if "<enum 'TaskValidValues'>" == f'{arg}':
                rval = True
                break
    return rval

def safe_is_sub(sub: Any, parent: Any) -> bool:
    """Is this a subclass of parent."""
    # Handle Python 3.10+ generic types (e.g., dict[str, Any])
    # These are types.GenericAlias and cannot be used with issubclass()
    if hasattr(sub, '__origin__'):
        # For generic types like dict[str, Any], check the origin (dict)
        sub = typing.get_origin(sub)
    is_class = inspect.isclass(sub)
    return is_class and issubclass(sub, parent)