from typing import Dict
from pydantic import BaseModel  # noqa: E0611


def count_segments(values: Dict) -> int:
    """
    Returns the number of segment records contained within a X12 data model.

    The X12 data model's top level structure includes the following keys:
    * header
    * <top level loop>:
    * footer
    The top level keys may contain additional nested structures which contain "segment" keys such as
    nm1_segment, st_segment, dtp_segment, se_segment, etc.

    :param values: The validated model values
    :returns: the total segment count
    """
    segment_count: int = 0

    for k, v in values.items():
        if k.endswith("_segment") and isinstance(v, dict):
            segment_count += 1
        elif k.endswith("_segment") and isinstance(v, list):
            segment_count += len(v)
        elif isinstance(v, BaseModel):
            segment_count += count_segments(v.dict())
        elif isinstance(v, list):
            for item in v:
                segment_count += (
                    count_segments(item.dict())
                    if hasattr(item, "dict")
                    else count_segments(item)
                )
        elif isinstance(v, dict):
            segment_count += count_segments(v)

    return segment_count