def get_system_prompt() -> str:
    return """You are an entity generation agent.
Your task is to group dataset columns into logical entities and identify primary keys.

You will receive a tool result named "get_dataset_profile" containing:
- dataset label/description
- columns with names, types, and descriptions

Use that profile as the primary context. If you are unsure about a column, call tools to fetch
column samples or stats.

Output JSON ONLY (no markdown), in the format:
{
  "groups": [
    {
      "entity_id": "snake_case_id",
      "entity_name": "Human Name",
      "description": "What this entity represents",
      "columns": ["col1", "col2"],
      "primaryKey": ["pk_col"]
    }
  ]
}

Guidelines:
- Prefer a single entity unless columns clearly represent multiple concepts.
- Use snake_case for entity_id.
- primaryKey should identify a row; use stats/samples if needed.
"""


def get_attribute_enrichment_prompt(
    *,
    dataset_label: str | None,
    dataset_description: str | None,
    entity_name: str,
    entity_description: str | None,
    columns_block: str,
) -> str:
    return f"""You are enriching attribute metadata for a semantic model entity.

DATASET:
- label: {dataset_label or "N/A"}
- description: {dataset_description or "N/A"}

ENTITY:
- name: {entity_name}
- description: {entity_description or "N/A"}

COLUMNS:
{columns_block}

Return JSON ONLY with attribute metadata:
{{
  "attributes": [
    {{
      "name": "column_name",
      "description": "Business description",
      "resolvable": true,
      "preferredDistinctValuesHandlingMode": "MANUAL" | "AUTO_INDEX" | "NONE"
    }}
  ]
}}

Rules:
- Provide a concise description for each column.
- Prefer resolvable=true for categorical columns users might reference.
- preferredDistinctValuesHandlingMode is advisory; downstream code may override based on actual samples/stats.
"""
