import json


def _schema_parts(parameters: dict[str, str]) -> list[str]:
    """Collect Spark schema fragments from Glue table parameters."""
    parts: list[tuple[int, str]] = []
    for key, value in parameters.items():
        if not key.startswith("spark.sql.sources.schema.part."):
            continue
        index = int(key.rsplit(".", 1)[1])
        parts.append((index, value))
    return [value for _, value in sorted(parts)]


def _struct_child_path(prefix: str, name: str) -> str:
    """Build the child path for a nested field."""
    return f"{prefix}.{name}" if prefix else name


def _flatten_spark_field(
    field: dict,
    prefix: str,
) -> dict[str, str]:
    """
    Flatten Spark field comments into dotted field paths.
    """
    name = field["name"]
    path = _struct_child_path(prefix, name)
    descriptions: dict[str, str] = {}

    metadata = field.get("metadata", {})
    comment = metadata.get("comment")
    if isinstance(comment, str) and comment:
        descriptions[path] = comment

    data_type = field.get("type")
    if isinstance(data_type, dict):
        type_name = data_type.get("type")
        if type_name == "struct":
            # Struct fields become dotted child paths.
            for child in data_type.get("fields", []):
                descriptions.update(_flatten_spark_field(child, path))
        elif type_name == "array":
            element_type = data_type.get("elementType")
            if (
                isinstance(element_type, dict)
                and element_type.get("type") == "struct"
            ):
                # Arrays of structs are flattened under the array field itself.
                for child in element_type.get("fields", []):
                    descriptions.update(_flatten_spark_field(child, path))

    return descriptions


def spark_field_descriptions(parameters: dict[str, str]) -> dict[str, str]:
    """
    Extract Spark field comments from Glue table parameters.

    Spark can persist nested field comments in Glue table properties under
    keys such as `spark.sql.sources.schema.part.0`,
    `spark.sql.sources.schema.part.1`, etc. Those values form one JSON schema.

    Example:
        parameters = {
            "spark.sql.sources.schema.part.0": (
                '{"type":"struct","fields":['
                '{"name":"activity","type":{"type":"struct","fields":['
                '{"name":"city","type":"string","metadata":{"comment":"city"}}'
                ']},"metadata":{"comment":"activity"}}]}'
            )
        }
        output = {
            "activity": "activity",
            "activity.city": "city",
        }
    """
    schema_json = "".join(_schema_parts(parameters))
    if not schema_json:
        return {}

    schema = json.loads(schema_json)
    descriptions: dict[str, str] = {}
    for field in schema.get("fields", []):
        descriptions.update(_flatten_spark_field(field, prefix=""))
    return descriptions
