from typing import Optional
from analogai.deepthink.application.usecases.chat.shared.relationship_category import RelationshipCategory


CATEGORY_MAP = {
    "which": RelationshipCategory.WHICH,
    "how_many": RelationshipCategory.HOW_MANY,
    "is": RelationshipCategory.SIMPLE,
    "are": RelationshipCategory.SIMPLE,
    "what": RelationshipCategory.WHAT,
    "when": RelationshipCategory.WHEN,
    "where": RelationshipCategory.WHERE,
    "how": RelationshipCategory.HOW,
    "how_it_works": RelationshipCategory.HOW,
    "who": RelationshipCategory.WHO,
    "how_much": RelationshipCategory.HOW_MUCH,
    "quantity": RelationshipCategory.HOW_MUCH,
    "price": RelationshipCategory.HOW_MUCH,
    "distance": RelationshipCategory.DISTANCE,
    "can": RelationshipCategory.CAN,
    "did": RelationshipCategory.SIMPLE,
    "does": RelationshipCategory.SIMPLE,
}


def map_category(category_str: Optional[str]) -> RelationshipCategory:
    if not category_str:
        return RelationshipCategory.SIMPLE
    return CATEGORY_MAP.get(category_str.strip().lower(), RelationshipCategory.SIMPLE)
